"""Report harness: GQA end-to-end latency + op/engine breakdown.

Isolated under ``scripts/paper/`` for the 1H codesign report only — it is
NOT a registered bench and does not touch other people's benches. It
reuses the existing GQA headline panels (the real GQA kernels wired in
``milestone_gqa_headline``) but, unlike that milestone (which records only
op-counts), it also harvests per-panel end-to-end latency and per-engine
occupancy from ``result.engine.op_log``.

Latency definition (same window convention as ``milestone_1h_gemm``):
    end_to_end_ns = max(r.t_end) - min(r.t_start)  over all op_log records.

Output: docs/report/1H-codesign-paper/figures/gqa_latency.json

Run:
    python scripts/paper/paper_gqa_latency.py
"""
from __future__ import annotations

import json
from pathlib import Path

from kernbench.benches.milestone_gqa_headline import (
    _PANEL_DISPATCH,
    _PANELS,
    _make_bench_fn,
    _summarize_op_log,
)

_REPORT_DIR = Path(__file__).resolve().parents[2] / "docs" / "report" / "1H-codesign-paper"
_FIG_DIR = _REPORT_DIR / "figures"
_OUT_JSON = _FIG_DIR / "gqa_latency.json"

# PE engine component suffixes whose occupancy we break out.
_ENGINES = ("pe_gemm", "pe_math", "pe_dma", "pe_fetch_store", "pe_ipcq", "pe_cpu")


def _occupancy_ns(op_log, suffix: str) -> float:
    return sum(
        r.t_end - r.t_start
        for r in op_log
        if r.component_id.endswith("." + suffix)
    )


def _end_to_end_ns(op_log) -> float:
    if not op_log:
        return 0.0
    return max(r.t_end for r in op_log) - min(r.t_start for r in op_log)


def _run_panel(panel: str, topology: str) -> dict:
    from kernbench.runtime_api.bench_runner import run_bench
    from kernbench.runtime_api.types import resolve_device
    from kernbench.sim_engine.engine import GraphEngine
    from kernbench.topology.builder import resolve_topology

    topo = resolve_topology(topology)
    result = run_bench(
        topology=topo,
        bench_fn=_make_bench_fn(panel),
        device=resolve_device(None),
        engine_factory=lambda t, d: GraphEngine(
            getattr(t, "topology_obj", t), enable_data=True,
        ),
    )
    if not result.completion.ok:
        raise RuntimeError(f"panel {panel!r} failed: {result.completion}")

    op_log = result.engine.op_log
    kind, params = _PANEL_DISPATCH[panel]
    return {
        "panel": panel,
        "kind": kind,
        **params,
        "latency_ns": _end_to_end_ns(op_log),
        "op_log_summary": _summarize_op_log(op_log),
        "engine_occupancy_ns": {
            eng: _occupancy_ns(op_log, eng) for eng in _ENGINES
        },
    }


def main() -> None:
    topology = "topology.yaml"
    rows = [_run_panel(panel, topology) for panel in _PANELS]
    _FIG_DIR.mkdir(parents=True, exist_ok=True)
    out = {"version": 1, "panels": list(_PANELS), "rows": rows}
    _OUT_JSON.write_text(json.dumps(out, indent=2))
    print(f"wrote {_OUT_JSON}")
    for r in rows:
        s = r["op_log_summary"]
        print(
            f"  {r['panel']:24s} latency={r['latency_ns']:10.1f} ns  "
            f"gemm={s['gemm_count']:3d} ipcq={s['ipcq_copy_count']:3d} "
            f"dma_rd={s['dma_read_count']:3d} dma_wr={s['dma_write_count']:2d}"
        )


if __name__ == "__main__":
    main()