#!/usr/bin/env python3
"""Combine the frozen simulator entry audit with the serving trace smoke."""

from __future__ import annotations

import json
from pathlib import Path

ROOT = Path(__file__).resolve().parent
ENTRY = ROOT / "results/entry-audit.json"
SMOKE = ROOT / "results/serving-smoke.json"


def main() -> None:
    entry = json.loads(ENTRY.read_text())
    smoke = json.loads(SMOKE.read_text())
    simulator = entry["cells"]["tp2"]
    critical = smoke["critical_rank"]

    sim_total = float(simulator["total_ms"])
    real_total = float(simulator["real_total_ms"])
    sim_moe = float(simulator["moe_grouped_gemm_ms"])
    serving_total = float(critical["execute_wall_ms"])
    serving_moe = float(critical["components_ms"]["moe"])
    total_residual = real_total - sim_total
    moe_residual = serving_moe - sim_moe
    counterfactual_total = sim_total + moe_residual
    payload = {
        "schema": "frontier-tp2-prefill-serving-smoke-verdict.v1",
        "contract": {
            "simulator_sample": "mean of first nine q8192/ctx0 single-request chunks",
            "serving_sample": "longest execute window on critical TP rank",
            "moe_mapping": (
                "standalone moe_grouped_gemm and serving MoE both include "
                "expert prepare/finalize plus expert GEMMs"
            ),
            "real_anchor_ms": real_total,
        },
        "measurements_ms": {
            "simulator_total": sim_total,
            "serving_execute_wall": serving_total,
            "real_anchor": real_total,
            "simulator_moe_grouped_gemm": sim_moe,
            "serving_moe": serving_moe,
            "simulator_non_moe": sim_total - sim_moe,
            "serving_non_moe": serving_total - serving_moe,
        },
        "counterfactual": {
            "total_residual_ms": total_residual,
            "moe_residual_ms": moe_residual,
            "moe_residual_fraction": moe_residual / total_residual,
            "moe_scale": serving_moe / sim_moe,
            "moe_only_counterfactual_total_ms": counterfactual_total,
            "moe_only_counterfactual_bias": (
                counterfactual_total - real_total
            )
            / real_total,
        },
        "gates": {
            "serving_reproduces_anchor_within_5pct": (
                abs(serving_total - real_total) / real_total <= 0.05
            ),
            "moe_explains_at_least_70pct": moe_residual / total_residual >= 0.70,
            "moe_only_counterfactual_within_5pct": (
                abs(counterfactual_total - real_total) / real_total <= 0.05
            ),
        },
    }
    payload["decision"] = (
        "inject_tp2_prefill_moe_and_replay"
        if all(payload["gates"].values())
        else "stop_moe_injection_and_profile_whole_layer"
    )
    output = ROOT / "results/serving-smoke-verdict.json"
    output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    print(json.dumps(payload, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()