"""Tests for inference functions: depth, edges, segmentation."""

from __future__ import annotations

from unittest.mock import MagicMock
from types import SimpleNamespace

import numpy as np
import torch

from src.augmentor.inference import (
    compute_edges_from_depth,
    infer_chmv2_batch,
    infer_depth_batch,
    infer_segmentation_batch,
    postprocess_segmentation,
    _SOBEL_X,
    _SOBEL_Y,
)


# ---------------------------------------------------------------------------
# CHMv2 — mock model
# ---------------------------------------------------------------------------

class TestInferChmv2Batch:
    def _make_mock(self, B, H, W):
        model = MagicMock()
        model.parameters = MagicMock(return_value=iter([torch.nn.Parameter(torch.zeros(1))]))

        def forward(pixel_values=None):
            bb = pixel_values.shape[0]
            return SimpleNamespace(predicted_depth=torch.rand(bb, H, W))

        model.side_effect = forward

        processor = MagicMock()

        def preprocess(images=None, return_tensors=None):
            bb = len(images)
            return {"pixel_values": torch.rand(bb, 3, H, W)}

        processor.side_effect = preprocess

        def post_process(outputs, target_sizes=None):
            bb = outputs.predicted_depth.shape[0]
            return [
                {"predicted_depth": torch.rand(target_sizes[i][0], target_sizes[i][1])}
                for i in range(bb)
            ]

        processor.post_process_depth_estimation = post_process
        return model, processor

    def test_output_shape(self, sample_batch: torch.Tensor) -> None:
        B, _, H, W = sample_batch.shape
        model, processor = self._make_mock(B, H, W)
        result = infer_chmv2_batch(model, processor, sample_batch, torch.device("cpu"))
        assert result.shape == (B, 1, H, W)

    def test_output_normalized(self, sample_batch: torch.Tensor) -> None:
        B, _, H, W = sample_batch.shape
        model, processor = self._make_mock(B, H, W)
        result = infer_chmv2_batch(model, processor, sample_batch, torch.device("cpu"))
        assert result.min() >= 0.0
        assert result.max() <= 1.0 + 1e-6


# ---------------------------------------------------------------------------
# Edges (Sobel) — pure computation, no model needed
# ---------------------------------------------------------------------------

class TestComputeEdges:
    def test_output_shape(self, sample_depth: torch.Tensor) -> None:
        edges = compute_edges_from_depth(sample_depth)
        assert edges.shape == sample_depth.shape

    def test_output_range(self, sample_depth: torch.Tensor) -> None:
        edges = compute_edges_from_depth(sample_depth)
        assert edges.min() >= 0.0
        assert edges.max() <= 1.0 + 1e-6

    def test_flat_depth_gives_zero_edges(self) -> None:
        flat = torch.ones(1, 1, 32, 32) * 0.5
        edges = compute_edges_from_depth(flat)
        assert edges.max().item() < 1e-6

    def test_single_image(self) -> None:
        depth = torch.rand(1, 1, 16, 16)
        edges = compute_edges_from_depth(depth)
        assert edges.shape == (1, 1, 16, 16)

    def test_batch_consistency(self) -> None:
        """Batched result matches per-image results."""
        depth = torch.rand(4, 1, 32, 32)
        edges_batch = compute_edges_from_depth(depth)
        for i in range(4):
            edges_single = compute_edges_from_depth(depth[i : i + 1])
            torch.testing.assert_close(edges_batch[i], edges_single[0], atol=1e-5, rtol=1e-5)


class TestSobelKernelCache:
    def test_kernels_are_module_level(self) -> None:
        assert _SOBEL_X.shape == (1, 1, 3, 3)
        assert _SOBEL_Y.shape == (1, 1, 3, 3)
        assert _SOBEL_X.dtype == torch.float32


# ---------------------------------------------------------------------------
# Depth — mock model
# ---------------------------------------------------------------------------

