tests/test_multiview_pipeline.py

from __future__ import annotations

import unittest
from collections.abc import Sequence

from PIL import Image

from src.clients.medgemma_client import MedGemmaDetector
from src.schemas.detection import (
    BBox as LocalizationBox,
    DetectionResult as PipelineResult,
    LocalizedFinding,
    ViewLocalization,
)
from src.visualization import draw_view_localizations
from tests.test_parsing import REAL_MEDGEMMA_DISCOVERY_RESPONSE


def make_test_image(color: tuple[int, int, int] = (128, 128, 128)) -> Image.Image:
    return Image.new("RGB", (64, 48), color)


def discovery_response(findings: list[dict[str, str]]) -> str:
    import json
    return f"```json\n{json.dumps({'findings': findings})}\n```"


def localization_response(boxes: list[dict[str, object]]) -> str:
    import json
    return f"```json\n{json.dumps(boxes)}\n```"


# ---------------------------------------------------------------------------
# Mock clients
# ---------------------------------------------------------------------------

class MockMedGemmaClient:
    def __init__(self) -> None:
        self.calls: list[dict[str, object]] = []

    def generate(self, *, images: Sequence[Image.Image], prompt: str, max_new_tokens: int) -> str:
        self.calls.append({"images": list(images), "prompt": prompt, "max_new_tokens": max_new_tokens})
        call_index = len(self.calls) - 1

        if call_index == 0:
            return discovery_response([
                {"finding": "pneumothorax", "anatomical_location": "right apex", "certainty": "positive"},
                {"finding": "rib fracture", "anatomical_location": "right lateral 8th rib", "certainty": "positive"},
            ])

        localization_index = call_index - 1
        if localization_index == 0:
            return localization_response([{"box_2d": [100, 200, 150, 250], "label": "pneumothorax"}])
        if localization_index == 1:
            return localization_response([])
        if localization_index == 2:
            return localization_response([{"box_2d": [300, 400, 350, 450], "label": "rib fracture"}])
        return localization_response([])


class SingleFindingMockMedGemmaClient:
    def __init__(self) -> None:
        self.calls: list[dict[str, object]] = []

    def generate(self, *, images: Sequence[Image.Image], prompt: str, max_new_tokens: int) -> str:
        self.calls.append({"images": list(images), "prompt": prompt, "max_new_tokens": max_new_tokens})
        if len(self.calls) == 1:
            return discovery_response([
                {"finding": "pulmonary nodule", "anatomical_location": "right upper lobe", "certainty": "positive"}
            ])
        return localization_response([{"box_2d": [120, 610, 290, 790], "label": "pulmonary nodule"}])


class RealMedGemmaDiscoveryMockClient:
    def generate(self, *, images: Sequence[Image.Image], prompt: str, max_new_tokens: int) -> str:
        return REAL_MEDGEMMA_DISCOVERY_RESPONSE


class BroadFractureMockMedGemmaClient:
    def __init__(self) -> None:
        self.calls: list[dict[str, object]] = []

    def generate(self, *, images: Sequence[Image.Image], prompt: str, max_new_tokens: int) -> str:
        self.calls.append({"images": list(images), "prompt": prompt, "max_new_tokens": max_new_tokens})
        if len(self.calls) == 1:
            return discovery_response([
                {"finding": "fracture", "anatomical_location": "right humerus", "certainty": "positive"}
            ])
        return localization_response([{"box_2d": [100, 100, 700, 500], "label": "fracture"}])


class MalformedLocalizationMockMedGemmaClient:
    def __init__(self) -> None:
        self.calls: list[dict[str, object]] = []

    def generate(self, *, images: Sequence[Image.Image], prompt: str, max_new_tokens: int) -> str:
        self.calls.append({"images": list(images), "prompt": prompt, "max_new_tokens": max_new_tokens})
        if len(self.calls) == 1:
            return discovery_response([
                {"finding": "nodule", "anatomical_location": "right upper lobe", "certainty": "positive"}
            ])
        return "not valid json at all"


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------

