tests/test_benchmark.py

from simpliscribe.benchmark import load_cases
from simpliscribe.benchmark import load_case_bundle
from simpliscribe.benchmark import evaluate_quality_gates
from simpliscribe.benchmark import run_case
from simpliscribe.benchmark import run_benchmark
from simpliscribe.benchmark import score_case

import pandas as pd
import pytest


def test_score_case_counts_matching_fields():
    case = {
        "id": "case-1",
        "label": "Simple case",
        "raw_text": "Paracetamol 650 tab od 5 days",
        "expected_medications": [
            {
                "name": "Paracetamol",
                "type": "Tablet",
                "dosage": "650 mg",
                "frequency": "once daily",
                "duration": "5 days",
            }
        ],
    }
    actual = [
        {
            "name": "Paracetamol",
            "type": "Tablet",
            "dosage": "650 mg",
            "frequency": "once daily",
            "duration": "5 days",
        }
    ]

    result = score_case(case, actual)

    assert result.matched_fields == 5
    assert result.total_fields == 5
    assert result.score == 1.0


def test_run_benchmark_returns_summary(tmp_path):
    cases_path = tmp_path / "cases.json"
    cases_path.write_text(
        """
        [
          {
            "id": "case-1",
            "label": "Simple once-daily tablet",
            "raw_text": "Paracetamol 650 tab od 5 days",
            "expected_medications": [
              {
                "name": "Paracetamol",
                "type": "Tablet",
                "dosage": "650 mg",
                "frequency": "once daily",
                "duration": "5 days"
              }
            ]
          }
        ]
        """.strip(),
        encoding="utf-8",
    )

    cases = load_cases(cases_path)
    result = run_benchmark(cases)

    assert result["case_count"] == 1
    assert result["success_count"] == 1
    assert result["failure_count"] == 0
    assert result["cases"][0]["label"] == "Simple once-daily tablet"
    assert result["cases"][0]["actual"][0]["frequency"] == "once daily"


def test_run_benchmark_preserves_capsule_and_bedtime_in_mixed_case(tmp_path):
        cases_path = tmp_path / "cases.json"
        cases_path.write_text(
                """
                [
                    {
                        "id": "case-2",
                        "label": "Two-line mixed medicines",
                        "raw_text": "Amoxycillin 500 cap bd 5 days\\nCetirizine 10 tab hs 3 days",
                        "expected_medications": [
                            {
                                "name": "Amoxycillin",
                                "type": "Capsule",
                                "dosage": "500 mg",
                                "frequency": "twice daily",
                                "duration": "5 days"
                            },
                            {
                                "name": "Cetirizine",
                                "type": "Tablet",
                                "dosage": "10 mg",
                                "frequency": "at bedtime",
                                "duration": "3 days"
                            }
                        ]
                    }
                ]
                """.strip(),
                encoding="utf-8",
        )

        cases = load_cases(cases_path)
        result = run_benchmark(cases)

        assert result["average_score"] == 1.0
        assert result["cases"][0]["actual"][0]["type"] == "Capsule"
        assert result["cases"][0]["actual"][1]["frequency"] == "at bedtime"


def test_run_case_returns_failed_score_when_inference_errors(monkeypatch):
    case = {
        "id": "case-error",
        "label": "Broken case",
        "raw_text": "Broken OCR text",
        "expected_medications": [
            {"name": "Paracetamol", "type": "Tablet", "dosage": "650 mg", "frequency": "once daily", "duration": "5 days"}
        ],
    }

    def fake_structure_medications(_: str):
        raise RuntimeError("Synthetic benchmark failure")

    monkeypatch.setattr("simpliscribe.benchmark.structure_medications", fake_structure_medications)

    result = run_case(case, retries=0)

    assert result.score == 0.0
    assert result.error == "Synthetic benchmark failure"


