#!/usr/bin/env python3
"""Compare one complete Qwen30 Frontier latency surface with real vLLM."""

from __future__ import annotations

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


CONFIGS = tuple(f"tp{tp}_mns{mns}" for tp in (1, 2, 4) for mns in (8, 16, 32, 64))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--sim-root", type=Path, required=True)
    parser.add_argument("--real-audit", type=Path, required=True)
    parser.add_argument("--json-output", type=Path, required=True)
    parser.add_argument("--markdown-output", type=Path, required=True)
    return parser.parse_args()


def number(value: Any, field: str) -> float:
    if not isinstance(value, (int, float)) or isinstance(value, bool):
        raise ValueError(f"{field} is not numeric: {value!r}")
    result = float(value)
    if not math.isfinite(result) or result < 0:
        raise ValueError(f"{field} is invalid: {value!r}")
    return result


def parse_config(name: str) -> tuple[int, int]:
    tp, mns = name.split("_", 1)
    return int(tp.removeprefix("tp")), int(mns.removeprefix("mns"))


def rank(values: dict[str, float]) -> list[str]:
    return [name for name, _ in sorted(values.items(), key=lambda item: (item[1], item[0]))]


def pairwise(sim: dict[str, float], real: dict[str, float]) -> dict[str, int | float | None]:
    concordant = discordant = ties = 0
    for left, right in combinations(CONFIGS, 2):
        sim_delta = sim[left] - sim[right]
        real_delta = real[left] - real[right]
        if math.isclose(sim_delta, 0.0, abs_tol=1e-9) or math.isclose(real_delta, 0.0, abs_tol=1e-9):
            ties += 1
        elif (sim_delta > 0) == (real_delta > 0):
            concordant += 1
        else:
            discordant += 1
    informative = concordant + discordant
    return {
        "concordant_pairs": concordant,
        "discordant_pairs": discordant,
        "tied_pairs": ties,
        "informative_pairs": informative,
        "agreement": concordant / informative if informative else None,
    }


def main() -> None:
    args = parse_args()
    root = args.sim_root.resolve()
    audit = json.loads(args.real_audit.read_text())
    metrics = tuple(audit.get("applicable_metrics") or [])
    if not metrics or set(audit.get("configs") or {}) != set(CONFIGS):
        raise ValueError("real audit is incomplete")
    expected_requests = int(audit["trace_manifests"]["tp1"]["requests"])
    cells: dict[str, Any] = {}
    for name in CONFIGS:
        tp, mns = parse_config(name)
        path = root / "runs" / name / f"tp{tp}" / "result.json"
        result = json.loads(path.read_text())
        if result.get("status") != "completed" or result.get("request_count") != expected_requests:
            raise ValueError(f"incomplete simulator result: {path}")
        if result.get("config") != {"tp": tp, "mns": mns, "name": name}:
            raise ValueError(f"simulator config drift: {path}")
        score = result.get("score")
        if not isinstance(score, dict):
            raise ValueError(f"simulator score missing: {path}")
        values: dict[str, Any] = {}
        for metric in metrics:
            prefix = metric.removesuffix("_ms")
            values[metric] = {
                statistic: number(score.get(f"{prefix}_{statistic}_ms"), f"{path}:{metric}:{statistic}")
                for statistic in ("mean", "p90")
            }
        cells[name] = {"path": str(path), "metrics": values}
    selection: dict[str, Any] = {}
    for metric in metrics:
        for statistic in ("mean", "p90"):
            sim_values = {name: cells[name]["metrics"][metric][statistic] for name in CONFIGS}
            real_values = {
                name: number(
                    audit["configs"][name]["metrics"][metric][f"pooled_{statistic}_ms"],
                    f"real:{name}:{metric}:{statistic}",
                )
                for name in CONFIGS
            }
            sim_ranking, real_ranking = rank(sim_values), rank(real_values)
            sim_choice, real_best = sim_ranking[0], real_ranking[0]
            selection[f"{metric}:{statistic}"] = {
                "sim_winner": sim_choice,
                "real_winner": real_best,
                "winner_match": sim_choice == real_best,
                "selected_config_real_regret": (real_values[sim_choice] - real_values[real_best]) / real_values[real_best],
                "sim_ranking": sim_ranking,
                "real_ranking": real_ranking,
                **pairwise(sim_values, real_values),
            }
    payload = {
        "schema": "qwen30-latency-case-frontier-real-comparison-v1",
        "sim_root": str(root),
        "real_audit": str(args.real_audit.resolve()),
        "prefill_only": bool(audit["prefill_only"]),
        "applicable_metrics": metrics,
        "selection": selection,
        "cells": cells,
    }
    lines = ["# Qwen30 Frontier vs real latency selection", "", "| Objective | Frontier | Real | Match | Regret | Pairwise |", "|---|---|---|---:|---:|---:|"]
    for objective, row in selection.items():
        agreement = "N/A" if row["agreement"] is None else f"{row['agreement']:.1%}"
        lines.append(f"| {objective} | {row['sim_winner']} | {row['real_winner']} | {'yes' if row['winner_match'] else 'no'} | {row['selected_config_real_regret']:.1%} | {agreement} |")
    args.json_output.parent.mkdir(parents=True, exist_ok=True)
    args.markdown_output.parent.mkdir(parents=True, exist_ok=True)
    args.json_output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    args.markdown_output.write_text("\n".join(lines) + "\n")
    print(json.dumps(selection, sort_keys=True))


if __name__ == "__main__":
    main()
