"""Summarize an Nsight Systems CUDA trace into stable H3 component categories."""

from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path


def category(name: str) -> str:
    if any(token in name for token in (
        "qk_int_sv_f8_attn_kernel",
        "MeanScaleKernel",
        "TransposePadPermuteKernel",
        "QuantInt8Kernel",
    )):
        return "sage2"
    if "cutlass3x_sm120_bstensorop" in name:
        return "nvfp4_gemms"
    if any(token in name for token in (
        "partial_absmax_",
        "quantize_nvfp4_",
        "quantize_nvfp4_kernel",
        "final_scale_",
        "FillFunctor<unsigned char>",
    )):
        return "nvfp4_packing"
    if any(token in name for token in (
        "rope_kernel",
        "vectorized_layer_norm_kernel",
        "MeanOps<c10::BFloat16",
    )):
        return "norm_and_rope"
    if "_gate_add_kernel" in name:
        return "remaining_gate_add"
    return "other"


def merge(intervals: list[tuple[int, int]]) -> list[tuple[int, int]]:
    output: list[list[int]] = []
    for start, end in sorted(intervals):
        if not output or start > output[-1][1]:
            output.append([start, end])
        else:
            output[-1][1] = max(output[-1][1], end)
    return [(start, end) for start, end in output]


def total(intervals: list[tuple[int, int]]) -> int:
    return sum(end - start for start, end in merge(intervals))


def intersection_total(
    left: list[tuple[int, int]], right: list[tuple[int, int]],
) -> int:
    left = merge(left)
    right = merge(right)
    i = j = result = 0
    while i < len(left) and j < len(right):
        start = max(left[i][0], right[j][0])
        end = min(left[i][1], right[j][1])
        result += max(0, end - start)
        if left[i][1] <= right[j][1]:
            i += 1
        else:
            j += 1
    return result


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--gpu-trace", type=Path, required=True)
    parser.add_argument("--kernel-exec-trace", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    with args.gpu_trace.open(newline="", encoding="utf-8") as handle:
        gpu_rows = list(csv.DictReader(handle))
    kernels = [row for row in gpu_rows if row["GrdX"]]
    gpu_intervals = [
        (int(row["Start (ns)"]), int(row["Start (ns)"]) + int(row["Duration (ns)"]))
        for row in kernels
    ]
    all_gpu_intervals = [
        (int(row["Start (ns)"]), int(row["Start (ns)"]) + int(row["Duration (ns)"]))
        for row in gpu_rows
    ]
    first_gpu = min(start for start, _ in all_gpu_intervals)
    last_gpu = max(end for _, end in all_gpu_intervals)
    kernel_time = sum(int(row["Duration (ns)"]) for row in kernels)
    component_ns = {
        name: 0 for name in (
            "sage2", "nvfp4_gemms", "nvfp4_packing", "norm_and_rope",
            "remaining_gate_add", "other",
        )
    }
    component_launches = component_ns.copy()
    for row in kernels:
        name = category(row["Name"])
        component_ns[name] += int(row["Duration (ns)"])
        component_launches[name] += 1

    ordered = sorted(gpu_intervals)
    gaps = [
        max(0, ordered[index][0] - ordered[index - 1][1])
        for index in range(1, len(ordered))
    ]
    positive_gaps = [gap for gap in gaps if gap]

    with args.kernel_exec_trace.open(newline="", encoding="utf-8") as handle:
        launch_rows = list(csv.DictReader(handle))
    api_intervals = [
        (int(row["API Start (ns)"]), int(row["API Start (ns)"]) + int(row["API Dur (ns)"]))
        for row in launch_rows
    ]
    launch_api_time = total(api_intervals)
    launch_api_gpu_overlap = intersection_total(api_intervals, gpu_intervals)
    gpu_span = last_gpu - first_gpu

    report = {
        "source_gpu_trace": str(args.gpu_trace),
        "source_kernel_exec_trace": str(args.kernel_exec_trace),
        "gpu_span_ns": gpu_span,
        "gpu_operation_count": len(gpu_rows),
        "kernel_count": len(kernels),
        "kernel_time_ns": kernel_time,
        "kernel_busy_percent_of_span": kernel_time / gpu_span * 100.0,
        "launch_gaps": {
            "positive_gap_count": len(positive_gaps),
            "total_ns": sum(positive_gaps),
            "average_ns": sum(positive_gaps) / len(positive_gaps) if positive_gaps else 0,
            "maximum_ns": max(positive_gaps, default=0),
        },
        "cpu_gpu_overlap": {
            "scope": "CUDA kernel-launch API intervals intersected with GPU kernel intervals",
            "launch_api_union_ns": launch_api_time,
            "launch_api_gpu_overlap_ns": launch_api_gpu_overlap,
            "launch_api_overlap_percent": (
                launch_api_gpu_overlap / launch_api_time * 100.0 if launch_api_time else 0
            ),
        },
        "components": {
            name: {
                "milliseconds": value / 1.0e6,
                "percent_of_kernel_time": value / kernel_time * 100.0,
                "launches": component_launches[name],
            }
            for name, value in component_ns.items()
        },
        "classification_policy": {
            "sage2": "mainloop plus MeanScale/TransposePadPermute/QuantInt8 preparation",
            "nvfp4_gemms": "SM120 block-scaled CUTLASS GEMMs",
            "nvfp4_packing": "absmax, final-scale, zero-fill and NVFP4 quantization kernels",
            "norm_and_rope": "layer norm, mean reduction and fused RMS/RoPE kernels",
            "remaining_gate_add": "fused residual gate/add kernels",
            "other": "all unmatched kernels",
        },
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()