#!/usr/bin/env python3
"""Trial-aware verdict for the seven structured-attention trace replays."""

from __future__ import annotations

import csv
import json
import math
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parent
REPO = ROOT.parents[1]
S3_REAL = REPO / "runs/frontier-s3-real-v0"
V5 = REPO / "runs/frontier-prefill-kvgrowth-fix-v0"

CELLS = {
    "tp1_rho0p00125": {
        "real": "frontier-tp1-real-r0p00125-t*",
        "old": V5 / "sim-replay-tp1/v5/tp1_rho0p00125",
    },
    "tp1_rho0p0025": {
        "real": "frontier-tp1-real-r0p0025-t*",
        "old": V5 / "sim-replay-tp1/v5/tp1_rho0p0025",
    },
    "tp2_rho0p0025": {
        "real": "frontier-s3-real-full-r0p0025-tp2-t*",
        "old": V5 / "sim-replay/tp2_rho0p0025",
    },
    "tp2_rho0p005": {
        "real": "frontier-s3-real-full-r0p005-tp2-t*",
        "old": V5 / "sim-replay/tp2_rho0p005",
    },
    "tp4_rho0p0025": {
        "real": "frontier-s3-real-full-r0p0025-tp4-t*",
        "old": V5 / "sim-replay/tp4_rho0p0025",
    },
    "tp4_rho0p005": {
        "real": "frontier-s3-real-full-r0p005-tp4-t*",
        "old": V5 / "sim-replay/tp4_rho0p005",
    },
    "tp4_rho0p01": {
        "real": "frontier-s3-real-full-r0p01-tp4-t*",
        "old": V5 / "sim-replay/tp4_rho0p01",
    },
}
METRICS = {
    "ttft": ("ttft_ms", "ttft"),
    "tpot": ("tpot_ms", "tpot"),
    "e2e": ("e2e_ms", "request_e2e_time"),
}
QUANTILES = {"mean": None, "p50": 0.5, "p90": 0.9, "p99": 0.99}


def percentile(values: list[float], quantile: float) -> float:
    ordered = sorted(values)
    position = (len(ordered) - 1) * quantile
    lower, upper = math.floor(position), math.ceil(position)
    if lower == upper:
        return ordered[lower]
    return (
        ordered[lower] * (upper - position)
        + ordered[upper] * (position - lower)
    )


def summarize(values: list[float]) -> dict[str, float]:
    return {
        name: (
            sum(values) / len(values)
            if quantile is None
            else percentile(values, quantile)
        )
        for name, quantile in QUANTILES.items()
    }


def load_real_trials(pattern: str) -> list[list[dict[str, Any]]]:
    trials = []
    for run_root in sorted((S3_REAL / "fleet-artifacts").glob(pattern)):
        results = list(
            run_root.glob(
                "artifacts/outputs/full-real/*/*/trial-*/results/result.json"
            )
        )
        if len(results) != 1:
            raise ValueError(f"expected one result in {run_root}, got {results}")
        trials.append(json.loads(results[0].read_text())["requests"])
    if len(trials) != 2:
        raise ValueError(f"expected two real trials for {pattern}, got {len(trials)}")
    return trials


def load_sim(root: Path) -> list[dict[str, str]]:
    matches = list((root / "metrics").rglob("request_metrics.csv"))
    if len(matches) != 1:
        raise ValueError(f"expected one request_metrics.csv below {root}: {matches}")
    rows = list(csv.DictReader(matches[0].open()))
    rows.sort(key=lambda row: int(float(row["Request Id"])))
    return rows


def distribution_bias(
    real_rows: list[dict[str, Any]],
    sim_rows: list[dict[str, str]],
) -> dict[str, dict[str, float]]:
    output: dict[str, dict[str, float]] = {}
    for metric, (real_key, sim_key) in METRICS.items():
        pairs = [
            (float(real[real_key]), float(sim[sim_key]))
            for real, sim in zip(real_rows, sim_rows)
            if real.get("success")
        ]
        real_summary = summarize([pair[0] for pair in pairs])
        sim_summary = summarize([pair[1] for pair in pairs])
        output[metric] = {
            name: (sim_summary[name] - real_summary[name]) / real_summary[name]
            for name in QUANTILES
        }
    return output


def paired_relative_error(
    real_rows: list[dict[str, Any]],
    sim_rows: list[dict[str, str]],
) -> dict[str, dict[str, float]]:
    output: dict[str, dict[str, float]] = {}
    for metric, (real_key, sim_key) in METRICS.items():
        errors = [
            (float(sim[sim_key]) - float(real[real_key])) / float(real[real_key])
            for real, sim in zip(real_rows, sim_rows)
            if real.get("success") and float(real[real_key]) != 0
        ]
        output[metric] = summarize(errors)
    return output


def aggregate_trial_bias(
    trial_biases: list[dict[str, dict[str, float]]],
) -> dict[str, dict[str, dict[str, float]]]:
    return {
        metric: {
            quantile: {
                "mean": sum(values) / len(values),
                "min": min(values),
                "max": max(values),
            }
            for quantile in QUANTILES
            for values in [
                [trial[metric][quantile] for trial in trial_biases]
            ]
        }
        for metric in METRICS
    }


