#!/usr/bin/env python3
"""Compare pooled real and Frontier Qwen235 latency-selection surfaces."""

from __future__ import annotations

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


CONFIGS = ("tp4_ep1_mns64", "tp4_ep1_mns128", "tp8_ep8_mns64", "tp8_ep8_mns128")
CASES = ("fixed-pd", "fixed-po", "trace-pd", "trace-po")


def percentile(values: list[float], fraction: float):
    if not values:
        return None
    return sorted(values)[math.ceil(len(values) * fraction) - 1]


def summarize(values: list[float]):
    if not values:
        return {"mean": None, "p90": None}
    return {"mean": sum(values) / len(values), "p90": percentile(values, 0.9)}


def real_surface(root: Path, case: str):
    surface = {}
    for config in CONFIGS:
        requests = []
        trials = []
        for trial in (1, 2, 3):
            path = root / "real" / case / "real" / config / f"trial{trial}" / "results" / "result.json"
            payload = json.loads(path.read_text())
            summary = payload["summary"]
            if summary["failed"] or summary["completed"] != len(payload["requests"]):
                raise ValueError(f"invalid real result: {path}")
            if any(not row["success"] for row in payload["requests"]):
                raise ValueError(f"failed request: {path}")
            requests.extend(payload["requests"])
            trials.append(str(path.resolve()))
        ttft = summarize([float(row["ttft_ms"]) for row in requests])
        tpot = summarize([float(row["tpot_ms"]) for row in requests if row["tpot_ms"] is not None])
        e2e = summarize([float(row["e2e_ms"]) for row in requests])
        surface[config] = {
            "ttft_mean_ms": ttft["mean"], "ttft_p90_ms": ttft["p90"],
            "tpot_mean_ms": tpot["mean"], "tpot_p90_ms": tpot["p90"],
            "e2e_mean_ms": e2e["mean"], "e2e_p90_ms": e2e["p90"],
            "request_samples": len(requests), "trials": trials,
        }
    return surface


def sim_surface(root: Path, case: str):
    payload = json.loads((root / "sim" / case / "frontier_surface.json").read_text())
    surface = {}
    for result in payload["results"]:
        if result["status"] != "completed":
            continue
        config = result["config"]["name"]
        if config in surface:
            raise ValueError(f"duplicate simulator config: {config}")
        surface[config] = {key: value for key, value in result["metrics"].items() if key.endswith("_ms")}
    if set(surface) != set(CONFIGS):
        raise ValueError(f"simulator coverage failure for {case}: {sorted(surface)}")
    return surface


def compare_metric(real: dict, sim: dict, metric: str):
    applicable = [config for config in CONFIGS if real[config][metric] is not None and sim[config][metric] is not None]
    real_winner = min(applicable, key=lambda config: (real[config][metric], config))
    sim_winner = min(applicable, key=lambda config: (sim[config][metric], config))
    regret = real[sim_winner][metric] / real[real_winner][metric] - 1
    informative = agreement = 0
    reversals = []
    for left, right in itertools.combinations(applicable, 2):
        real_direction = (real[left][metric] > real[right][metric]) - (real[left][metric] < real[right][metric])
        sim_direction = (sim[left][metric] > sim[right][metric]) - (sim[left][metric] < sim[right][metric])
        if real_direction and sim_direction:
            informative += 1
            agreement += int(real_direction == sim_direction)
            if real_direction != sim_direction:
                reversals.append([left, right])
    return {
        "real_winner": real_winner,
        "sim_winner": sim_winner,
        "winner_match": real_winner == sim_winner,
        "selected_real_regret": regret,
        "informative_pairs": informative,
        "pair_direction_agreement": agreement / informative if informative else None,
        "reversals": reversals,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--campaign-root", type=Path, required=True)
    parser.add_argument("--json-output", type=Path, required=True)
    parser.add_argument("--markdown-output", type=Path, required=True)
    args = parser.parse_args()
    cases = {}
    lines = ["# Qwen235 vLLM 0.20 Frontier vs real", "", "| case | metric | Frontier winner | real winner | match | regret | pair agreement |", "|---|---|---|---|---:|---:|---:|"]
    for case in CASES:
        real = real_surface(args.campaign_root, case)
        sim = sim_surface(args.campaign_root, case)
        metrics = ["ttft_mean_ms", "ttft_p90_ms", "e2e_mean_ms", "e2e_p90_ms"]
        if case.endswith("pd"):
            metrics[2:2] = ["tpot_mean_ms", "tpot_p90_ms"]
        comparisons = {}
        for metric in metrics:
            item = compare_metric(real, sim, metric)
            comparisons[metric] = item
            lines.append(
                f"| {case} | {metric} | {item['sim_winner']} | {item['real_winner']} | "
                f"{'yes' if item['winner_match'] else 'no'} | {item['selected_real_regret']:.1%} | "
                f"{item['pair_direction_agreement']:.1%} |"
            )
        cases[case] = {"real": real, "sim": sim, "comparison": comparisons}
    payload = {"schema": "qwen235-v020-simulator-real-comparison-v1", "cases": cases}
    args.json_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")


if __name__ == "__main__":
    main()