#!/usr/bin/env python3
"""Aggregate fresh-process decode profiles and enforce the repeat gate."""

from __future__ import annotations

import argparse
import json
import statistics
from pathlib import Path


COMPONENT_NAMES = (
    "attention",
    "linear_norm_rope",
    "router",
    "moe",
    "collective",
    "output_head",
    "other",
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--results-dir", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    cells = []
    all_resolved = True
    for tp in (2, 4):
        for batch in (2, 4, 6, 8):
            repeats = []
            for repeat in (1, 2, 3):
                path = args.results_dir / f"tp{tp}-b{batch}-r{repeat}.json"
                if repeat == 3 and not path.exists():
                    continue
                payload = json.loads(path.read_text())
                repeats.append(
                    {
                        "repeat": repeat,
                        "source": str(path),
                        "execute_mean_ms": payload["rank_summary"][
                            "slowest_rank_execute_mean_ms"
                        ],
                        "component_rank_mean_ms": payload["rank_summary"][
                            "component_rank_mean_ms"
                        ],
                    }
                )
            values = [row["execute_mean_ms"] for row in repeats]
            initial_values = values[:2]
            initial_mean = statistics.fmean(initial_values)
            initial_cv_pct = (
                statistics.pstdev(initial_values) / initial_mean * 100.0
            )
            needs_repeat_3 = initial_cv_pct > 10.0
            repeat_gate_resolved = not needs_repeat_3 or len(repeats) == 3
            all_resolved = all_resolved and repeat_gate_resolved
            cells.append(
                {
                    "tp": tp,
                    "batch": batch,
                    "repeats": repeats,
                    "median_execute_ms": statistics.median(values),
                    "initial_repeat_cv_pct": initial_cv_pct,
                    "all_repeat_cv_pct": (
                        statistics.pstdev(values)
                        / statistics.fmean(values)
                        * 100.0
                    ),
                    "needs_repeat_3": needs_repeat_3,
                    "repeat_gate_resolved": repeat_gate_resolved,
                    "median_component_ms": {
                        name: statistics.median(
                            row["component_rank_mean_ms"][name] for row in repeats
                        )
                        for name in COMPONENT_NAMES
                    },
                }
            )
    payload = {
        "schema": "frontier-decode-batch-grid.v1",
        "contract": {
            "workload": "Qwen3-30B-A3B BF16, 2048->128, graph-on, MNS=16",
            "timing": "slowest-rank execute mean over 16 pure-decode steps",
            "repeat_aggregation": "median of two fresh processes; three if unstable",
            "stability_gate": "r1/r2 population CV <=10%; else require r3 and take median",
        },
        "all_repeat_gates_resolved": all_resolved,
        "cells": cells,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    if not all_resolved:
        raise SystemExit("one or more cells require repeat 3")


if __name__ == "__main__":
    main()