def legacy_pooled_bias(
    real_trials: list[list[dict[str, Any]]],
    sim_rows: list[dict[str, str]],
) -> dict[str, dict[str, float]]:
    output: dict[str, dict[str, float]] = {}
    for metric, (real_key, sim_key) in METRICS.items():
        real_values = [
            float(row[real_key])
            for trial in real_trials
            for row in trial[: len(sim_rows)]
            if row.get("success")
        ]
        sim_values = [float(row[sim_key]) for row in sim_rows]
        real_summary = summarize(real_values)
        sim_summary = summarize(sim_values)
        output[metric] = {
            name: (sim_summary[name] - real_summary[name]) / real_summary[name]
            for name in QUANTILES
        }
    return output


def waiting_p99(sim_rows: list[dict[str, str]]) -> float:
    return percentile(
        [float(row["request_waiting_time_total"]) for row in sim_rows], 0.99
    )


def main() -> None:
    results: dict[str, Any] = {}
    flat_rows: list[dict[str, Any]] = []
    for label, paths in CELLS.items():
        real_trials = load_real_trials(paths["real"])
        old_sim = load_sim(paths["old"])
        new_sim = load_sim(ROOT / "replay" / label)
        old_trial_bias = [
            distribution_bias(trial, old_sim) for trial in real_trials
        ]
        new_trial_bias = [
            distribution_bias(trial, new_sim) for trial in real_trials
        ]
        old_legacy = legacy_pooled_bias(real_trials, old_sim)
        new_legacy = legacy_pooled_bias(real_trials, new_sim)
        wait_p99 = waiting_p99(new_sim)
        results[label] = {
            "old": {
                "trialwise_distribution_bias": old_trial_bias,
                "trialwise_distribution_bias_summary": aggregate_trial_bias(
                    old_trial_bias
                ),
                "legacy_pooled_distribution_bias": old_legacy,
            },
            "new": {
                "trialwise_distribution_bias": new_trial_bias,
                "trialwise_distribution_bias_summary": aggregate_trial_bias(
                    new_trial_bias
                ),
                "paired_relative_error": [
                    paired_relative_error(trial, new_sim)
                    for trial in real_trials
                ],
                "legacy_pooled_distribution_bias": new_legacy,
                "waiting_p99_ms": wait_p99,
                "validity": (
                    "PASS_SUBCRITICAL"
                    if wait_p99 < 1000
                    else "GATE_FAIL_DIAGNOSTIC"
                ),
            },
        }
        for metric in METRICS:
            for quantile in QUANTILES:
                flat_rows.append(
                    {
                        "cell": label,
                        "metric": metric,
                        "quantile": quantile,
                        "old_bias": old_legacy[metric][quantile],
                        "new_bias": new_legacy[metric][quantile],
                        "abs_bias_delta_pp": 100
                        * (
                            abs(new_legacy[metric][quantile])
                            - abs(old_legacy[metric][quantile])
                        ),
                        "validity": results[label]["new"]["validity"],
                    }
                )

    tp1_checks = []
    for cell in ("tp1_rho0p00125", "tp1_rho0p0025"):
        for quantile in ("mean", "p99"):
            old = results[cell]["old"]["legacy_pooled_distribution_bias"]["ttft"][
                quantile
            ]
            new = results[cell]["new"]["legacy_pooled_distribution_bias"]["ttft"][
                quantile
            ]
            tp1_checks.append(abs(old) - abs(new) >= 0.05)

    regressions = [
        row
        for row in flat_rows
        if row["cell"].startswith(("tp2", "tp4"))
        and row["metric"] in ("ttft", "e2e")
        and row["abs_bias_delta_pp"] > 5
    ]
    checks = {
        "tp1_ttft_mean_p99_improve_ge_5pp": all(tp1_checks),
        "tp2_tp4_ttft_e2e_no_abs_regression_gt_5pp": not regressions,
        "regressions": regressions,
    }
    checks["trace_gate"] = (
        checks["tp1_ttft_mean_p99_improve_ge_5pp"]
        and checks["tp2_tp4_ttft_e2e_no_abs_regression_gt_5pp"]
    )

    payload = {
        "schema": "frontier-attn-structured-trial-aware-verdict-v1",
        "metric_note": (
            "Primary values are per-real-trial distribution biases with request "
            "alignment by index. legacy_pooled reproduces the old milestone "
            "quantile convention only for direct comparison."
        ),
        "cells": results,
        "checks": checks,
    }
    output = ROOT / "results"
    output.mkdir(parents=True, exist_ok=True)
    (output / "trace-verdict.json").write_text(json.dumps(payload, indent=2))
    with (output / "trace-verdict.csv").open("w", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=list(flat_rows[0]))
        writer.writeheader()
        writer.writerows(flat_rows)

    print(json.dumps(checks, indent=2))
    for label, result in results.items():
        old = result["old"]["legacy_pooled_distribution_bias"]["ttft"]
        new = result["new"]["legacy_pooled_distribution_bias"]["ttft"]
        print(
            f"{label}: TTFT mean {old['mean']:+.1%}->{new['mean']:+.1%}, "
            f"p99 {old['p99']:+.1%}->{new['p99']:+.1%}, "
            f"waiting_p99={result['new']['waiting_p99_ms']:.0f}ms "
            f"{result['new']['validity']}"
        )


if __name__ == "__main__":
    main()