#!/usr/bin/env python3
"""Compare multiple audited CV-Bench runs at the paired sample level."""

from __future__ import annotations

import argparse
import json
import math
from itertools import combinations
from pathlib import Path


def _json(path: Path) -> dict:
    return json.loads(path.read_text(encoding="utf-8"))


def _jsonl_by_id(path: Path) -> dict[str, dict]:
    rows = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line]
    result = {row["sample_id"]: row for row in rows}
    if len(result) != len(rows):
        raise AssertionError(f"Duplicate sample IDs: {path}")
    return result


def _mcnemar_exact_p(improved: int, regressed: int) -> float:
    discordant = improved + regressed
    if discordant == 0:
        return 1.0
    tail = sum(math.comb(discordant, k) for k in range(min(improved, regressed) + 1))
    return min(1.0, 2.0 * tail / (2**discordant))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--eval-root", type=Path, required=True)
    parser.add_argument("--experiment", default="stage2_covt_sam_vggt_covt4_v2")
    parser.add_argument("--steps", type=int, nargs="+", required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    runs: dict[int, dict] = {}
    results: dict[int, dict[str, dict]] = {}
    for step in args.steps:
        directory = args.eval_root / f"{args.experiment}-checkpoint-{step}"
        validation = _json(directory / "validation_report.json")
        if validation.get("assessment") != "ready_to_share":
            raise AssertionError(f"Run has not passed validation: checkpoint-{step}")
        metrics = _json(directory / "metrics.json")["metrics"]
        swanlab = _json(directory / "swanlab.json")
        results[step] = _jsonl_by_id(directory / "results.jsonl")
        runs[step] = {
            "correct": int(validation["correct_total"]),
            "overall_accuracy": metrics["eval/cvbench_accuracy"],
            "task_macro_accuracy": metrics["eval/cvbench_task_macro_accuracy"],
            "2d_accuracy": metrics["eval/cvbench_2d_accuracy"],
            "2d_count_accuracy": metrics["eval/cvbench_2d_count_accuracy"],
            "2d_relation_accuracy": metrics["eval/cvbench_2d_relation_accuracy"],
            "3d_accuracy": metrics["eval/cvbench_3d_accuracy"],
            "3d_depth_accuracy": metrics["eval/cvbench_3d_depth_accuracy"],
            "3d_distance_accuracy": metrics["eval/cvbench_3d_distance_accuracy"],
            "parse_rate": metrics["eval/cvbench_parse_rate"],
            "errors": int(metrics["eval/cvbench_errors"]),
            "swanlab_run_id": swanlab["run_id"],
            "swanlab_run_url": swanlab["run_url"],
        }

    reference_ids = set(results[args.steps[0]])
    if len(reference_ids) != 500 or any(set(results[step]) != reference_ids for step in args.steps):
        raise AssertionError("Checkpoint runs do not cover the same 500 samples")

    pairwise = []
    for earlier, later in combinations(args.steps, 2):
        improved = sum(
            not results[earlier][sample_id]["correct"] and results[later][sample_id]["correct"]
            for sample_id in reference_ids
        )
        regressed = sum(
            results[earlier][sample_id]["correct"] and not results[later][sample_id]["correct"]
            for sample_id in reference_ids
        )
        pairwise.append(
            {
                "earlier_step": earlier,
                "later_step": later,
                "accuracy_delta_percentage_points": 100
                * (runs[later]["overall_accuracy"] - runs[earlier]["overall_accuracy"]),
                "improved_samples": improved,
                "regressed_samples": regressed,
                "net_correct_delta": improved - regressed,
                "mcnemar_exact_p": _mcnemar_exact_p(improved, regressed),
            }
        )

    payload = {
        "assessment": "ready_to_share_with_subset_and_sample_size_caveat",
        "population": "fixed CV-Bench-500: 250 2D + 250 3D",
        "scoring": "greedy option exact-match; no LLM judge",
        "runs": {str(step): runs[step] for step in args.steps},
        "pairwise": pairwise,
        "best_overall_step": max(args.steps, key=lambda step: runs[step]["overall_accuracy"]),
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()