#!/usr/bin/env python3
"""Compare a graph-aligned Frontier Qwen30 surface with audited real vLLM data.

The input surface may have been produced by independent TP jobs.  Therefore
this script deliberately reads the per-cell ``result.json`` files rather than
the runner's root manifest, which is only a convenience artifact and can be
overwritten by concurrent dispatch.
"""

from __future__ import annotations

import argparse
import hashlib
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))
METRICS = ("ttft_ms", "tpot_ms", "e2e_ms")
STATISTICS = ("mean", "p90")


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 read_json(path: Path) -> dict[str, Any]:
    with path.open(encoding="utf-8") as source:
        payload = json.load(source)
    if not isinstance(payload, dict):
        raise ValueError(f"expected JSON object: {path}")
    return payload


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


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


def parse_config(name: str) -> tuple[int, int]:
    try:
        tp_part, mns_part = name.split("_", 1)
        return int(tp_part.removeprefix("tp")), int(mns_part.removeprefix("mns"))
    except ValueError as error:
        raise ValueError(f"invalid config name: {name}") from error


def load_sim_cell(root: Path, name: str) -> dict[str, Any]:
    tp, mns = parse_config(name)
    path = root / "runs" / name / f"tp{tp}" / "result.json"
    result = read_json(path)
    if result.get("status") != "completed":
        raise ValueError(f"{path}: status is not completed")
    config = result.get("config")
    if config != {"tp": tp, "mns": mns, "name": name}:
        raise ValueError(f"{path}: config drift: {config!r}")
    if result.get("trace_label") != f"tp{tp}":
        raise ValueError(f"{path}: trace label does not match TP")
    if result.get("request_count") != 129:
        raise ValueError(f"{path}: expected 129 requests")
    score = result.get("score")
    if not isinstance(score, dict):
        raise ValueError(f"{path}: missing score")
    metrics: dict[str, dict[str, float]] = {}
    for metric in METRICS:
        prefix = metric.removesuffix("_ms")
        metrics[metric] = {
            statistic: numeric(score.get(f"{prefix}_{statistic}_ms"), f"{path}:{metric}:{statistic}")
            for statistic in STATISTICS
        }
    return {
        "config": name,
        "tp": tp,
        "mns": mns,
        "trace_label": result["trace_label"],
        "trace_sha256": result.get("trace_sha256"),
        "offered_request_rate": numeric(result.get("offered_request_rate"), f"{path}:rate"),
        "offered_request_rate_per_gpu": numeric(
            result.get("offered_request_rate_per_gpu"), f"{path}:rate_per_gpu"
        ),
        "request_count": result["request_count"],
        "elapsed_seconds": numeric(result.get("elapsed_seconds"), f"{path}:elapsed"),
        "result_path": str(path.resolve()),
        "result_sha256": sha256_file(path),
        "metrics": metrics,
    }


def load_real_cell(audit: dict[str, Any], name: str) -> dict[str, dict[str, float]]:
    configs = audit.get("configs")
    if not isinstance(configs, dict) or name not in configs:
        raise ValueError(f"real audit lacks {name}")
    result = configs[name]
    metrics = result.get("metrics")
    if not isinstance(metrics, dict):
        raise ValueError(f"real audit {name} lacks metrics")
    parsed: dict[str, dict[str, float]] = {}
    for metric in METRICS:
        values = metrics.get(metric)
        if not isinstance(values, dict):
            raise ValueError(f"real audit {name} lacks {metric}")
        parsed[metric] = {
            "mean": numeric(values.get("pooled_mean_ms"), f"real:{name}:{metric}:mean"),
            "p90": numeric(values.get("pooled_p90_ms"), f"real:{name}:{metric}:p90"),
        }
    return parsed


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


