#!/usr/bin/env python3
"""Compare the measured-prefill-MoE TP2 replays with structured baseline."""

from __future__ import annotations

import importlib.util
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent
REPO = ROOT.parents[1]
ATTN = REPO / "runs/frontier-attn-structured-v0"
S3_REAL = REPO / "runs/frontier-s3-real-v0"
CELLS = {
    "tp2_rho0p0025": "frontier-s3-real-full-r0p0025-tp2-t*",
    "tp2_rho0p005": "frontier-s3-real-full-r0p005-tp2-t*",
}


def load_analysis_module():
    path = ATTN / "analyze_trace_verdict.py"
    spec = importlib.util.spec_from_file_location("attention_verdict", path)
    module = importlib.util.module_from_spec(spec)
    sys.path.insert(0, str(ATTN))
    spec.loader.exec_module(module)
    return module


def main() -> None:
    analysis = load_analysis_module()
    cells = {}
    for label, pattern in CELLS.items():
        real_trials = analysis.load_real_trials(pattern)
        structured = analysis.load_sim(ATTN / "replay" / label)
        moe_corrected = analysis.load_sim(ROOT / "replay" / label)
        structured_biases = [
            analysis.distribution_bias(trial, structured)
            for trial in real_trials
        ]
        corrected_biases = [
            analysis.distribution_bias(trial, moe_corrected)
            for trial in real_trials
        ]
        cells[label] = {
            "structured_attention": {
                "trialwise_distribution_bias": structured_biases,
                "trialwise_distribution_bias_summary": (
                    analysis.aggregate_trial_bias(structured_biases)
                ),
                "legacy_pooled_distribution_bias": (
                    analysis.legacy_pooled_bias(real_trials, structured)
                ),
                "waiting_p99_ms": analysis.waiting_p99(structured),
            },
            "structured_attention_plus_prefill_moe": {
                "trialwise_distribution_bias": corrected_biases,
                "trialwise_distribution_bias_summary": (
                    analysis.aggregate_trial_bias(corrected_biases)
                ),
                "legacy_pooled_distribution_bias": (
                    analysis.legacy_pooled_bias(real_trials, moe_corrected)
                ),
                "paired_relative_error": [
                    analysis.paired_relative_error(trial, moe_corrected)
                    for trial in real_trials
                ],
                "waiting_p99_ms": analysis.waiting_p99(moe_corrected),
            },
        }

    low = cells["tp2_rho0p0025"]
    before = low["structured_attention"]["legacy_pooled_distribution_bias"]
    after = low["structured_attention_plus_prefill_moe"][
        "legacy_pooled_distribution_bias"
    ]
    checked = [
        (metric, quantile)
        for metric in ("ttft", "e2e")
        for quantile in ("mean", "p50", "p99")
    ]
    gates = {
        "subcritical_waiting_below_1s": (
            low["structured_attention_plus_prefill_moe"]["waiting_p99_ms"]
            < 1000
        ),
        "subcritical_ttft_e2e_abs_bias_not_worse": all(
            abs(after[metric][quantile])
            <= abs(before[metric][quantile]) + 0.01
            for metric, quantile in checked
        ),
        "subcritical_mean_ttft_abs_bias_improves_3pp": (
            abs(before["ttft"]["mean"]) - abs(after["ttft"]["mean"])
            >= 0.03
        ),
    }
    payload = {
        "schema": "frontier-tp2-prefill-serving-replay-verdict.v1",
        "cells": cells,
        "gates": gates,
        "decision": (
            "keep_tp2_prefill_moe_correction"
            if all(gates.values())
            else "reject_global_scale_and_fit_shape_conditioned_curve"
        ),
    }
    output = ROOT / "results/replay-verdict.json"
    output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    print(json.dumps({"gates": gates, "decision": payload["decision"]}, indent=2))


if __name__ == "__main__":
    main()