#!/usr/bin/env python3
"""Compare the frozen Frontier surface with conservative real capacities."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
import re
from collections import defaultdict
from datetime import datetime
from pathlib import Path
from typing import Any


RUN_PATTERN = re.compile(r"qwen30-prefill-real-tp(?P<tp>\d+)-mns(?P<mns>\d+)-")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--fleet-artifacts", type=Path, required=True)
    parser.add_argument(
        "--simulator-manifest", type=Path, action="append", required=True
    )
    parser.add_argument("--output-root", type=Path, required=True)
    return parser.parse_args()


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def load_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text())


def sign(value: float) -> int:
    return (value > 0) - (value < 0)


def kendall_tau_b(real: list[float], simulated: list[float]) -> dict[str, Any]:
    if len(real) != len(simulated):
        raise ValueError("ranking vectors have different lengths")
    concordant = discordant = real_only_ties = simulated_only_ties = both_ties = 0
    for left in range(len(real)):
        for right in range(left + 1, len(real)):
            real_sign = sign(real[left] - real[right])
            sim_sign = sign(simulated[left] - simulated[right])
            if real_sign == 0 and sim_sign == 0:
                both_ties += 1
            elif real_sign == 0:
                real_only_ties += 1
            elif sim_sign == 0:
                simulated_only_ties += 1
            elif real_sign == sim_sign:
                concordant += 1
            else:
                discordant += 1
    denominator = math.sqrt(
        (concordant + discordant + real_only_ties)
        * (concordant + discordant + simulated_only_ties)
    )
    tau = (concordant - discordant) / denominator if denominator else None
    return {
        "kendall_tau_b": tau,
        "concordant": concordant,
        "discordant": discordant,
        "real_only_ties": real_only_ties,
        "simulator_only_ties": simulated_only_ties,
        "both_ties": both_ties,
    }


def result_files_by_anchor(root: Path) -> dict[tuple[str, str], Path]:
    selected: dict[tuple[str, str], Path] = {}
    digests: dict[tuple[str, str], str] = {}
    for path in root.glob("artifacts/**/round*/results/r*.json"):
        if path.name.startswith("warmup_"):
            continue
        round_name = path.parent.parent.name
        key = (round_name, path.name)
        digest = sha256(path)
        if key in digests and digests[key] != digest:
            raise RuntimeError(
                f"conflicting duplicate real anchor {key} under {root}"
            )
        digests[key] = digest
        if key not in selected or len(path.parts) < len(selected[key].parts):
            selected[key] = path
    return selected


def find_real_runs(root: Path) -> dict[str, Path]:
    candidates: dict[str, list[Path]] = defaultdict(list)
    for path in root.iterdir():
        if not path.is_dir():
            continue
        match = RUN_PATTERN.search(path.name)
        if not match:
            continue
        name = f"tp{int(match.group('tp'))}_mns{int(match.group('mns'))}"
        candidates[name].append(path)
    selected: dict[str, Path] = {}
    for name, paths in candidates.items():
        complete = [
            path
            for path in paths
            if (path / "remote_run" / "exit_code").is_file()
            and (path / "remote_run" / "exit_code").read_text().strip() == "0"
        ]
        measured_counts = {
            path: len(result_files_by_anchor(path))
            for path in complete
        }
        if not measured_counts:
            raise RuntimeError(f"no successful run for {name}")
        maximum = max(measured_counts.values())
        richest = [path for path, count in measured_counts.items() if count == maximum]
        if len(richest) != 1:
            raise RuntimeError(f"ambiguous richest successful run for {name}: {richest}")
        selected[name] = richest[0]
    if len(selected) != 12:
        raise RuntimeError(f"expected 12 real configs, got {sorted(selected)}")
    return selected


def campaign_resources(root: Path) -> dict[str, Any]:
    runs = []
    gpu_hours = 0.0
    for path in sorted(root.iterdir()):
        match = RUN_PATTERN.search(path.name)
        exit_code = path / "remote_run" / "exit_code"
        started_at = path / "remote_run" / "started_at"
        finished_at = path / "remote_run" / "finished_at"
        if (
            not path.is_dir()
            or not match
            or not exit_code.is_file()
            or exit_code.read_text().strip() != "0"
            or not started_at.is_file()
            or not finished_at.is_file()
        ):
            continue
        started = datetime.fromisoformat(started_at.read_text().strip())
        finished = datetime.fromisoformat(finished_at.read_text().strip())
        duration_seconds = (finished - started).total_seconds()
        tp = int(match.group("tp"))
        run_gpu_hours = duration_seconds * tp / 3600.0
        gpu_hours += run_gpu_hours
        runs.append(
            {
                "run": path.name,
                "tp": tp,
                "duration_seconds": duration_seconds,
                "gpu_hours": run_gpu_hours,
            }
        )
    return {
        "successful_fleet_jobs": len(runs),
        "gpu_hours": gpu_hours,
        "runs": runs,
    }


def parse_real_config(name: str, run_root: Path) -> dict[str, Any]:
    result_files = sorted(result_files_by_anchor(run_root).values())
    if len(result_files) not in {10, 16}:
        raise RuntimeError(f"expected 10 base or 16 refined anchors for {name}, got {len(result_files)}")
    by_rate: dict[float, list[dict[str, Any]]] = defaultdict(list)
    for path in result_files:
        payload = load_json(path)
        rate = float(payload["workload"]["offered_request_rate"])
        by_rate[rate].append(
            {
                "path": str(path.resolve()),
                "sha256": sha256(path),
                "summary": payload["summary"],
            }
        )
    tp = int(name.split("_")[0][2:])
    base_rates = [4.0, 8.0, 16.0, 32.0, 64.0]
    refined_rates = sorted({*base_rates, *(tp * value for value in (5.0, 6.0, 7.0))})
    if tuple(sorted(by_rate)) not in {tuple(base_rates), tuple(refined_rates)}:
        raise RuntimeError(f"unexpected rate grid for {name}: {sorted(by_rate)}")
    anchors = []
    for rate, rounds in sorted(by_rate.items()):
        if len(rounds) != 2:
            raise RuntimeError(f"expected two rounds for {name}@{rate}, got {len(rounds)}")
        round_feasible = [bool(row["summary"]["slo"]["feasible"]) for row in rounds]
        anchors.append(
            {
                "rate": rate,
                "rounds": rounds,
                "conservative_feasible": all(round_feasible),
                "round_feasible": round_feasible,
                "round_ttft_p95_ms": [
                    float(row["summary"]["ttft_p95_ms"]) for row in rounds
                ],
            }
        )
    feasible = [row["rate"] for row in anchors if row["conservative_feasible"]]
    capacity = max(feasible, default=0.0)
    return {
        "name": name,
        "tp": tp,
        "mns": int(name.split("_mns")[1]),
        "anchors": anchors,
        "capacity": capacity,
        "capacity_per_gpu": capacity / tp,
        "source_run": str(run_root.resolve()),
    }


def parse_simulator(manifests: list[Path]) -> tuple[dict[str, Any], list[dict[str, str]]]:
    configs: dict[str, Any] = {}
    sources = []
    for path in manifests:
        payload = load_json(path)
        if payload["status"] not in {"complete", "partial_not_decision_bearing"}:
            raise RuntimeError(f"simulator manifest has invalid status: {path}")
        sources.append({"path": str(path.resolve()), "sha256": sha256(path)})
        result_by_name = {
            result["config"]["name"]: result for result in payload["config_results"]
        }
        for capacity in payload["capacity"]:
            name = capacity["config"]["name"]
            result = result_by_name.get(name)
            if result is None:
                raise RuntimeError(f"missing simulator config result {name}")
            entry = configs.setdefault(
                name,
                {
                    "name": name,
                    "tp": int(capacity["config"]["tp"]),
                    "mns": int(capacity["config"]["mns"]),
                    "anchor_by_rate": {},
                },
            )
            for load in result["loads"]:
                rate = float(load["offered_request_rate"])
                if rate in entry["anchor_by_rate"]:
                    raise RuntimeError(f"duplicate simulator anchor {name}@{rate}")
                entry["anchor_by_rate"][rate] = {
                    "rate": rate,
                    "feasible": bool(load["score"]["feasible"]),
                    "pass_rate": float(load["score"]["pass_rate"]),
                    "ttft_p95_ms": float(load["score"]["ttft_p95_ms"]),
                }
    if len(configs) != 12:
        raise RuntimeError(f"expected 12 simulator configs, got {sorted(configs)}")
    for entry in configs.values():
        entry["anchors"] = [
            entry["anchor_by_rate"][rate] for rate in sorted(entry["anchor_by_rate"])
        ]
        del entry["anchor_by_rate"]
        feasible = [row["rate"] for row in entry["anchors"] if row["feasible"]]
        entry["capacity"] = max(feasible, default=0.0)
        entry["capacity_per_gpu"] = entry["capacity"] / entry["tp"]
    return configs, sources


def compare(real: dict[str, Any], simulated: dict[str, Any]) -> dict[str, Any]:
    names = sorted(real, key=lambda name: (real[name]["tp"], real[name]["mns"]))
    if set(names) != set(simulated):
        raise RuntimeError("real and simulator config sets differ")
    real_scores = [real[name]["capacity_per_gpu"] for name in names]
    sim_scores = [simulated[name]["capacity_per_gpu"] for name in names]
    real_best = max(real_scores)
    sim_best = max(sim_scores)
    real_top = [name for name in names if real[name]["capacity_per_gpu"] == real_best]
    sim_top = [
        name for name in names if simulated[name]["capacity_per_gpu"] == sim_best
    ]
    worst_sim_choice = min(real[name]["capacity_per_gpu"] for name in sim_top)
    best_sim_choice = max(real[name]["capacity_per_gpu"] for name in sim_top)
    tau = kendall_tau_b(real_scores, sim_scores)

    pairwise = {"all": {"comparable": 0, "correct": 0}, "within_tp": {}}
    for left in range(len(names)):
        for right in range(left + 1, len(names)):
            real_sign = sign(real_scores[left] - real_scores[right])
            sim_sign = sign(sim_scores[left] - sim_scores[right])
            if real_sign:
                pairwise["all"]["comparable"] += 1
                pairwise["all"]["correct"] += int(real_sign == sim_sign)
            if real[names[left]]["tp"] == real[names[right]]["tp"] and real_sign:
                key = f"tp{real[names[left]]['tp']}"
                bucket = pairwise["within_tp"].setdefault(
                    key, {"comparable": 0, "correct": 0}
                )
                bucket["comparable"] += 1
                bucket["correct"] += int(real_sign == sim_sign)
    for bucket in [pairwise["all"], *pairwise["within_tp"].values()]:
        bucket["accuracy"] = (
            bucket["correct"] / bucket["comparable"]
            if bucket["comparable"]
            else None
        )

    confusion = {"real_pass_sim_pass": 0, "real_pass_sim_fail": 0,
                 "real_fail_sim_pass": 0, "real_fail_sim_fail": 0}
    for name in names:
        real_anchors = {row["rate"]: row for row in real[name]["anchors"]}
        sim_anchors = {row["rate"]: row for row in simulated[name]["anchors"]}
        if set(real_anchors) != set(sim_anchors):
            raise RuntimeError(f"anchor grids differ for {name}")
        for rate in real_anchors:
            real_pass = real_anchors[rate]["conservative_feasible"]
            sim_pass = sim_anchors[rate]["feasible"]
            key = f"real_{'pass' if real_pass else 'fail'}_sim_{'pass' if sim_pass else 'fail'}"
            confusion[key] += 1
    return {
        "config_order": names,
        "real_top_set": real_top,
        "simulator_top_set": sim_top,
        "top_set_exact_match": real_top == sim_top,
        "top_set_overlap": sorted(set(real_top) & set(sim_top)),
        "top1_regret_best": (real_best - best_sim_choice) / real_best,
        "top1_regret_worst": (real_best - worst_sim_choice) / real_best,
        "real_best_capacity_per_gpu": real_best,
        "simulator_best_capacity_per_gpu": sim_best,
        "kendall": tau,
        "pairwise_non_tied": pairwise,
        "anchor_confusion": confusion,
    }


def write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
    with path.open("w", newline="") as handle:
        writer = csv.DictWriter(
            handle, fieldnames=list(rows[0]), lineterminator="\n"
        )
        writer.writeheader()
        writer.writerows(rows)


def plot(path: Path, rows: list[dict[str, Any]], metrics: dict[str, Any]) -> None:
    import matplotlib.pyplot as plt
    import numpy as np

    labels = [f"TP{row['tp']}\nMNS{row['mns']}" for row in rows]
    x = np.arange(len(rows))
    width = 0.36
    figure, axes = plt.subplots(
        1, 2, figsize=(13.5, 5.0), gridspec_kw={"width_ratios": [3.3, 1.0]}
    )
    axis = axes[0]
    axis.bar(x - width / 2, [row["real"] for row in rows], width, label="Real vLLM")
    axis.bar(
        x + width / 2,
        [row["simulator"] for row in rows],
        width,
        label="Frontier profile-only",
    )
    axis.set_xticks(x, labels, fontsize=8)
    axis.set_ylabel("Max tested SLO-feasible request rate / GPU")
    axis.set_title("Qwen3-30B-A3B prefill-only: config ranking")
    axis.grid(axis="y", alpha=0.25)
    axis.legend(frameon=False, ncols=2)
    for separator in (3.5, 7.5):
        axis.axvline(separator, color="0.75", linewidth=0.8)
    tau = metrics["kendall"]["kendall_tau_b"]
    annotation = f"worst regret={metrics['top1_regret_worst'] * 100:.1f}%"
    annotation += (
        f"\nKendall τ-b={tau:.3f}"
        if tau is not None
        else "\nKendall τ-b=undefined"
    )
    axis.text(
        0.01,
        0.98,
        annotation,
        transform=axis.transAxes,
        va="top",
        fontsize=9,
        bbox={"facecolor": "white", "edgecolor": "0.8", "alpha": 0.9},
    )

    confusion = metrics["anchor_confusion"]
    matrix = np.array(
        [
            [confusion["real_pass_sim_pass"], confusion["real_pass_sim_fail"]],
            [confusion["real_fail_sim_pass"], confusion["real_fail_sim_fail"]],
        ]
    )
    image = axes[1].imshow(matrix, cmap="Blues", vmin=0)
    axes[1].set_xticks([0, 1], ["Sim pass", "Sim fail"])
    axes[1].set_yticks([0, 1], ["Real pass", "Real fail"])
    axes[1].set_title(f"{int(matrix.sum())} anchor decisions")
    for row in range(2):
        for column in range(2):
            axes[1].text(column, row, int(matrix[row, column]), ha="center", va="center")
    figure.colorbar(image, ax=axes[1], fraction=0.047, pad=0.04)
    figure.suptitle("ISL=2048, OSL=1, TTFT≤1256 ms, 95% pass gate; two real rounds")
    figure.tight_layout()
    figure.savefig(path, dpi=180, bbox_inches="tight")
    plt.close(figure)


def main() -> None:
    args = parse_args()
    real_runs = find_real_runs(args.fleet_artifacts.resolve())
    real = {name: parse_real_config(name, path) for name, path in real_runs.items()}
    simulated, simulator_sources = parse_simulator(
        [path.resolve() for path in args.simulator_manifest]
    )
    metrics = compare(real, simulated)
    rows = [
        {
            "config": name,
            "tp": real[name]["tp"],
            "mns": real[name]["mns"],
            "real": real[name]["capacity_per_gpu"],
            "simulator": simulated[name]["capacity_per_gpu"],
        }
        for name in metrics["config_order"]
    ]
    resources = campaign_resources(args.fleet_artifacts.resolve())
    resources["fresh_server_anchors"] = sum(
        len(config["anchors"]) * 2 for config in real.values()
    )
    resources["measured_requests"] = resources["fresh_server_anchors"] * 64
    resources["warmup_requests"] = sum(
        min(32, max(4, math.ceil(anchor["rate"] * 2.0))) * 2
        for config in real.values()
        for anchor in config["anchors"]
    )
    args.output_root.mkdir(parents=True, exist_ok=True)
    payload = {
        "schema": "qwen30-prefill-fidelity-comparison-v1",
        "objective": "maximum_tested_slo_feasible_offered_request_rate_per_gpu",
        "contract": {
            "model": "Qwen3-30B-A3B",
            "input_tokens": 2048,
            "output_tokens": 1,
            "prefix_caching": False,
            "ttft_slo_ms": 1256.0,
            "target_pass_rate": 0.95,
            "real_anchor_merge": "both_fresh_server_rounds_must_pass",
        },
        "metrics": metrics,
        "real_campaign_resources": resources,
        "real": real,
        "simulator": simulated,
        "simulator_sources": simulator_sources,
    }
    (args.output_root / "comparison.json").write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n"
    )
    write_csv(args.output_root / "capacity.csv", rows)
    plot(args.output_root / "qwen30-prefill-ranking.png", rows, metrics)
    print(json.dumps(metrics, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()