def pairwise_agreement(sim: dict[str, float], real: dict[str, float]) -> dict[str, int]:
    concordant = discordant = ties = 0
    for left, right in combinations(sorted(sim), 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
    return {
        "concordant_pairs": concordant,
        "discordant_pairs": discordant,
        "tied_pairs": ties,
        "informative_pairs": concordant + discordant,
    }


def comparison_summary(cells: dict[str, dict[str, Any]], real_audit: dict[str, Any]) -> dict[str, Any]:
    summaries: dict[str, Any] = {}
    for metric in METRICS:
        for statistic in STATISTICS:
            sim_values = {
                name: cells[name]["metrics"][metric][f"sim_{statistic}_ms"]
                for name in CONFIGS
            }
            real_values = {
                name: load_real_cell(real_audit, name)[metric][statistic] for name in CONFIGS
            }
            sim_ranking = ranking(sim_values)
            real_ranking = ranking(real_values)
            pairwise = pairwise_agreement(sim_values, real_values)
            summaries[f"{metric}:{statistic}"] = {
                "sim_winner": sim_ranking[0],
                "real_winner": real_ranking[0],
                "winner_match": sim_ranking[0] == real_ranking[0],
                "sim_ranking": sim_ranking,
                "real_ranking": real_ranking,
                **pairwise,
                "pairwise_agreement_fraction": (
                    pairwise["concordant_pairs"] / pairwise["informative_pairs"]
                    if pairwise["informative_pairs"]
                    else None
                ),
            }
    return summaries


def render_markdown(payload: dict[str, Any]) -> str:
    lines = [
        "# Frontier piecewise graph-profile vs. real vLLM",
        "",
        "Each cell uses its TP-normalized, 129-request trace. Frontier values are one deterministic simulation; "
        "real values pool three fresh-server trials (387 requests/cell).",
        "",
        "| Config | TTFT sim / real mean (ms) | TPOT sim / real mean (ms) | E2E sim / real mean (ms) |",
        "|---|---:|---:|---:|",
    ]
    for name in CONFIGS:
        cell = payload["cells"][name]
        entries = []
        for metric in METRICS:
            values = cell["metrics"][metric]
            entries.append(f"{values['sim_mean_ms']:.1f} / {values['real_mean_ms']:.1f}")
        lines.append(f"| {name} | " + " | ".join(entries) + " |")

    lines.extend(["", "## Selection agreement", "", "| Target | Frontier winner | Real winner | Match | Pairwise agreement |", "|---|---|---|---:|---:|"])
    for target, summary in payload["selection"].items():
        agreement = summary["pairwise_agreement_fraction"]
        agreement_text = "n/a" if agreement is None else f"{agreement:.1%}"
        lines.append(
            f"| {target} | {summary['sim_winner']} | {summary['real_winner']} | "
            f"{'yes' if summary['winner_match'] else 'no'} | {agreement_text} "
            f"({summary['concordant_pairs']}/{summary['informative_pairs']}) |"
        )
    lines.append("")
    return "\n".join(lines)


def main() -> None:
    args = parse_args()
    sim_root = args.sim_root.resolve()
    real_path = args.real_audit.resolve()
    real_audit = read_json(real_path)
    if set(real_audit.get("configs", {})) != set(CONFIGS):
        raise ValueError("real audit config set does not match the 12-cell surface")

    cells: dict[str, dict[str, Any]] = {}
    for name in CONFIGS:
        sim = load_sim_cell(sim_root, name)
        real = load_real_cell(real_audit, name)
        metrics = {}
        for metric in METRICS:
            metrics[metric] = {
                f"sim_{statistic}_ms": sim["metrics"][metric][statistic]
                for statistic in STATISTICS
            }
            metrics[metric].update(
                {
                    f"real_{statistic}_ms": real[metric][statistic]
                    for statistic in STATISTICS
                }
            )
            metrics[metric].update(
                {
                    f"{statistic}_ratio_sim_over_real": sim["metrics"][metric][statistic]
                    / real[metric][statistic]
                    for statistic in STATISTICS
                }
            )
        cells[name] = {key: value for key, value in sim.items() if key != "metrics"}
        cells[name]["metrics"] = metrics

    payload = {
        "schema": "frontier-qwen30-piecewise-graph-comparison-v1",
        "contract": {
            "configs": list(CONFIGS),
            "request_count_per_cell": 129,
            "real_trials_per_cell": 3,
            "trace_policy": "TP-normalized arrivals with full 16-token prefix blocks only",
            "graph_semantics": "Frontier piecewise; CUDA_EVENT for prefill/mixed and KERNEL_ONLY for captured decode",
        },
        "sim_root": str(sim_root),
        "real_audit": str(real_path),
        "real_audit_sha256": sha256_file(real_path),
        "cells": cells,
        "selection": comparison_summary(cells, real_audit),
    }
    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", encoding="utf-8")
    args.markdown_output.write_text(render_markdown(payload), encoding="utf-8")
    print(json.dumps(payload["selection"], indent=2, sort_keys=True))


if __name__ == "__main__":
    main()