from __future__ import annotations

import json
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]


class ModelComparisonTest(unittest.TestCase):
    def test_aligned_candidate_reports_paired_improvement(self) -> None:
        labels = [0, 0, 0, 0, 1, 1, 1, 1]
        baseline_scores = [0.1, 0.4, 0.3, 0.2, 0.3, 0.6, 0.5, 0.4]
        candidate_scores = [0.05, 0.10, 0.15, 0.20, 0.80, 0.85, 0.90, 0.95]
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            baseline = root / "baseline.jsonl"
            candidate = root / "candidate.jsonl"
            output = root / "comparison.json"
            for path, scores in ((baseline, baseline_scores), (candidate, candidate_scores)):
                path.write_text(
                    "".join(
                        json.dumps(
                            {
                                "record_id": f"row-{index}",
                                "group_id": f"group-{index}",
                                "label": label,
                                "probability": score,
                            }
                        )
                        + "\n"
                        for index, (label, score) in enumerate(zip(labels, scores, strict=True))
                    ),
                    encoding="utf-8",
                )
            completed = subprocess.run(
                [
                    sys.executable,
                    str(ROOT / "scripts/compare_models.py"),
                    "--baseline",
                    str(baseline),
                    "--candidate",
                    str(candidate),
                    "--output",
                    str(output),
                    "--fpr-budget",
                    "0.25",
                    "--bootstrap-samples",
                    "50",
                ],
                cwd=ROOT,
                check=False,
                capture_output=True,
                text=True,
            )
            self.assertEqual(completed.returncode, 0, completed.stderr)
            report = json.loads(output.read_text(encoding="utf-8"))
        self.assertEqual(report["count"], 8)
        self.assertEqual(report["group_count"], 8)
        self.assertGreater(report["candidate_minus_baseline_95ci"]["roc_auc"]["estimate"], 0.0)
        self.assertEqual(report["candidate"]["metrics"]["roc_auc"], 1.0)


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