def test_run_case_supports_file_based_ocr(monkeypatch, tmp_path):
    image_path = tmp_path / "sample.png"
    image_path.write_bytes(b"mock-image")

    case = {
        "id": "case-file",
        "label": "File-based case",
        "file_path": "sample.png",
        "expected_medications": [
            {
                "name": "Paracetamol",
                "type": "Tablet",
                "dosage": "650 mg",
                "frequency": "once daily",
                "duration": "5 days",
            }
        ],
    }

    monkeypatch.setattr("simpliscribe.benchmark.extract_ocr_text", lambda path: "Paracetamol 650 tab od 5 days")

    result = run_case(case, base_dir=tmp_path, retries=0)

    assert result.source_kind == "file"
    assert result.source_path == "sample.png"
    assert result.raw_text == "Paracetamol 650 tab od 5 days"
    assert result.actual[0]["frequency"] == "once daily"


def test_load_cases_supports_parquet_ground_truth(tmp_path):
    parquet_path = tmp_path / "synthetic.parquet"
    dataframe = pd.DataFrame(
        [
            {
                "ground_truth": "<s_ocr> doctor_name: Dr. A medications: - Paracetamol 650 mg - Take once daily - Cetirizine 10 mg - At bedtime signature: Dr. A </s>"
            }
        ]
    )
    dataframe.to_parquet(parquet_path)

    cases = load_cases(parquet_path)

    assert len(cases) == 1
    assert cases[0]["raw_text"] == "Paracetamol 650 mg Take once daily\nCetirizine 10 mg At bedtime"
    assert cases[0]["expected_medications"][0]["name"] == "Paracetamol"
    assert cases[0]["expected_medications"][0]["dosage"] == "650 mg"
    assert cases[0]["expected_medications"][0]["frequency"] == "once daily"
    assert cases[0]["expected_medications"][1]["frequency"] == "at bedtime"


def test_load_cases_supports_parquet_ground_truth_without_signature(tmp_path):
    parquet_path = tmp_path / "synthetic_no_signature.parquet"
    dataframe = pd.DataFrame(
        [
            {
                "ground_truth": "<s_ocr> doctor_name: Dr. A medications: - Paracetamol 650 mg - Take once daily </s>"
            }
        ]
    )
    dataframe.to_parquet(parquet_path)

    cases = load_cases(parquet_path)

    assert len(cases) == 1
    assert cases[0]["expected_medications"] == [
        {
            "name": "Paracetamol",
            "type": "",
            "dosage": "650 mg",
            "frequency": "once daily",
            "duration": "N/A",
        }
    ]


def test_load_cases_preserves_hyphenated_medication_names(tmp_path):
    parquet_path = tmp_path / "synthetic_hyphenated.parquet"
    dataframe = pd.DataFrame(
        [
            {
                "ground_truth": "<s_ocr> medications: - Co-trimoxazole 500/125 mg - Take twice daily signature: Dr. A </s>"
            }
        ]
    )
    dataframe.to_parquet(parquet_path)

    cases = load_cases(parquet_path)

    assert len(cases) == 1
    assert cases[0]["expected_medications"][0]["name"] == "Co-Trimoxazole"
    assert cases[0]["expected_medications"][0]["dosage"] == "500/125 mg"
    assert cases[0]["expected_medications"][0]["frequency"] == "twice daily"


def test_load_cases_supports_complex_dosage_formats(tmp_path):
    parquet_path = tmp_path / "synthetic_complex_dosage.parquet"
    dataframe = pd.DataFrame(
        [
            {
                "ground_truth": "<s_ocr> medications: - Ambroxol 5 mg/5 ml - As directed signature: Dr. A </s>"
            }
        ]
    )
    dataframe.to_parquet(parquet_path)

    cases = load_cases(parquet_path)

    assert len(cases) == 1
    assert cases[0]["expected_medications"][0]["name"] == "Ambroxol"
    assert cases[0]["expected_medications"][0]["dosage"] == "5 mg/5 ml"
    assert cases[0]["expected_medications"][0]["frequency"] == "Refer to prescription"


