"""
tests/test_classifier.py
────────────────────────
Unit and integration tests for the medical-ai module.

Run from the repo root:
    pytest tests/ -v

Import paths assume the project root (medical-ai/) is the working directory.
"""

import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import numpy as np
import pytest
import torch
from PIL import Image

from model import BreastCancerClassifier, BreastCancerInferencePipeline
from utils import ImagePreprocessor


# ── Fixtures ──────────────────────────────────────────────────────────────────

@pytest.fixture(scope="module")
def dummy_pil():
    """224×224 RGB PIL image filled with mid-grey."""
    return Image.fromarray(
        np.full((224, 224, 3), 128, dtype=np.uint8), mode="RGB"
    )

@pytest.fixture(scope="module")
def dummy_np():
    """224×224×3 uint8 numpy array."""
    return np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)

@pytest.fixture(scope="module")
def dummy_tensor():
    """Properly-shaped (1, 3, 224, 224) float32 tensor."""
    return torch.rand(1, 3, 224, 224)

@pytest.fixture(scope="module")
def model():
    return BreastCancerClassifier(pretrained=False)

@pytest.fixture(scope="module")
def pipeline():
    return BreastCancerInferencePipeline(weights_path=None, device="cpu")


# ── model/model.py ────────────────────────────────────────────────────────────

class TestBreastCancerClassifier:

    def test_output_keys(self, model):
        x = torch.rand(1, 3, 224, 224)
        out = model(x)
        assert "logits" in out and "probs" in out

    def test_logits_shape(self, model):
        x = torch.rand(2, 3, 224, 224)
        out = model(x)
        assert out["logits"].shape == (2, 2), "Logits must be (B, 2)"

    def test_probs_sum_to_one(self, model):
        x = torch.rand(4, 3, 224, 224)
        out = model(x)
        sums = out["probs"].sum(dim=1)
        assert torch.allclose(sums, torch.ones(4), atol=1e-5)

    def test_probs_in_range(self, model):
        x = torch.rand(1, 3, 224, 224)
        out = model(x)
        assert out["probs"].min() >= 0.0
        assert out["probs"].max() <= 1.0

    def test_feature_maps_shape(self, model):
        """Validates Grad-CAM hook compatibility."""
        x = torch.rand(1, 3, 224, 224)
        fm = model.get_feature_maps(x)
        assert fm.shape == (1, 1024, 7, 7)

    def test_deterministic_output(self, model):
        """Identical inputs must produce identical outputs at eval time."""
        model.eval()
        x = torch.rand(1, 3, 224, 224)
        with torch.no_grad():
            out_a = model(x)
            out_b = model(x)
        assert torch.allclose(out_a["logits"], out_b["logits"])

    def test_batch_inference(self, model):
        x = torch.rand(8, 3, 224, 224)
        out = model(x)
        assert out["logits"].shape == (8, 2)

    def test_freeze_backbone(self):
        m = BreastCancerClassifier(pretrained=False, freeze_backbone=True)
        for param in m.features.parameters():
            assert not param.requires_grad, "Backbone params should be frozen"


# ── utils/preprocessing.py ───────────────────────────────────────────────────

class TestImagePreprocessor:

    def test_pil_input(self, dummy_pil):
        t = ImagePreprocessor()(dummy_pil)
        assert t.shape == (1, 3, 224, 224)
        assert t.dtype == torch.float32

    def test_numpy_input(self, dummy_np):
        t = ImagePreprocessor()(dummy_np)
        assert t.shape == (1, 3, 224, 224)

    def test_tensor_input(self, dummy_tensor):
        t = ImagePreprocessor()(dummy_tensor)
        assert t.shape == (1, 3, 224, 224)

    def test_normalization_shifts_range(self, dummy_pil):
        """ImageNet normalization should shift values outside raw [0,1]."""
        t = ImagePreprocessor()(dummy_pil)
        assert not (t.min() >= 0 and t.max() <= 1)

    def test_invalid_type_raises(self):
        with pytest.raises(TypeError):
            ImagePreprocessor()({"not": "an image"})

    def test_grayscale_numpy_converted_to_rgb(self):
        grey = np.full((224, 224), 128, dtype=np.uint8)
        t = ImagePreprocessor()(grey)
        assert t.shape == (1, 3, 224, 224)


# ── model/inference.py ────────────────────────────────────────────────────────

class TestInferencePipeline:

    def test_output_schema(self, pipeline, dummy_pil):
        result = pipeline.predict(dummy_pil)
        assert "prediction" in result
        assert "confidence" in result
        assert "logits"     in result

    def test_prediction_is_valid_label(self, pipeline, dummy_pil):
        result = pipeline.predict(dummy_pil)
        assert result["prediction"] in ("benign", "malignant")

    def test_confidence_range(self, pipeline, dummy_pil):
        result = pipeline.predict(dummy_pil)
        assert 0.0 <= result["confidence"] <= 1.0

    def test_logits_shape(self, pipeline, dummy_pil):
        result = pipeline.predict(dummy_pil)
        assert result["logits"].shape == (1, 2)

    def test_batch_predict_length(self, pipeline, dummy_pil, dummy_np):
        results = pipeline.predict_batch([dummy_pil, dummy_np])
        assert len(results) == 2

    def test_deterministic_inference(self, pipeline, dummy_pil):
        r1 = pipeline.predict(dummy_pil)
        r2 = pipeline.predict(dummy_pil)
        assert r1["prediction"] == r2["prediction"]
        assert r1["confidence"] == r2["confidence"]
        assert torch.allclose(r1["logits"], r2["logits"])

    def test_numpy_input_accepted(self, pipeline, dummy_np):
        result = pipeline.predict(dummy_np)
        assert result["prediction"] in ("benign", "malignant")

    def test_logits_detached_from_graph(self, pipeline, dummy_pil):
        result = pipeline.predict(dummy_pil)
        assert not result["logits"].requires_grad

    def test_missing_weights_raises(self):
        with pytest.raises(FileNotFoundError):
            BreastCancerInferencePipeline(
                weights_path="model/nonexistent_weights.pth",
                device="cpu",
            )