#!/usr/bin/env python3
"""Validate and compare frozen T0 Frontier and two-round real surfaces."""

from __future__ import annotations

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


SLOS = ("tpot_40ms", "tpot_120ms", "tpot_150ms", "tpot_180ms")
RATE_LATTICE = (0.10, 0.20, 0.40, 0.80, 1.20, 1.60, 2.40, 3.20)


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def write_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")


def rate_key(rate: float) -> str:
    return f"r{rate:.2f}".replace(".", "p")


def find_real_result(roots: list[Path], name: str, round_id: int, rate: float) -> Path:
    relative = Path(name) / f"round{round_id}/results" / f"{rate_key(rate)}.json"
    matches = [root / relative for root in roots if (root / relative).is_file()]
    if len(matches) != 1:
        raise ValueError(f"expected one real result for {relative}, got {matches}")
    return matches[0]


def capacity(loads: list[dict[str, Any]], slo: str, field: str) -> float | None:
    values = [float(load["rate"]) for load in loads if bool(load[field][slo]["feasible"])]
    return max(values) if values else None


def real_boundary_status(loads: list[dict[str, Any]], slo: str) -> str:
    labels = {
        float(load["rate"]): bool(load["real_conservative"][slo]["feasible"])
        for load in loads
    }
    ordered = [(rate, labels[rate]) for rate in RATE_LATTICE if rate in labels]
    if any(not left and right for (_, left), (_, right) in zip(ordered, ordered[1:])):
        return "non_monotonic_requires_full_lattice"
    if len(labels) == len(RATE_LATTICE):
        return "complete_lattice"
    for lower, upper in zip(RATE_LATTICE, RATE_LATTICE[1:]):
        if labels.get(lower) is True and labels.get(upper) is False:
            return "adjacent_transition_bracketed"
    if labels.get(RATE_LATTICE[-1]) is True:
        return "upper_lattice_reached"
    if labels.get(RATE_LATTICE[0]) is False and not any(labels.values()):
        return "lowest_anchor_infeasible"
    return "unbracketed_requires_expansion"


def kendall_tau_b(left: list[float], right: list[float]) -> float | None:
    concordant = discordant = left_ties = right_ties = 0
    for i in range(len(left)):
        for j in range(i + 1, len(left)):
            x = (left[i] > left[j]) - (left[i] < left[j])
            y = (right[i] > right[j]) - (right[i] < right[j])
            if x == 0 and y == 0:
                continue
            if x == 0:
                left_ties += 1
            elif y == 0:
                right_ties += 1
            elif x == y:
                concordant += 1
            else:
                discordant += 1
    denominator = math.sqrt(
        (concordant + discordant + left_ties)
        * (concordant + discordant + right_ties)
    )
    return (concordant - discordant) / denominator if denominator else None


