#!/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_stable = True
    for tp in (2, 4):
        for batch in (2, 4, 6, 8):
            repeats = []
            for repeat in (1, 2):
                path = args.results_dir / f"tp{tp}-b{batch}-r{repeat}.json"
                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]
            mean = statistics.fmean(values)
            cv_pct = statistics.pstdev(values) / mean * 100.0
            stable = cv_pct <= 10.0
            all_stable = all_stable and stable
            cells.append(
                {
                    "tp": tp,
                    "batch": batch,
                    "repeats": repeats,
                    "median_execute_ms": statistics.median(values),
                    "repeat_cv_pct": cv_pct,
                    "stable": stable,
                    "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 engine processes",
            "stability_gate": "population CV <= 10%; add repeat 3 otherwise",
        },
        "all_cells_stable": all_stable,
        "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_stable:
        raise SystemExit("one or more cells failed the repeat stability gate")


if __name__ == "__main__":
    main()