#!/usr/bin/env python3
"""Score a replacement-profile S2 sweep against the frozen real oracle."""

from __future__ import annotations

import argparse
import importlib.util
import json
import math
import sys
from pathlib import Path
from typing import Any


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--shard-metrics", type=Path, action="append", required=True)
    parser.add_argument("--ground-truth", type=Path, required=True)
    parser.add_argument("--historical-metrics", type=Path, required=True)
    parser.add_argument("--historical-analyzer", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def load_module(path: Path):
    spec = importlib.util.spec_from_file_location("simfid_s2_analyzer", path)
    if spec is None or spec.loader is None:
        raise ImportError(path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def ranking(scores: dict[str, float]) -> list[dict[str, Any]]:
    return [
        {"rank": index + 1, "cell": cell, "score": score}
        for index, (cell, score) in enumerate(
            sorted(scores.items(), key=lambda item: (-item[1], item[0]))
        )
    ]


def robust_loao(
    analyzer: Any,
    runs: list[dict[str, Any]],
    mode: str,
    reading: str,
    real_scores: dict[str, float],
) -> dict[str, Any]:
    """Historical LOAO with an explicit null range for all-tied concordance."""
    anchors = sorted(
        {(row["cell_id"], int(row["probe_index"])) for row in runs},
        key=lambda item: (analyzer.CELL_ORDER.index(item[0]), item[1]),
    )
    replicates = []
    undefined = []
    for cell, probe in anchors:
        scores, detail = analyzer.cell_scores(
            runs, mode, reading, removed=(cell, probe)
        )
        if scores is None:
            undefined.append(
                {"removed_cell_id": cell, "removed_probe_index": probe, **detail}
            )
            continue
        metric = analyzer.rank_metrics(real_scores, scores)
        replicates.append(
            {
                "removed_cell_id": cell,
                "removed_probe_index": probe,
                "top1_optimistic_regret": metric["top1"]["optimistic_regret"],
                "top1_worst_case_regret": metric["top1"]["worst_case_regret"],
                "top5_minimum_exact_five_overlap": metric["top5"]["minimum_exact_five_overlap"],
                "top5_maximum_exact_five_overlap": metric["top5"]["maximum_exact_five_overlap"],
                "top5_optimistic_regret": metric["top5"]["optimistic_regret"],
                "top5_worst_case_regret": metric["top5"]["worst_case_regret"],
                "tau_b": metric["kendall_tau_b"]["tau_b"],
                "pairwise_exact_sign_accuracy": metric["pairwise_direction"]["exact_sign_accuracy"],
                "pairwise_non_tied_concordance": metric["pairwise_direction"]["non_tied_concordance"],
                "trap_reproduced": metric["named_interactions"]["trap_reproduced"],
                "tp2_mns32_unique_global_best": metric["named_interactions"]["tp2_mns32_unique_global_best"],
            }
        )
    scalar_keys = (
        "top1_optimistic_regret",
        "top1_worst_case_regret",
        "top5_minimum_exact_five_overlap",
        "top5_maximum_exact_five_overlap",
        "top5_optimistic_regret",
        "top5_worst_case_regret",
        "tau_b",
        "pairwise_exact_sign_accuracy",
        "pairwise_non_tied_concordance",
    )
    ranges = {}
    for key in scalar_keys:
        values = [row[key] for row in replicates if row[key] is not None]
        ranges[key] = {
            "min": min(values) if values else None,
            "max": max(values) if values else None,
        }
    return {
        "replicate_count": len(anchors),
        "defined_replicates": len(replicates),
        "undefined_replicates": len(undefined),
        "undefined": undefined,
        "ranges": ranges,
        "trap_reproduced_count": sum(row["trap_reproduced"] for row in replicates),
        "tp2_mns32_unique_global_best_count": sum(
            row["tp2_mns32_unique_global_best"] for row in replicates
        ),
        "replicates": replicates,
        "range_semantics": "deterministic LOAO sensitivity ranges, not confidence intervals",
    }


def main() -> None:
    args = parse_args()
    analyzer = load_module(args.historical_analyzer)
    ground = json.loads(args.ground_truth.read_text())
    cells = {str(cell["cell_id"]): cell for cell in ground["cells"]}
    if set(cells) != set(analyzer.CELL_ORDER):
        raise ValueError("ground-truth cells differ from the frozen 3x4 surface")

    result_rows: list[dict[str, Any]] = []
    sources = []
    for path in args.shard_metrics:
        payload = json.loads(path.read_text())
        if payload["status"] != "PASS":
            raise ValueError(f"shard did not pass: {path}")
        sources.append({"path": str(path.resolve()), "runs": len(payload["results"])})
        result_rows.extend(payload["results"])
    if len(result_rows) != 92:
        raise ValueError(f"expected 92 replacement-profile probes, found {len(result_rows)}")

    seen: set[tuple[str, int]] = set()
    runs = []
    for row in result_rows:
        cell_id = str(row["cell"])
        probe_index = int(row["probe_index"])
        key = (cell_id, probe_index)
        if key in seen:
            raise ValueError(f"duplicate probe {key}")
        seen.add(key)
        real_probe = cells[cell_id]["probe_history"][probe_index]
        if not math.isclose(
            float(real_probe["sampling_u"]), float(row["sampling_u"]), rel_tol=0, abs_tol=1e-15
        ):
            raise ValueError(f"sampling_u mismatch for {key}")
        if int(real_probe["request_count"]) != int(row["selected_count"]):
            raise ValueError(f"request count mismatch for {key}")
        runs.append(
            {
                **row,
                "cell_id": cell_id,
                "mode": "vllm020-profile-only",
                "probe_index": probe_index,
                "sampling_u": float(row["sampling_u"]),
                "request_count": int(row["selected_count"]),
                "tensor_parallel_size": int(row["tensor_parallel_size"]),
                "real_anchor": {
                    "feasible": bool(real_probe["feasible"]),
                    "pass_rate": float(real_probe["pass_rate"]),
                    "request_count": int(real_probe["request_count"]),
                },
            }
        )

    historical = json.loads(args.historical_metrics.read_text())
    real_scores = {cell: float(value) for cell, value in historical["real_scores"].items()}
    analyses = {}
    for reading in analyzer.READINGS:
        scores, detail = analyzer.cell_scores(runs, "vllm020-profile-only", reading)
        if scores is None:
            raise ValueError(f"undefined {reading} scores: {detail}")
        analysis = {
            "reading": reading,
            "simulated_scores": scores,
            "ranking": ranking(scores),
            "cell_score_details": detail["cells"],
            "metrics": analyzer.rank_metrics(real_scores, scores),
            "loao": robust_loao(
                analyzer, runs, "vllm020-profile-only", reading, real_scores
            ),
        }
        if reading == "SLO-gated":
            analysis["false_feasibility"] = analyzer.false_feasibility(
                runs, "vllm020-profile-only"
            )
        analyses[reading] = analysis

    historical_modes = {}
    for label, key in (
        ("historical-profile-only", "uncalibrated/SLO-gated"),
        ("historical-per-tp-calibration", "frozen-calibrated/SLO-gated"),
    ):
        value = historical["analyses"][key]
        historical_modes[label] = {
            "simulated_scores": value["simulated_scores"],
            "ranking": ranking(value["simulated_scores"]),
            "metrics": value["metrics"],
            "false_feasibility": value["false_feasibility"],
        }

    output = {
        "schema": "frontier-qwen30-s2-profile-ablation.v1",
        "scope": {
            "cells": len(analyzer.CELL_ORDER),
            "probes": len(runs),
            "trace_horizon_seconds": 60,
            "calibration_a_tp": 1.0,
        },
        "sources": {
            "replacement_profile_shards": sources,
            "ground_truth": str(args.ground_truth.resolve()),
            "historical_metrics": str(args.historical_metrics.resolve()),
        },
        "real_scores": real_scores,
        "historical_modes": historical_modes,
        "vllm020_profile_only": analyses,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(output, indent=2, sort_keys=True) + "\n")
    print(args.output)


if __name__ == "__main__":
    main()