def sign(left: float, right: float) -> int:
    return (left > right) - (left < right)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--frontier-freeze", type=Path, required=True)
    parser.add_argument("--real-plan", type=Path, required=True)
    parser.add_argument("--real-root", type=Path, action="append", required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    args = parser.parse_args()

    freeze = json.loads(args.frontier_freeze.read_text())
    plan = json.loads(args.real_plan.read_text())
    real_roots = [root.resolve() for root in args.real_root]
    if freeze.get("status") != "frozen_before_real_surface" or len(freeze.get("config_results") or []) != 8:
        raise ValueError("Frontier freeze is incomplete")
    if plan.get("frontier_freeze", {}).get("sha256") != sha256(args.frontier_freeze):
        raise ValueError("real plan does not point to this Frontier freeze")
    sim_by_name = {item["config"]["name"]: item for item in freeze["config_results"]}

    cells = []
    for cell in plan["cells"]:
        config = cell["config"]
        name = config["name"]
        sim_loads = {float(load["offered_request_rate"]): load for load in sim_by_name[name]["loads"]}
        loads = []
        for rate in cell["rates"]:
            round_summaries = []
            files = []
            for round_id in (1, 2):
                path = find_real_result(real_roots, name, round_id, float(rate))
                payload = json.loads(path.read_text())
                if payload.get("schema") != "qwen235b-t0-rate-anchor-v1":
                    raise ValueError(f"bad real result schema: {path}")
                if payload["summary"]["completed"] != 64 or payload["summary"]["failed"] != 0:
                    raise ValueError(f"incomplete real anchor: {path}")
                if float(payload["workload"]["offered_request_rate"]) != float(rate):
                    raise ValueError(f"offered-rate drift: {path}")
                round_summaries.append(payload["summary"]["slos"])
                files.append({"path": str(path), "sha256": sha256(path)})
            conservative = {
                slo: {
                    "feasible": all(summary[slo]["feasible"] for summary in round_summaries),
                    "round_pass_rates": [summary[slo]["pass_rate"] for summary in round_summaries],
                }
                for slo in SLOS
            }
            loads.append(
                {
                    "rate": float(rate),
                    "real_conservative": conservative,
                    "sim": sim_loads[float(rate)]["slos"],
                    "real_files": files,
                }
            )
        cells.append({"config": config, "loads": loads})

    comparisons = {}
    for slo in SLOS:
        records = []
        for cell in cells:
            real = capacity(cell["loads"], slo, "real_conservative")
            sim_values = [
                float(load["offered_request_rate"])
                for load in sim_by_name[cell["config"]["name"]]["loads"]
                if bool(load["slos"][slo]["feasible"])
            ]
            sim = max(sim_values) if sim_values else None
            tp = int(cell["config"]["tp"])
            boundary = real_boundary_status(cell["loads"], slo)
            records.append(
                {
                    "config": cell["config"],
                    "real_capacity": real,
                    "sim_capacity": sim,
                    "real_capacity_per_gpu": real / tp if real is not None else None,
                    "sim_capacity_per_gpu": sim / tp if sim is not None else None,
                    "real_boundary_status": boundary,
                    "expansion_required": boundary in {
                        "non_monotonic_requires_full_lattice",
                        "unbracketed_requires_expansion",
                    },
                }
            )
        rankable = [
            row
            for row in records
            if not row["expansion_required"]
            and row["real_capacity_per_gpu"] is not None
            and row["sim_capacity_per_gpu"] is not None
        ]
        tau = kendall_tau_b(
            [row["real_capacity_per_gpu"] for row in rankable],
            [row["sim_capacity_per_gpu"] for row in rankable],
        )
        pairwise = {
            "all_pairs": 0,
            "exact_sign_matches": 0,
            "real_non_tie_pairs": 0,
            "real_non_tie_direction_matches": 0,
        }
        for i, left in enumerate(rankable):
            for right in rankable[i + 1 :]:
                real_sign = sign(left["real_capacity_per_gpu"], right["real_capacity_per_gpu"])
                sim_sign = sign(left["sim_capacity_per_gpu"], right["sim_capacity_per_gpu"])
                pairwise["all_pairs"] += 1
                pairwise["exact_sign_matches"] += real_sign == sim_sign
                if real_sign:
                    pairwise["real_non_tie_pairs"] += 1
                    pairwise["real_non_tie_direction_matches"] += real_sign == sim_sign
        real_best = max((row["real_capacity_per_gpu"] for row in records if row["real_capacity_per_gpu"] is not None), default=None)
        sim_best = max((row["sim_capacity_per_gpu"] for row in records if row["sim_capacity_per_gpu"] is not None), default=None)
        sim_top = [row for row in records if sim_best is not None and row["sim_capacity_per_gpu"] == sim_best]
        real_top = [row for row in records if real_best is not None and row["real_capacity_per_gpu"] == real_best]
        optimistic_regret = worst_regret = None
        if real_best is not None and sim_top and all(row["real_capacity_per_gpu"] is not None for row in sim_top):
            regrets = [(real_best - row["real_capacity_per_gpu"]) / real_best for row in sim_top]
            optimistic_regret = min(regrets)
            worst_regret = max(regrets)
        confusion = {"anchors": 0, "match": 0, "false_feasible": 0, "false_infeasible": 0}
        for cell in cells:
            for load in cell["loads"]:
                real_feasible = bool(load["real_conservative"][slo]["feasible"])
                sim_feasible = bool(load["sim"][slo]["feasible"])
                confusion["anchors"] += 1
                confusion["match"] += real_feasible == sim_feasible
                confusion["false_feasible"] += sim_feasible and not real_feasible
                confusion["false_infeasible"] += real_feasible and not sim_feasible
        comparisons[slo] = {
            "records": records,
            "kendall_tau_b": tau,
            "pairwise": pairwise,
            "anchor_confusion": confusion,
            "real_top_set": [row["config"]["name"] for row in real_top],
            "sim_top_set": [row["config"]["name"] for row in sim_top],
            "top_set_intersection": sorted(
                {row["config"]["name"] for row in real_top}
                & {row["config"]["name"] for row in sim_top}
            ),
            "top_set_exact_match": {
                row["config"]["name"] for row in real_top
            } == {row["config"]["name"] for row in sim_top},
            "optimistic_regret": optimistic_regret,
            "worst_tie_break_regret": worst_regret,
        }

    run_costs = []
    for root in real_roots:
        config_names = [child.name for child in root.iterdir() if child.is_dir() and child.name in sim_by_name]
        if len(config_names) != 1:
            raise ValueError(f"expected one config directory in real root {root}, got {config_names}")
        config = sim_by_name[config_names[0]]["config"]
        remote_run = root.parents[2] / "remote_run"
        started = datetime.fromisoformat((remote_run / "started_at").read_text().strip())
        finished = datetime.fromisoformat((remote_run / "finished_at").read_text().strip())
        wall_seconds = (finished - started).total_seconds()
        run_costs.append(
            {
                "run_id": root.parents[2].name,
                "config": config_names[0],
                "gpu_count": int(config["tp"]),
                "wall_seconds": wall_seconds,
                "h20_gpu_hours": wall_seconds * int(config["tp"]) / 3600,
                "started_at": started.isoformat(),
                "finished_at": finished.isoformat(),
            }
        )
    fresh_server_anchors = 2 * sum(len(cell["loads"]) for cell in cells)
    real_execution_cost = {
        "accepted_fleet_jobs": len(run_costs),
        "fresh_server_anchors": fresh_server_anchors,
        "measured_requests": fresh_server_anchors * 64,
        "warmup_requests": 2
        * sum(
            min(32, max(4, math.ceil(float(load["rate"]) * 20)))
            for cell in cells
            for load in cell["loads"]
        ),
        "accepted_h20_gpu_hours": sum(run["h20_gpu_hours"] for run in run_costs),
        "campaign_wall_span_seconds": (
            max(datetime.fromisoformat(run["finished_at"]) for run in run_costs)
            - min(datetime.fromisoformat(run["started_at"]) for run in run_costs)
        ).total_seconds(),
        "runs": run_costs,
    }

    output = {
        "schema": "qwen235b-t0-simulator-real-comparison-v1",
        "frontier_freeze_sha256": sha256(args.frontier_freeze),
        "real_plan_sha256": sha256(args.real_plan),
        "real_execution_cost": real_execution_cost,
        "cells": cells,
        "comparisons": comparisons,
    }
    args.output_root.mkdir(parents=True, exist_ok=True)
    write_json(args.output_root / "comparison.json", output)
    with (args.output_root / "capacity.csv").open("w", newline="") as target:
        writer = csv.DictWriter(
            target,
            fieldnames=[
                "slo",
                "config",
                "tp",
                "mns",
                "mbt",
                "real_capacity_per_gpu",
                "sim_capacity_per_gpu",
                "real_boundary_status",
                "expansion_required",
            ],
        )
        writer.writeheader()
        for slo, comparison in comparisons.items():
            for row in comparison["records"]:
                writer.writerow(
                    {
                        "slo": slo,
                        "config": row["config"]["name"],
                        "tp": row["config"]["tp"],
                        "mns": row["config"]["mns"],
                        "mbt": row["config"]["mbt"],
                        "real_capacity_per_gpu": row["real_capacity_per_gpu"],
                        "sim_capacity_per_gpu": row["sim_capacity_per_gpu"],
                        "real_boundary_status": row["real_boundary_status"],
                        "expansion_required": row["expansion_required"],
                    }
                )
    print(args.output_root / "comparison.json")


if __name__ == "__main__":
    main()