#!/usr/bin/env python3
"""Summarize graph-on vLLM Torch traces by decode component."""

from __future__ import annotations

import argparse
import gzip
import json
import math
import statistics
from collections import Counter, defaultdict
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("--trace-root", type=Path, required=True)
    parser.add_argument("--label", required=True)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def classify(name: str, occurrences: int, steps: int) -> str:
    lower = name.lower()
    if any(token in lower for token in ("nccl", "allreduce", "all_reduce")):
        return "collective"
    if "topkgating" in lower or "fused_topk" in lower:
        return "router"
    if any(
        token in lower
        for token in (
            "fused_moe",
            "moefcgemm",
            "tensorrt_llm::kernels::cutlass_kernels",
            "groupproblemshape",
            "memcpy32_post",
        )
    ):
        return "moe"
    if any(
        token in lower
        for token in (
            "flashattn",
            "flashattnfwd",
            "reshape_and_cache",
            "prepare_varlen_num_blocks",
        )
    ):
        return "attention"
    if "nvjet" in lower and occurrences <= steps * 2:
        return "output_head"
    if any(
        token in lower
        for token in (
            "nvjet",
            "cublaslt",
            "rms_norm",
            "rsqrt",
            "triton_red_fused_2",
            "triton_poi_fused_3",
            "triton_red_fused_0",
            "triton_poi_fused_1",
        )
    ):
        return "linear_norm_rope"
    return "other"


def stats(values: list[float]) -> dict[str, float | int]:
    ordered = sorted(values)
    return {
        "n": len(values),
        "mean_ms": statistics.fmean(values),
        "population_std_ms": statistics.pstdev(values),
        "p50_ms": statistics.median(ordered),
        "p95_ms": ordered[math.ceil(0.95 * len(ordered)) - 1],
    }


def load_events(path: Path) -> list[dict]:
    opener = gzip.open if path.suffix == ".gz" else open
    with opener(path, "rt") as source:
        return json.load(source)["traceEvents"]


def analyze_rank(path: Path) -> dict:
    events = load_events(path)
    kernels = [event for event in events if event.get("cat") == "kernel"]
    all_windows = sorted(
        (
            event
            for event in events
            if event.get("cat") == "gpu_user_annotation"
            and str(event.get("name", "")).startswith("execute_")
        ),
        key=lambda event: float(event["ts"]),
    )
    if not all_windows:
        raise ValueError(f"{path}: no GPU execute annotations")
    window_names = Counter(str(window["name"]) for window in all_windows)
    selected_name = window_names.most_common(1)[0][0]
    windows = [
        window for window in all_windows if str(window["name"]) == selected_name
    ]

    selected: list[dict] = []
    step_kernels: list[list[dict]] = []
    for window in windows:
        start = float(window["ts"])
        end = start + float(window["dur"])
        current = [
            kernel for kernel in kernels if start <= float(kernel["ts"]) < end
        ]
        selected.extend(current)
        step_kernels.append(current)

    occurrences = Counter(str(kernel["name"]) for kernel in selected)
    component_steps: dict[str, list[float]] = defaultdict(list)
    busy_steps: list[float] = []
    wall_steps = [float(window["dur"]) / 1000.0 for window in windows]
    for current in step_kernels:
        per_component: dict[str, float] = defaultdict(float)
        for kernel in current:
            name = str(kernel["name"])
            per_component[classify(name, occurrences[name], len(windows))] += (
                float(kernel["dur"]) / 1000.0
            )
        for name in COMPONENT_NAMES:
            component_steps[name].append(per_component[name])
        busy_steps.append(sum(per_component.values()))

    return {
        "trace": str(path),
        "selected_execute_annotation": selected_name,
        "execute_annotation_histogram": dict(sorted(window_names.items())),
        "steps": len(windows),
        "execute_wall": stats(wall_steps),
        "gpu_kernel_busy": stats(busy_steps),
        "non_kernel_gap": stats(
            [wall - busy for wall, busy in zip(wall_steps, busy_steps)]
        ),
        "components": {
            name: stats(values) for name, values in component_steps.items()
        },
    }


def main() -> None:
    args = parse_args()
    traces = sorted(args.trace_root.rglob("*.pt.trace.json*"))
    if not traces:
        raise SystemExit(f"no traces under {args.trace_root}")
    ranks = [analyze_rank(path) for path in traces]
    payload = {
        "schema": "frontier-decode-batch-trace.v1",
        "label": args.label,
        "contract": {
            "timing": "CUDA graph-on GPU execute annotations",
            "component_time": "sum of CUDA kernel durations inside execute range",
            "tp_aggregation": "per-rank; slowest rank mean approximates critical path",
        },
        "ranks": ranks,
        "rank_summary": {
            "ranks": len(ranks),
            "slowest_rank_execute_mean_ms": max(
                rank["execute_wall"]["mean_ms"] for rank in ranks
            ),
            "slowest_rank_kernel_busy_mean_ms": max(
                rank["gpu_kernel_busy"]["mean_ms"] for rank in ranks
            ),
            "component_rank_mean_ms": {
                name: statistics.fmean(
                    rank["components"][name]["mean_ms"] for rank in ranks
                )
                for name in COMPONENT_NAMES
            },
        },
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")


if __name__ == "__main__":
    main()