"""Tests for model loading, synthetic inference, and benchmarking.

All models are mocked — no real weights needed.
"""

from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import MagicMock

import numpy as np
import pytest
import torch
from torch.utils.data import DataLoader

from src.augmentor.inference import (
    compute_edges_from_depth,
    infer_chmv2_batch,
    infer_depth_batch,
    infer_segmentation_batch,
)
from src.utils.benchmark import (
    BenchmarkResult,
    SyntheticImageDataset,
    benchmark_chmv2,
    benchmark_depth,
    benchmark_edges,
    benchmark_full_pipeline,
    benchmark_segmentation,
)

from src.conf.hardware_conf import HardwareConfig

DEVICE = torch.device("cpu")


# ---------------------------------------------------------------------------
# Mock model factories
# ---------------------------------------------------------------------------

def _mock_da3(H: int = 64, W: int = 64) -> MagicMock:
    """Mock DA3 model with inference() method."""
    model = MagicMock()

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

    model.inference = inference
    return model


def _mock_transformers_depth(H: int = 64, W: int = 64) -> MagicMock:
    """Mock transformers-style depth model (DA V2 fallback)."""
    mock_param = torch.nn.Parameter(torch.zeros(1))
    model = MagicMock(spec=[])
    model.parameters = MagicMock(return_value=iter([mock_param]))

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

    model.return_value = None
    model.side_effect = lambda *a, **kw: forward(**kw)
    # spec=[] means hasattr(model, "inference") is False
    return model