def test_score_case_penalizes_extra_medications():
    expected = {
        "name": "Cetirizine",
        "type": "Tablet",
        "dosage": "10 mg",
        "frequency": "once daily",
        "duration": "5 days",
    }
    extra = {**expected, "name": "Paracetamol"}

    result = score_case({"id": "extra-drug", "expected_medications": [expected]}, [expected, extra])

    assert result.matched_fields == 5
    assert result.total_fields == 10
    assert result.score == 0.5
    assert result.actual_count == 2


def test_score_case_does_not_score_unspecified_expected_fields():
    case = {"id": "name-only", "expected_medications": [{"name": "Cetirizine"}]}
    actual = [{
        "name": "Cetirizine",
        "type": "Tablet",
        "dosage": "10 mg",
        "frequency": "once daily",
        "duration": "5 days",
    }]

    result = score_case(case, actual)

    assert result.matched_fields == 1
    assert result.total_fields == 1
    assert result.score == 1.0


def test_score_case_matches_medications_independent_of_order():
    cetirizine = {"name": "Cetirizine", "dosage": "10 mg", "frequency": "at bedtime"}
    paracetamol = {"name": "Paracetamol", "dosage": "650 mg", "frequency": "once daily"}

    result = score_case(
        {"id": "reordered", "expected_medications": [cetirizine, paracetamol]},
        [paracetamol, cetirizine],
    )

    assert result.matched_fields == 6
    assert result.total_fields == 6
    assert result.score == 1.0
    assert result.field_results[0]["actual_medication_index"] == 1


def test_load_versioned_golden_bundle(tmp_path):
    path = tmp_path / "golden.json"
    path.write_text(
        '{"schema_version":"1.0","dataset":"test-set","cases":['
        '{"id":"case-1","raw_text":"Example","expected_medications":[]}]}',
        encoding="utf-8",
    )

    metadata, cases = load_case_bundle(path)

    assert metadata["schema_version"] == "1.0"
    assert metadata["dataset"] == "test-set"
    assert cases[0]["id"] == "case-1"


def test_load_versioned_golden_bundle_rejects_duplicate_ids(tmp_path):
    path = tmp_path / "golden.json"
    path.write_text(
        '{"schema_version":"1.0","cases":['
        '{"id":"duplicate","expected_medications":[]},'
        '{"id":"duplicate","expected_medications":[]}]}',
        encoding="utf-8",
    )

    with pytest.raises(ValueError, match="Duplicate benchmark case id"):
        load_case_bundle(path)


def test_run_benchmark_reports_detection_field_and_safety_metrics(monkeypatch):
    cases = [
        {
            "id": "positive",
            "raw_text": "positive",
            "expected_medications": [{"name": "A", "dosage": "10 mg", "requires_review": True}],
        },
        {"id": "negative", "raw_text": "negative", "expected_medications": []},
        {"id": "unreadable", "raw_text": "unreadable", "expected_rejection": True, "expected_medications": []},
    ]
    outputs = {
        "positive": {"medications": [{"name": "A", "dosage": "10 mg", "requires_review": True}]},
        "negative": {"medications": [{"name": "Hallucination"}]},
        "unreadable": {"medications": []},
    }
    monkeypatch.setattr("simpliscribe.benchmark.structure_medications", lambda text: outputs[text])

    result = run_benchmark(cases)
    detection = result["metrics"]["medication_detection"]

    assert detection == {
        "true_positives": 1,
        "false_positives": 1,
        "false_negatives": 0,
        "precision": 0.5,
        "recall": 1.0,
        "f1": 0.6667,
        "hallucination_rate": 0.5,
    }
    assert result["metrics"]["field_accuracy"]["dosage"]["accuracy"] == 1.0
    assert result["metrics"]["review_flag_recall"] == 1.0
    assert result["metrics"]["unreadable_rejection_rate"] == 1.0
    assert evaluate_quality_gates(result, 0.7, 0.4)