class TestMultiViewDiscovery(unittest.TestCase):
    def test_discover_findings_normalizes_real_medgemma_response(self) -> None:
        detector = MedGemmaDetector(client=RealMedGemmaDiscoveryMockClient())  # type: ignore[arg-type]

        discovery = detector.discover_findings(
            [make_test_image(), make_test_image(color=(64, 64, 64))]
        )

        self.assertEqual(len(discovery.findings), 2)
        self.assertEqual(discovery.findings[0].finding, "nodule")
        self.assertEqual(discovery.findings[0].anatomical_location, "right upper lobe")
        self.assertEqual(discovery.findings[1].anatomical_location, "left upper lobe")

    def test_a_single_multi_view_discovery_call(self) -> None:
        mock_client = MockMedGemmaClient()
        detector = MedGemmaDetector(client=mock_client)  # type: ignore[arg-type]

        detector.detect(
            [make_test_image(), make_test_image(color=(64, 64, 64))],
            ["/study/0.png", "/study/1.png"],
        )

        discovery_calls = [c for c in mock_client.calls if len(c["images"]) > 1]  # type: ignore[arg-type]
        self.assertEqual(len(discovery_calls), 1)
        self.assertEqual(len(discovery_calls[0]["images"]), 2)  # type: ignore[arg-type]

    def test_b_localization_call_count(self) -> None:
        mock_client = MockMedGemmaClient()
        detector = MedGemmaDetector(client=mock_client)  # type: ignore[arg-type]

        detector.detect(
            [make_test_image(), make_test_image(color=(64, 64, 64))],
            ["/study/0.png", "/study/1.png"],
        )

        localization_calls = [c for c in mock_client.calls if len(c["images"]) == 1]  # type: ignore[arg-type]
        self.assertEqual(len(localization_calls), 4)

    def test_c_view_localization_metadata(self) -> None:
        mock_client = MockMedGemmaClient()
        detector = MedGemmaDetector(client=mock_client)  # type: ignore[arg-type]

        result, _ = detector.detect(
            [make_test_image(), make_test_image(color=(64, 64, 64))],
            ["/study/0.png", "/study/1.png"],
        )

        for finding in result.findings:
            self.assertEqual(len(finding.localizations), 2)
            for view in finding.localizations:
                self.assertIsInstance(view.image_index, int)
                self.assertTrue(view.image_path)
                self.assertIsInstance(view.status, str)
                self.assertIsInstance(view.boxes, list)
                self.assertIsInstance(view.candidate_boxes, list)

    def test_d_finding_visible_in_one_view_only(self) -> None:
        mock_client = MockMedGemmaClient()
        detector = MedGemmaDetector(client=mock_client)  # type: ignore[arg-type]

        result, _ = detector.detect(
            [make_test_image(), make_test_image(color=(64, 64, 64))],
            ["/study/0.png", "/study/1.png"],
        )

        pneumothorax = result.findings[0]
        view_0 = pneumothorax.localizations[0]
        view_1 = pneumothorax.localizations[1]

        self.assertEqual(view_0.status, "localized")
        self.assertEqual(len(view_0.boxes), 1)
        self.assertEqual(view_0.boxes[0].label, "pneumothorax")
        self.assertEqual(view_1.status, "abstained")
        self.assertEqual(view_1.boxes, [])

    def test_f_direct_image_mode(self) -> None:
        mock_client = SingleFindingMockMedGemmaClient()
        detector = MedGemmaDetector(client=mock_client)  # type: ignore[arg-type]

        result, processed_images = detector.detect(
            [make_test_image()],
            ["/single/0.png"],
        )

        discovery_calls = [
            c for c in mock_client.calls
            if len(c["images"]) == 1 and c is mock_client.calls[0]  # type: ignore[arg-type]
        ]
        self.assertEqual(len(discovery_calls), 1)
        self.assertEqual(len(mock_client.calls) - 1, 1)
        self.assertEqual(len(processed_images), 1)
        self.assertEqual(len(result.input_images), 1)
        self.assertEqual(len(result.findings[0].localizations), 1)

    def test_g_hidden_reference_not_passed_to_generate(self) -> None:
        mock_client = MockMedGemmaClient()
        detector = MedGemmaDetector(client=mock_client)  # type: ignore[arg-type]

        detector.detect(
            [make_test_image(), make_test_image(color=(64, 64, 64))],
            ["/study/0.png", "/study/1.png"],
        )

        for call in mock_client.calls:
            prompt = str(call["prompt"])
            self.assertNotIn("reference_report", prompt)
            self.assertNotIn("reference_findings", prompt)
            self.assertNotIn("Hidden IU X-Ray Reference", prompt)


class TestLocalizationGateStates(unittest.TestCase):
    def test_d_explicit_medgemma_abstention(self) -> None:
        mock_client = MockMedGemmaClient()
        detector = MedGemmaDetector(client=mock_client)  # type: ignore[arg-type]

        result, _ = detector.detect(
            [make_test_image(), make_test_image(color=(64, 64, 64))],
            ["/study/0.png", "/study/1.png"],
        )

        pneumothorax_view_1 = result.findings[0].localizations[1]
        self.assertEqual(pneumothorax_view_1.status, "abstained")
        self.assertEqual(pneumothorax_view_1.boxes, [])
        self.assertEqual(pneumothorax_view_1.candidate_boxes, [])

    def test_e_focal_candidate_fully_rejected(self) -> None:
        mock_client = BroadFractureMockMedGemmaClient()
        detector = MedGemmaDetector(client=mock_client)  # type: ignore[arg-type]

        result, _ = detector.detect([make_test_image()], ["/study/0.png"])

        view = result.findings[0].localizations[0]
        self.assertEqual(view.status, "rejected_by_quality_gate")
        self.assertEqual(view.boxes, [])
        self.assertEqual(len(view.candidate_boxes), 1)
        self.assertTrue(view.rejection_reasons)

    def test_f_parser_failure(self) -> None:
        mock_client = MalformedLocalizationMockMedGemmaClient()
        detector = MedGemmaDetector(client=mock_client)  # type: ignore[arg-type]

        result, _ = detector.detect([make_test_image()], ["/study/0.png"])

        view = result.findings[0].localizations[0]
        self.assertEqual(view.status, "parser_error")
        self.assertEqual(view.boxes, [])
        self.assertEqual(view.candidate_boxes, [])
        self.assertTrue(view.rejection_reasons)