class TestInferDepthBatch:
    def test_da3_api_path(self, sample_batch: torch.Tensor) -> None:
        """Test DA3 branch with model.inference()."""
        B, _, H, W = sample_batch.shape

        def fake_inference(image_list, process_res=None):
            depth = np.random.rand(len(image_list), H, W).astype(np.float32)
            return SimpleNamespace(depth=depth)

        model = MagicMock()
        model.inference = fake_inference

        result = infer_depth_batch(model, sample_batch, torch.device("cpu"))
        assert result.shape == (B, 1, H, W)
        assert result.min() >= 0.0
        assert result.max() <= 1.0 + 1e-6

    def test_transformers_api_path(self, sample_batch: torch.Tensor) -> None:
        """Test DA V2 fallback with model(pixel_values=x).predicted_depth."""
        B, _, H, W = sample_batch.shape

        # Mock a transformers-style model — spec=[] ensures hasattr(model, "inference") is False.
        mock_param = torch.nn.Parameter(torch.zeros(1))  # float32
        model = MagicMock(spec=[])
        model.parameters = MagicMock(return_value=iter([mock_param]))

        pred_depth = torch.rand(B, H, W)
        model.return_value = SimpleNamespace(predicted_depth=pred_depth)

        result = infer_depth_batch(model, sample_batch, torch.device("cpu"))
        assert result.shape == (B, 1, H, W)
        assert result.min() >= 0.0
        assert result.max() <= 1.0 + 1e-6

    def test_per_frame_normalization(self) -> None:
        """Each frame should be independently normalized to [0, 1]."""
        B, H, W = 2, 16, 16

        def fake_inference(image_list, process_res=None):
            depth = np.random.rand(len(image_list), H, W).astype(np.float32) * 100
            return SimpleNamespace(depth=depth)

        model = MagicMock()
        model.inference = fake_inference

        batch = torch.rand(B, 3, H, W)
        result = infer_depth_batch(model, batch, torch.device("cpu"))
        for i in range(B):
            assert result[i].min().item() < 0.01
            assert result[i].max().item() > 0.99


# ---------------------------------------------------------------------------
# Segmentation — mock model
# ---------------------------------------------------------------------------

