"""Generate a markdown summary for representative H3 block profiles."""

from __future__ import annotations

import argparse
import json
from pathlib import Path


MAIN_TIMINGS = (
    "norm1",
    "modulate_msa",
    "attn_qkv_proj",
    "attn_qkv_split_view",
    "attn_qk_rms_rope",
    "attn_q_bshd_contiguous",
    "attn_k_bshd_contiguous",
    "attn_v_bshd_contiguous",
    "attention_kernel",
    "attn_output_reshape",
    "attn_out_proj",
    "gate_msa",
    "norm2",
    "modulate_mlp",
    "mlp_fc1",
    "mlp_swiglu",
    "mlp_fc2",
    "gate_mlp",
    "block_total",
)

LINEARS = ("attn_qkv_proj", "attn_out_proj", "mlp_fc1", "mlp_fc2")
LINEAR_STAGES = (
    "activation_scale",
    "activation_quant_pack",
    "gemm",
    "pre_quant_scale",
    "packed_weight_wrapper",
    "bias_cast",
    "scale_to_device",
    "activation_quant_wrap",
    "flatten_contiguous",
    "slice_reshape",
)


def ms(value: float) -> float:
    return value * 1000.0


def row(timings: dict, name: str) -> tuple[float, float, float]:
    item = timings[name]
    return ms(item["mean_s"]), ms(item["p50_s"]), ms(item["p95_s"])


def add_table(lines: list[str], timings: dict, names: tuple[str, ...]) -> None:
    lines.extend(["| Stage | Mean ms | P50 ms | P95 ms |", "| --- | ---: | ---: | ---: |"])
    for name in names:
        if name not in timings:
            continue
        mean, p50, p95 = row(timings, name)
        lines.append(f"| `{name}` | {mean:.3f} | {p50:.3f} | {p95:.3f} |")


def linear_table(lines: list[str], timings: dict) -> None:
    lines.extend(["| Linear | Scale P50 ms | Pack P50 ms | GEMM P50 ms | Total P50 ms |", "| --- | ---: | ---: | ---: | ---: |"])
    for linear in LINEARS:
        prefix = f"linear.{linear}"
        scale = ms(timings[f"{prefix}.activation_scale"]["p50_s"])
        pack = ms(timings[f"{prefix}.activation_quant_pack"]["p50_s"])
        gemm = ms(timings[f"{prefix}.gemm"]["p50_s"])
        total = ms(timings[linear]["p50_s"])
        lines.append(f"| `{linear}` | {scale:.3f} | {pack:.3f} | {gemm:.3f} | {total:.3f} |")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    data = json.loads(args.input.read_text(encoding="utf-8"))
    timings = data["timings"]
    block_p50 = ms(timings["block_total"]["p50_s"])
    attention_layout_p50 = sum(ms(timings[name]["p50_s"]) for name in ("attn_q_bshd_contiguous", "attn_k_bshd_contiguous", "attn_v_bshd_contiguous", "attention_kernel", "attn_output_reshape"))
    mlp_linear_p50 = sum(ms(timings[name]["p50_s"]) for name in ("mlp_fc1", "mlp_fc2"))
    norm_mod_gate_p50 = sum(ms(timings[name]["p50_s"]) for name in ("norm1", "modulate_msa", "gate_msa", "norm2", "modulate_mlp", "gate_mlp"))
    qk_rope_p50 = ms(timings["attn_qk_rms_rope"]["p50_s"])

    lines = [
        "# H3 Block Profile Baseline",
        "",
        f"Input: block `{data['block_index']}`, hidden shape `{data['hidden_shape']}`, segments `{data['segments']}`.",
        f"Config: attention `{data['attention']}`, `{data['warmup']}` warmups, `{data['iterations']}` iterations.",
        "",
        "## Summary",
        "",
        f"- Block total p50: `{block_p50:.3f} ms`; p95: `{ms(timings['block_total']['p95_s']):.3f} ms`.",
        f"- Sol attention plus BSHD layout p50: `{attention_layout_p50:.3f} ms`.",
        f"- MLP linears p50: `{mlp_linear_p50:.3f} ms`; SwiGLU p50: `{ms(timings['mlp_swiglu']['p50_s']):.3f} ms`.",
        f"- Norm/modulate/gate p50: `{norm_mod_gate_p50:.3f} ms`.",
        f"- Q/K RMS+RoPE p50: `{qk_rope_p50:.3f} ms`.",
        "- The next likely targets are Q/K RMS+RoPE, Q/K/V contiguous extraction, and the modulation/gating helpers; Sol output reshape remains negligible.",
        "",
        "## Main Stages",
        "",
    ]
    add_table(lines, timings, MAIN_TIMINGS)
    lines.extend(["", "## NVFP4 Linear Breakdown", ""])
    linear_table(lines, timings)
    for linear in LINEARS:
        lines.extend(["", f"## `{linear}` Stages", ""])
        add_table(lines, timings, tuple(f"linear.{linear}.{stage}" for stage in LINEAR_STAGES))
    lines.extend(["", "## Source", "", f"- JSON: `{args.input}`"])
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text("\n".join(lines) + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()