def _mock_segformer(num_classes: int = 5, H: int = 64, W: int = 64):
    """Mock SegFormer model + config dict."""
    mock_param = torch.nn.Parameter(torch.zeros(1))
    model = MagicMock()
    model.parameters = lambda: iter([mock_param])

    def forward(pixel_values=None):
        B = pixel_values.shape[0]
        return SimpleNamespace(logits=torch.randn(B, num_classes, H // 4, W // 4))

    model.side_effect = forward
    processor = SimpleNamespace(
        image_mean=[0.485, 0.456, 0.406],
        image_std=[0.229, 0.224, 0.225],
    )
    seg_config = {"type": "segformer", "processor": processor, "num_classes": num_classes}
    return model, seg_config


def _mock_segearth(num_classes: int = 5, H: int = 64, W: int = 64):
    """Mock SegEarth-OV3 model + config dict."""
    model = MagicMock()

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

    model.predict = predict
    seg_config = {
        "type": "segearth-ov3",
        "prompts": [f"class_{i}" for i in range(num_classes)],
        "num_classes": num_classes,
    }
    return model, seg_config


def _mock_chmv2(H: int = 64, W: int = 64) -> tuple[MagicMock, MagicMock]:
    """Mock CHMv2 model + processor."""
    mock_param = torch.nn.Parameter(torch.zeros(1))
    model = MagicMock()
    model.parameters = lambda: iter([mock_param])

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

    model.side_effect = forward

    processor = MagicMock()

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

    processor.side_effect = preprocess

    def post_process(outputs, target_sizes=None):
        B = outputs.predicted_depth.shape[0]
        results = []
        for i in range(B):
            th, tw = target_sizes[i] if target_sizes else (H, W)
            results.append({"predicted_depth": torch.rand(th, tw)})
        return results

    processor.post_process_depth_estimation = post_process
    return model, processor


# ===========================================================================
# Model loading tests
# ===========================================================================

class TestModelLoading:
    """Test that model loading functions work with mock imports."""

    def test_da3_model_loads(self) -> None:
        """DA3 mock model initializes and has inference method."""
        model = _mock_da3()
        assert hasattr(model, "inference")
        result = model.inference([np.zeros((64, 64, 3), dtype=np.uint8)])
        assert result.depth.shape == (1, 64, 64)

    def test_transformers_depth_model_loads(self) -> None:
        """Transformers fallback mock model initializes."""
        model = _mock_transformers_depth()
        assert not hasattr(model, "inference")
        params = list(model.parameters())
        assert len(params) > 0

    def test_segformer_model_loads(self) -> None:
        """SegFormer mock loads with processor and config."""
        model, seg_config = _mock_segformer()
        assert seg_config["type"] == "segformer"
        assert "processor" in seg_config
        assert seg_config["num_classes"] == 5

    def test_chmv2_model_loads(self) -> None:
        """CHMv2 mock loads with model + processor."""
        model, processor = _mock_chmv2()
        assert callable(processor.post_process_depth_estimation)

    def test_segearth_model_loads(self) -> None:
        """SegEarth mock loads with predict method."""
        model, seg_config = _mock_segearth()
        assert hasattr(model, "predict")
        assert seg_config["type"] == "segearth-ov3"
        assert len(seg_config["prompts"]) == 5


# ===========================================================================
# Synthetic inference tests
# ===========================================================================

class TestDepthInference:
    """Depth inference on synthetic data."""

    @pytest.mark.parametrize("batch_size", [1, 2, 4])
    def test_da3_output_shape(self, batch_size: int) -> None:
        H = W = 64
        model = _mock_da3(H, W)
        images = torch.rand(batch_size, 3, H, W)
        result = infer_depth_batch(model, images, DEVICE)
        assert result.shape == (batch_size, 1, H, W)

    def test_da3_output_normalized(self) -> None:
        model = _mock_da3(64, 64)
        images = torch.rand(2, 3, 64, 64)
        result = infer_depth_batch(model, images, DEVICE)
        assert result.min() >= 0.0
        assert result.max() <= 1.0 + 1e-6

    def test_transformers_fallback_output(self) -> None:
        H = W = 32
        mock_param = torch.nn.Parameter(torch.zeros(1))
        model = MagicMock(spec=[])
        model.parameters = MagicMock(return_value=iter([mock_param]))
        pred = torch.rand(2, H, W)
        model.return_value = SimpleNamespace(predicted_depth=pred)

        images = torch.rand(2, 3, H, W)
        result = infer_depth_batch(model, images, DEVICE)
        assert result.shape == (2, 1, H, W)
        assert result.min() >= 0.0
        assert result.max() <= 1.0 + 1e-6

    @pytest.mark.parametrize("image_size", [32, 64, 128])
    def test_different_resolutions(self, image_size: int) -> None:
        model = _mock_da3(image_size, image_size)
        images = torch.rand(2, 3, image_size, image_size)
        result = infer_depth_batch(model, images, DEVICE)
        assert result.shape == (2, 1, image_size, image_size)


class TestChmv2Inference:
    """CHMv2 inference on synthetic data."""

    @pytest.mark.parametrize("batch_size", [1, 2, 4])
    def test_output_shape(self, batch_size: int) -> None:
        H = W = 64
        model, processor = _mock_chmv2(H, W)
        images = torch.rand(batch_size, 3, H, W)
        result = infer_chmv2_batch(model, processor, images, DEVICE)
        assert result.shape == (batch_size, 1, H, W)

    def test_output_normalized(self) -> None:
        H = W = 64
        model, processor = _mock_chmv2(H, W)
        images = torch.rand(2, 3, H, W)
        result = infer_chmv2_batch(model, processor, images, DEVICE)
        assert result.min() >= 0.0
        assert result.max() <= 1.0 + 1e-6

    @pytest.mark.parametrize("image_size", [32, 64, 128])
    def test_different_resolutions(self, image_size: int) -> None:
        model, processor = _mock_chmv2(image_size, image_size)
        images = torch.rand(2, 3, image_size, image_size)
        result = infer_chmv2_batch(model, processor, images, DEVICE)
        assert result.shape == (2, 1, image_size, image_size)


class TestEdgesInference:
    """Edges (Sobel) computation on synthetic data."""

    @pytest.mark.parametrize("batch_size", [1, 2, 8])
    def test_output_shape(self, batch_size: int) -> None:
        depth = torch.rand(batch_size, 1, 64, 64)
        edges = compute_edges_from_depth(depth)
        assert edges.shape == depth.shape

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

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

    def test_deterministic(self) -> None:
        depth = torch.rand(2, 1, 32, 32)
        e1 = compute_edges_from_depth(depth)
        e2 = compute_edges_from_depth(depth)
        torch.testing.assert_close(e1, e2)


class TestSegmentationInference:
    """Segmentation inference on synthetic data."""

    @pytest.mark.parametrize("batch_size", [1, 2, 4])
    def test_segformer_output_shape(self, batch_size: int) -> None:
        H = W = 64
        model, seg_config = _mock_segformer(5, H, W)
        images = torch.rand(batch_size, 3, H, W)
        result = infer_segmentation_batch(model, seg_config, images, DEVICE)
        assert result.shape == (batch_size, 1, H, W)
        assert result.dtype == torch.uint8

    def test_segformer_class_range(self) -> None:
        num_classes = 5
        model, seg_config = _mock_segformer(num_classes, 64, 64)
        images = torch.rand(4, 3, 64, 64)
        result = infer_segmentation_batch(model, seg_config, images, DEVICE)
        assert result.min() >= 0
        assert result.max() < num_classes

    @pytest.mark.parametrize("batch_size", [1, 2])
    def test_segearth_output_shape(self, batch_size: int) -> None:
        H = W = 64
        model, seg_config = _mock_segearth(5, H, W)
        images = torch.rand(batch_size, 3, H, W)
        result = infer_segmentation_batch(model, seg_config, images, DEVICE)
        assert result.shape == (batch_size, 1, H, W)
        assert result.dtype == torch.uint8

    def test_segearth_failure_returns_zeros(self) -> None:
        model = MagicMock()
        model.predict.side_effect = RuntimeError("OOM")
        seg_config = {"type": "segearth-ov3", "prompts": ["bg"]}
        images = torch.rand(2, 3, 32, 32)
        result = infer_segmentation_batch(model, seg_config, images, DEVICE)
        assert result.shape == (2, 1, 32, 32)
        assert (result == 0).all()


# ===========================================================================
# DataLoader integration tests
# ===========================================================================

class TestSyntheticDataset:
    """Test the synthetic dataset and DataLoader."""

    def test_dataset_length(self) -> None:
        ds = SyntheticImageDataset(num_images=50, image_size=64)
        assert len(ds) == 50

    def test_dataset_item_shape(self) -> None:
        ds = SyntheticImageDataset(num_images=10, image_size=128)
        item = ds[0]
        assert item["image_raw"].shape == (3, 128, 128)
        assert item["image_raw"].dtype == torch.float32
        assert item["image_raw"].min() >= 0.0
        assert item["image_raw"].max() <= 1.0

    def test_dataloader_batching(self) -> None:
        ds = SyntheticImageDataset(num_images=10, image_size=64)
        loader = DataLoader(ds, batch_size=4, shuffle=False)
        batches = list(loader)
        assert len(batches) == 3  # 4 + 4 + 2
        assert batches[0]["image_raw"].shape == (4, 3, 64, 64)
        assert batches[2]["image_raw"].shape == (2, 3, 64, 64)

    @pytest.mark.parametrize("num_images", [1, 7, 16])
    def test_various_counts(self, num_images: int) -> None:
        ds = SyntheticImageDataset(num_images=num_images, image_size=32)
        loader = DataLoader(ds, batch_size=4, shuffle=False)
        total = sum(b["image_raw"].shape[0] for b in loader)
        assert total == num_images

    def test_depth_inference_from_loader(self) -> None:
        """Full path: DataLoader → inference → check output."""
        model = _mock_da3(64, 64)
        ds = SyntheticImageDataset(num_images=8, image_size=64)
        loader = DataLoader(ds, batch_size=4, shuffle=False)
        all_depths = []
        for batch in loader:
            depths = infer_depth_batch(model, batch["image_raw"], DEVICE)
            all_depths.append(depths)
        combined = torch.cat(all_depths, dim=0)
        assert combined.shape == (8, 1, 64, 64)

    def test_segformer_inference_from_loader(self) -> None:
        model, seg_config = _mock_segformer(5, 64, 64)
        ds = SyntheticImageDataset(num_images=6, image_size=64)
        loader = DataLoader(ds, batch_size=3, shuffle=False)
        all_segs = []
        for batch in loader:
            segs = infer_segmentation_batch(model, seg_config, batch["image_raw"], DEVICE)
            all_segs.append(segs)
        combined = torch.cat(all_segs, dim=0)
        assert combined.shape == (6, 1, 64, 64)
        assert combined.dtype == torch.uint8


# ===========================================================================
# Benchmark result tests
# ===========================================================================

class TestBenchmarkResult:
    """Test BenchmarkResult calculations."""

    def test_fps_calculation(self) -> None:
        r = BenchmarkResult("test", num_images=100, total_time_s=2.0)
        assert r.fps == pytest.approx(50.0)

    def test_avg_latency(self) -> None:
        r = BenchmarkResult("test", num_images=100, total_time_s=2.0)
        assert r.avg_latency_ms == pytest.approx(20.0)

    def test_median_batch_latency(self) -> None:
        r = BenchmarkResult("test", num_images=100, total_time_s=2.0,
                            batch_times_s=[0.01, 0.02, 0.03, 0.04, 0.05])
        assert r.median_batch_latency_ms == pytest.approx(30.0)

    def test_p95_batch_latency(self) -> None:
        times = [0.01] * 95 + [0.1] * 5
        r = BenchmarkResult("test", num_images=100, total_time_s=2.0,
                            batch_times_s=times)
        assert r.p95_batch_latency_ms > r.median_batch_latency_ms  # p95 > median

    def test_zero_images(self) -> None:
        r = BenchmarkResult("test", num_images=0, total_time_s=0.0)
        assert r.fps == 0.0
        assert r.avg_latency_ms == 0.0

    def test_empty_batch_times(self) -> None:
        r = BenchmarkResult("test", num_images=10, total_time_s=1.0)
        assert r.median_batch_latency_ms == 0.0
        assert r.p95_batch_latency_ms == 0.0


# ===========================================================================
# Benchmark function tests (small runs with mocks)
# ===========================================================================

class TestBenchmarkDepth:
    """Benchmark depth with mocked model."""

    def test_runs_and_returns_result(self) -> None:
        model = _mock_da3(32, 32)
        result = benchmark_depth(model, DEVICE, num_images=8,
                                 batch_size=4, image_size=32, warmup_batches=1)
        assert isinstance(result, BenchmarkResult)
        assert result.num_images == 8
        assert result.total_time_s > 0
        assert result.fps > 0
        assert len(result.batch_times_s) == 2  # 8 images / 4 batch

    def test_single_image(self) -> None:
        model = _mock_da3(32, 32)
        result = benchmark_depth(model, DEVICE, num_images=1,
                                 batch_size=1, image_size=32, warmup_batches=0)
        assert result.num_images == 1
        assert len(result.batch_times_s) == 1


class TestBenchmarkChmv2:
    """Benchmark CHMv2 with mocked model."""

    def test_runs_and_returns_result(self) -> None:
        model, processor = _mock_chmv2(32, 32)
        result = benchmark_chmv2(model, processor, DEVICE, num_images=8,
                                 batch_size=4, image_size=32, warmup_batches=1)
        assert isinstance(result, BenchmarkResult)
        assert result.num_images == 8
        assert result.total_time_s > 0
        assert result.fps > 0
        assert len(result.batch_times_s) == 2

    def test_single_image(self) -> None:
        model, processor = _mock_chmv2(32, 32)
        result = benchmark_chmv2(model, processor, DEVICE, num_images=1,
                                 batch_size=1, image_size=32, warmup_batches=0)
        assert result.num_images == 1
        assert len(result.batch_times_s) == 1


class TestBenchmarkEdges:
    """Benchmark edges computation."""

    def test_runs_and_returns_result(self) -> None:
        result = benchmark_edges(num_images=16, batch_size=8,
                                 image_size=32, warmup_batches=1)
        assert isinstance(result, BenchmarkResult)
        assert result.num_images == 16
        assert result.fps > 0
        assert len(result.batch_times_s) == 2  # 16 / 8

    def test_odd_count(self) -> None:
        result = benchmark_edges(num_images=7, batch_size=3,
                                 image_size=32, warmup_batches=0)
        assert result.num_images == 7
        assert len(result.batch_times_s) == 3  # ceil(7/3)


class TestBenchmarkSegmentation:
    """Benchmark segmentation with mocked model."""

    def test_segformer(self) -> None:
        model, seg_config = _mock_segformer(5, 32, 32)
        result = benchmark_segmentation(model, seg_config, DEVICE,
                                        num_images=8, batch_size=4,
                                        image_size=32, warmup_batches=1)
        assert isinstance(result, BenchmarkResult)
        assert result.num_images == 8
        assert result.fps > 0

    def test_segearth(self) -> None:
        model, seg_config = _mock_segearth(3, 32, 32)
        result = benchmark_segmentation(model, seg_config, DEVICE,
                                        num_images=4, batch_size=2,
                                        image_size=32, warmup_batches=0)
        assert result.num_images == 4


class TestBenchmarkFullPipeline:
    """Benchmark full pipeline with all stages mocked."""

    def test_runs_all_stages(self) -> None:
        depth_model = _mock_da3(32, 32)
        seg_model, seg_config = _mock_segformer(5, 32, 32)

        results = benchmark_full_pipeline(
            depth_model, seg_model, seg_config, DEVICE,
            num_images=8, batch_size=4, image_size=32, warmup_batches=1,
        )

        assert "depth" in results
        assert "edges" in results
        assert "segmentation" in results
        assert "total" in results

        for key, r in results.items():
            assert isinstance(r, BenchmarkResult)
            assert r.num_images == 8
            assert r.total_time_s > 0
            assert r.fps > 0

    def test_total_ge_sum_of_stages(self) -> None:
        """Total time >= sum of individual stages (includes overhead)."""
        depth_model = _mock_da3(32, 32)
        seg_model, seg_config = _mock_segformer(3, 32, 32)

        results = benchmark_full_pipeline(
            depth_model, seg_model, seg_config, DEVICE,
            num_images=4, batch_size=2, image_size=32, warmup_batches=0,
        )

        stage_sum = (results["depth"].total_time_s +
                     results["edges"].total_time_s +
                     results["segmentation"].total_time_s)
        assert results["total"].total_time_s >= stage_sum * 0.95  # allow 5% tolerance

    def test_configurable_image_count(self) -> None:
        """num_images parameter controls how many images are processed."""
        depth_model = _mock_da3(32, 32)
        seg_model, seg_config = _mock_segformer(3, 32, 32)

        for n in [4, 10, 17]:
            results = benchmark_full_pipeline(
                depth_model, seg_model, seg_config, DEVICE,
                num_images=n, batch_size=4, image_size=32, warmup_batches=0,
            )
            assert results["total"].num_images == n

    def test_different_batch_sizes(self) -> None:
        depth_model = _mock_da3(32, 32)
        seg_model, seg_config = _mock_segformer(3, 32, 32)

        for bs in [1, 2, 8]:
            results = benchmark_full_pipeline(
                depth_model, seg_model, seg_config, DEVICE,
                num_images=8, batch_size=bs, image_size=32, warmup_batches=0,
            )
            expected_batches = (8 + bs - 1) // bs
            assert len(results["total"].batch_times_s) == expected_batches


# ===========================================================================
# find_batch_size tests
# ===========================================================================

class TestFindBatchSize:
    """Test empirical batch size probing."""

    def test_returns_config_batch_size_when_set(self) -> None:
        """If batch_size is set in config, find_batch_size returns it directly."""
        hw = HardwareConfig(batch_size=8)
        bs = hw.find_batch_size(
            inference_fn=lambda x: None,
            input_shape=(3, 64, 64),
            device=DEVICE,
        )
        assert bs == 8

    def test_cpu_returns_1(self) -> None:
        """On CPU (no CUDA), returns 1."""
        hw = HardwareConfig(batch_size=None)
        bs = hw.find_batch_size(
            inference_fn=lambda x: None,
            input_shape=(3, 64, 64),
            device=torch.device("cpu"),
        )
        assert bs == 1

    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
    def test_finds_batch_gt_1_on_gpu(self) -> None:
        """On a real GPU with a trivial model, batch should be > 1."""
        hw = HardwareConfig(batch_size=None)
        device = torch.device("cuda")

        def trivial_fn(x: torch.Tensor) -> None:
            _ = x.to(device)
            torch.cuda.synchronize()

        bs = hw.find_batch_size(
            inference_fn=trivial_fn,
            input_shape=(3, 32, 32),
            device=device,
            max_batch=16,
        )
        assert bs >= 1

    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
    def test_measures_per_sample_cost(self) -> None:
        """Heavier inference → smaller batch size."""
        hw = HardwareConfig(batch_size=None)
        device = torch.device("cuda")

        # Light: tiny tensor on GPU.
        def light_fn(x: torch.Tensor) -> None:
            _ = x.to(device)
            torch.cuda.synchronize()

        # Heavy: allocate a large intermediate on GPU.
        def heavy_fn(x: torch.Tensor) -> None:
            y = x.to(device)
            # ~200 MB intermediate per sample for 512x512.
            big = torch.randn(x.shape[0], 64, 512, 512, device=device)
            _ = big.sum()
            torch.cuda.synchronize()

        bs_light = hw.find_batch_size(
            inference_fn=light_fn,
            input_shape=(3, 32, 32),
            device=device,
            max_batch=64,
        )
        bs_heavy = hw.find_batch_size(
            inference_fn=heavy_fn,
            input_shape=(3, 32, 32),
            device=device,
            max_batch=64,
        )
        assert bs_light >= bs_heavy

    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
    def test_respects_max_batch(self) -> None:
        """Batch size never exceeds max_batch."""
        hw = HardwareConfig(batch_size=None)
        device = torch.device("cuda")

        bs = hw.find_batch_size(
            inference_fn=lambda x: x.to(device),
            input_shape=(3, 32, 32),
            device=device,
            max_batch=4,
        )
        assert bs <= 4

    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
    def test_result_is_power_of_2(self) -> None:
        """Batch size is always a power of 2."""
        hw = HardwareConfig(batch_size=None)
        device = torch.device("cuda")

        bs = hw.find_batch_size(
            inference_fn=lambda x: x.to(device),
            input_shape=(3, 64, 64),
            device=device,
        )
        assert bs >= 1
        assert (bs & (bs - 1)) == 0  # power of 2 check