class TestInferSegmentationBatch:
    def test_segformer_path(self, sample_batch: torch.Tensor) -> None:
        """Test SegFormer fallback branch."""
        B, _, H, W = sample_batch.shape

        mock_param = torch.nn.Parameter(torch.zeros(1))
        model = MagicMock()
        model.parameters.return_value = iter([mock_param])

        num_classes = 5
        logits = torch.randn(B, num_classes, H // 4, W // 4)
        model.return_value = SimpleNamespace(logits=logits)

        processor = SimpleNamespace(
            image_mean=[0.485, 0.456, 0.406],
            image_std=[0.229, 0.224, 0.225],
        )
        seg_config = {"type": "segformer", "processor": processor}

        result = infer_segmentation_batch(model, seg_config, sample_batch, torch.device("cpu"))
        assert result.shape == (B, 1, H, W)
        assert result.dtype == torch.uint8
        assert result.min() >= 0
        assert result.max() < num_classes

    def test_segearth_path(self, sample_batch: torch.Tensor) -> None:
        """Test SegEarth-OV3 branch with mock pipeline."""
        B, _, H, W = sample_batch.shape

        def fake_predict(pil_img, text_prompts=None):
            return np.random.randint(0, 3, (H, W), dtype=np.uint8)

        pipeline = MagicMock()
        pipeline.predict = fake_predict

        seg_config = {
            "type": "segearth-ov3",
            "prompts": ["bg", "building", "road"],
        }

        result = infer_segmentation_batch(pipeline, seg_config, sample_batch, torch.device("cpu"))
        assert result.shape == (B, 1, H, W)
        assert result.dtype == torch.uint8

    def test_segearth_handles_failure(self, sample_batch: torch.Tensor) -> None:
        """Failed prediction should produce zero tensor, not crash."""
        B, _, H, W = sample_batch.shape

        pipeline = MagicMock()
        pipeline.predict.side_effect = RuntimeError("OOM")

        seg_config = {"type": "segearth-ov3", "prompts": ["bg"]}
        result = infer_segmentation_batch(pipeline, seg_config, sample_batch, torch.device("cpu"))
        assert result.shape == (B, 1, H, W)
        assert (result == 0).all()


# ---------------------------------------------------------------------------
# Post-processing heuristics
# ---------------------------------------------------------------------------

class TestPostprocessSegmentation:
    def test_dark_water_reclassified(self) -> None:
        """Dark uniform background → water."""
        # Dark image (mean ~0.15, std ~0.10 — typical dark ocean)
        rgb = torch.full((1, 3, 32, 32), 0.15)
        rgb += torch.randn_like(rgb) * 0.05
        rgb.clamp_(0, 1)
        seg = torch.zeros(1, 1, 32, 32, dtype=torch.uint8)  # all background
        result = postprocess_segmentation(seg, rgb)
        assert (result == 4).all()  # water

    def test_bright_background_unchanged(self) -> None:
        """Bright varied background should NOT become water."""
        rgb = torch.rand(1, 3, 32, 32) * 0.4 + 0.4  # mean ~0.6, std ~0.12
        seg = torch.zeros(1, 1, 32, 32, dtype=torch.uint8)
        result = postprocess_segmentation(seg, rgb)
        assert (result == 0).all()  # stays background

    def test_non_background_unchanged(self) -> None:
        """Non-background classes untouched even if dark."""
        rgb = torch.full((1, 3, 32, 32), 0.1)
        seg = torch.full((1, 1, 32, 32), 3, dtype=torch.uint8)  # vegetation
        result = postprocess_segmentation(seg, rgb)
        assert (result == 3).all()

    def test_wetland_reclassify_green(self) -> None:
        """Green wetland → vegetation when reclassify_wetland=True."""
        rgb = torch.zeros(1, 3, 32, 32)
        rgb[0, 1] = 0.6  # green dominant
        rgb[0, 0] = 0.2  # low red
        seg = torch.full((1, 1, 32, 32), 14, dtype=torch.uint8)  # wetland
        result = postprocess_segmentation(seg, rgb, reclassify_wetland=True)
        assert (result == 3).all()  # vegetation

    def test_wetland_reclassify_brown(self) -> None:
        """Brown wetland → bare soil when reclassify_wetland=True."""
        rgb = torch.zeros(1, 3, 32, 32)
        rgb[0, 0] = 0.5  # red dominant
        rgb[0, 1] = 0.3  # low green
        seg = torch.full((1, 1, 32, 32), 14, dtype=torch.uint8)
        result = postprocess_segmentation(seg, rgb, reclassify_wetland=True)
        assert (result == 11).all()  # bare soil

    def test_wetland_unchanged_without_flag(self) -> None:
        """Wetland stays if reclassify_wetland=False."""
        rgb = torch.zeros(1, 3, 32, 32)
        rgb[0, 1] = 0.6
        seg = torch.full((1, 1, 32, 32), 14, dtype=torch.uint8)
        result = postprocess_segmentation(seg, rgb, reclassify_wetland=False)
        assert (result == 14).all()

    def test_uniform_water_reclassified(self) -> None:
        """Uniform muddy water (high mean, low std, high bg) → water."""
        # Muddy river: mean ~0.46, std ~0.05 — SegEarth misses without land context
        rgb = torch.full((1, 3, 32, 32), 0.46)
        rgb += torch.randn_like(rgb) * 0.02
        rgb.clamp_(0, 1)
        seg = torch.zeros(1, 1, 32, 32, dtype=torch.uint8)  # all background
        result = postprocess_segmentation(seg, rgb, fix_uniform_water=True)
        assert (result == 4).all()  # water

    def test_uniform_water_partial_bg(self) -> None:
        """Uniform water with some classified pixels — bg part → water."""
        rgb = torch.full((1, 3, 32, 32), 0.46)
        rgb += torch.randn_like(rgb) * 0.02
        rgb.clamp_(0, 1)
        seg = torch.zeros(1, 1, 32, 32, dtype=torch.uint8)
        # 10% non-background (e.g. boats)
        seg[0, 0, :3, :] = 1  # building (boats)
        result = postprocess_segmentation(seg, rgb, fix_uniform_water=True)
        assert (result[0, 0, :3, :] == 1).all()  # boats stay
        assert (result[0, 0, 3:, :] == 4).all()  # bg → water

    def test_uniform_water_disabled(self) -> None:
        """Uniform water fix disabled → bg stays."""
        rgb = torch.full((1, 3, 32, 32), 0.46)
        seg = torch.zeros(1, 1, 32, 32, dtype=torch.uint8)
        result = postprocess_segmentation(seg, rgb, fix_uniform_water=False)
        assert (result == 0).all()  # stays background

    def test_uniform_water_high_std_no_fix(self) -> None:
        """High-std image with bg → NOT reclassified (not uniform water)."""
        rgb = torch.rand(1, 3, 32, 32)  # std ~0.29
        seg = torch.zeros(1, 1, 32, 32, dtype=torch.uint8)
        result = postprocess_segmentation(seg, rgb, fix_uniform_water=True)
        assert (result == 0).all()  # stays background

    def test_uniform_water_low_bg_ratio_no_fix(self) -> None:
        """Low bg ratio + low std → NOT reclassified."""
        rgb = torch.full((1, 3, 32, 32), 0.46)
        rgb += torch.randn_like(rgb) * 0.02
        rgb.clamp_(0, 1)
        seg = torch.full((1, 1, 32, 32), 3, dtype=torch.uint8)  # all vegetation
        # Only 20% background
        seg[0, 0, :6, :] = 0
        result = postprocess_segmentation(seg, rgb, fix_uniform_water=True)
        assert (result[0, 0, :6, :] == 0).all()  # bg stays — ratio too low

    def test_blue_ocean_reclassified(self) -> None:
        """Blue ocean with high cross-channel variance → water via ch_std."""
        # Blue ocean: R=0.10, G=0.26, B=0.40, per-channel std ~0.015
        rgb = torch.zeros(1, 3, 32, 32)
        rgb[0, 0] = 0.10 + torch.randn(32, 32) * 0.01  # R
        rgb[0, 1] = 0.26 + torch.randn(32, 32) * 0.01  # G
        rgb[0, 2] = 0.40 + torch.randn(32, 32) * 0.01  # B
        rgb.clamp_(0, 1)
        seg = torch.zeros(1, 1, 32, 32, dtype=torch.uint8)  # all background
        # rgb.std() is ~0.13 (cross-channel), but per-channel std ~0.01
        result = postprocess_segmentation(seg, rgb, fix_uniform_water=True)
        assert (result == 4).all()  # water

    def test_bright_terrain_not_reclassified(self) -> None:
        """Bright terrain with low per-channel std → NOT water (mean > 0.60)."""
        # Bright terrain with strong cross-channel variance: R=0.90, G=0.65, B=0.50
        # Per-channel std < 0.05 but mean ~0.68 > 0.60, so Condition B rejected.
        # rgb_std ~0.16 > 0.10, so Condition A also doesn't trigger.
        rgb = torch.zeros(1, 3, 32, 32)
        rgb[0, 0] = 0.90 + torch.randn(32, 32) * 0.01
        rgb[0, 1] = 0.65 + torch.randn(32, 32) * 0.01
        rgb[0, 2] = 0.50 + torch.randn(32, 32) * 0.01
        rgb.clamp_(0, 1)
        seg = torch.zeros(1, 1, 32, 32, dtype=torch.uint8)
        result = postprocess_segmentation(seg, rgb, fix_uniform_water=True)
        assert (result == 0).all()  # stays background