class TestAnnotatedImageConsistency(unittest.TestCase):
    def test_e_annotated_image_contains_only_matching_view_boxes(self) -> None:
        findings = [
            LocalizedFinding(
                finding="pneumothorax",
                anatomical_location="right apex",
                certainty="positive",
                localizations=[
                    ViewLocalization(
                        image_index=0,
                        image_path="/study/0.png",
                        status="localized",
                        boxes=[LocalizationBox(box_2d=[100, 100, 200, 200], label="pneumothorax")],
                    ),
                    ViewLocalization(
                        image_index=1, image_path="/study/1.png", status="abstained", boxes=[]
                    ),
                ],
            ),
            LocalizedFinding(
                finding="rib fracture",
                anatomical_location="right lateral 8th rib",
                certainty="positive",
                localizations=[
                    ViewLocalization(
                        image_index=0, image_path="/study/0.png", status="abstained", boxes=[]
                    ),
                    ViewLocalization(
                        image_index=1,
                        image_path="/study/1.png",
                        status="localized",
                        boxes=[LocalizationBox(box_2d=[300, 300, 400, 400], label="rib fracture")],
                    ),
                ],
            ),
        ]

        blank = Image.new("RGB", (1000, 1000), (255, 255, 255))
        annotated_0 = draw_view_localizations(blank.copy(), 0, findings)
        annotated_1 = draw_view_localizations(blank.copy(), 1, findings)

        self.assertNotEqual(list(annotated_0.getdata()), list(blank.getdata()))
        self.assertNotEqual(list(annotated_1.getdata()), list(blank.getdata()))

        view_1_only = draw_view_localizations(blank.copy(), 1, findings[:1])
        self.assertEqual(list(view_1_only.getdata()), list(blank.getdata()))

        view_0_only_other = draw_view_localizations(blank.copy(), 0, findings[1:])
        self.assertEqual(list(view_0_only_other.getdata()), list(blank.getdata()))

    def test_g_visualizer_uses_accepted_boxes_only(self) -> None:
        accepted_box = LocalizationBox(box_2d=[100, 100, 200, 200], label="nodule")
        rejected_candidate = LocalizationBox(box_2d=[100, 100, 700, 500], label="fracture")
        findings = [
            LocalizedFinding(
                finding="nodule",
                anatomical_location="right upper lobe",
                certainty="positive",
                localizations=[
                    ViewLocalization(
                        image_index=0,
                        image_path="/study/0.png",
                        status="localized",
                        boxes=[accepted_box],
                        candidate_boxes=[accepted_box],
                    )
                ],
            ),
            LocalizedFinding(
                finding="fracture",
                anatomical_location="right humerus",
                certainty="positive",
                localizations=[
                    ViewLocalization(
                        image_index=0,
                        image_path="/study/0.png",
                        status="rejected_by_quality_gate",
                        boxes=[],
                        candidate_boxes=[rejected_candidate],
                        rejection_reasons=["bbox_area_ratio=0.2400 exceeds focal max_area_ratio=0.1200"],
                    )
                ],
            ),
        ]

        blank = Image.new("RGB", (1000, 1000), (255, 255, 255))
        annotated = draw_view_localizations(blank.copy(), 0, findings)
        self.assertNotEqual(list(annotated.getdata()), list(blank.getdata()))

        rejected_only = draw_view_localizations(blank.copy(), 0, [findings[1]])
        self.assertEqual(list(rejected_only.getdata()), list(blank.getdata()))


class TestPipelineResultSchema(unittest.TestCase):
    def test_pipeline_result_uses_input_images(self) -> None:
        result = PipelineResult(
            case_id="CXR3281_IM-1562",
            input_images=["/study/0.png", "/study/1.png"],
            findings=[],
        )
        dumped = result.model_dump()
        self.assertIn("input_images", dumped)
        self.assertNotIn("input_image", dumped)


if __name__ == "__main__":
    unittest.main()