"""milestone-1h-gqa: compute-bound prefill composite-command study.

Three command-form variants of a single-rank compute-bound prefill
attention kernel (``_gqa_prefill_compute_bound``), swept over context
length S_kv = T_q. Unlike the memory-bound decode study, prefill has a
large M = G·T_q, so the score / context GEMMs are compute-bound — the
regime where the composite command keeps the MAC array fed (DMA↔compute
pipelining) and the hand-tiled primitive starves it on the serial
load→dot path.

Records per (variant, context) the end-to-end latency, the GEMM-engine
busy time, and the MAC occupancy (gemm_busy / e2e) so the comparative
plot can show composite winning on both latency and utilization.

Runs in data mode (engine latency). Gated via the umbrella
``GQA_1H_SWEEPS=prefill_cb``.
"""
from __future__ import annotations

import json
from pathlib import Path

from kernbench.benches.gqa_helpers.long_ctx._gqa_prefill_compute_bound import (
    gqa_prefill_composite_ext_kernel,
    gqa_prefill_composite_kernel,
    gqa_prefill_primitive_kernel,
)
from kernbench.policy.placement.dp import DPPolicy

_OUTPUT_DIR = (
    Path(__file__).resolve().parents[2]
    / "1H_milestone_output" / "gqa" / "long_ctx"
)
_SWEEP_JSON = _OUTPUT_DIR / "sweep_prefill_compute_bound.json"

_VARIANT_KERNELS = {
    "primitive": gqa_prefill_primitive_kernel,
    "composite": gqa_prefill_composite_kernel,
    "composite_extended": gqa_prefill_composite_ext_kernel,
}
_VARIANTS = ("primitive", "composite", "composite_extended")

# Context length S_kv = T_q (prefill processes T_q tokens against S_kv=T_q
# keys). M = G·T_q = 8·T_q is tile-filling/compute-bound at every point.
_CTX_POINTS = (256, 512, 1024)

_H_Q, _H_KV, _D_HEAD = 8, 1, 128
_PEAK_TFLOPS = 8.0  # per-PE f16 GEMM peak (topology.yaml pe_gemm.peak_tflops_f16)


def _run_panel_fn(variant: str, ctx_len: int):
    kernel = _VARIANT_KERNELS[variant]
    panel = f"prefill_cb_{variant}_c{ctx_len}"

    def _bench_fn(ctx):
        dp = DPPolicy(cube="replicate", pe="replicate",
                      num_cubes=1, num_pes=1)
        q = ctx.zeros((ctx_len, _H_Q * _D_HEAD),
                      dtype="f16", dp=dp, name=f"{panel}_q")
        k = ctx.zeros((ctx_len, _H_KV * _D_HEAD),
                      dtype="f16", dp=dp, name=f"{panel}_k")
        v = ctx.zeros((ctx_len, _H_KV * _D_HEAD),
                      dtype="f16", dp=dp, name=f"{panel}_v")
        o = ctx.empty((ctx_len, _H_Q * _D_HEAD),
                      dtype="f16", dp=dp, name=f"{panel}_o")
        ctx.launch(panel, kernel, q, k, v, o,
                   ctx_len, ctx_len, _H_Q, _H_KV, _D_HEAD, 1, 1,
                   _auto_dim_remap=False)

    return _bench_fn


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 _engine_busy_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 _run_panel(variant: str, ctx_len: int, 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=_run_panel_fn(variant, ctx_len),
        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"gqa-prefill-cb {variant}@{ctx_len} failed: {result.completion}"
        )
    op_log = result.engine.op_log
    e2e = _end_to_end_ns(op_log)
    gemm = _engine_busy_ns(op_log, "pe_gemm")
    dma = _engine_busy_ns(op_log, "pe_dma")
    # Useful attention flops (Q·Kᵀ + P·V), single rank.
    G = _H_Q // _H_KV
    M = G * ctx_len
    useful_flops = 4.0 * M * _D_HEAD * ctx_len
    # MAC utilization = achieved / peak. ``achieved_tflops`` uses wall-clock
    # (useful_flops / e2e) so it is bounded by peak even when the composite
    # path overlaps many tile GEMMs (gemm_busy is a sum over overlapping ops
    # and is kept only for diagnostics).
    return {
        "variant": variant,
        "ctx_len": ctx_len,
        "M": M,
        "latency_ns": e2e,
        "gemm_busy_ns": gemm,
        "dma_busy_ns": dma,
        "achieved_tflops": (useful_flops / e2e / 1e3) if e2e > 0 else 0.0,
        "mac_util": (useful_flops / e2e / 1e3 / _PEAK_TFLOPS) if e2e > 0 else 0.0,
    }


def run_sweep(topology: str = "topology.yaml") -> int:
    """Drive all (variant, context) prefill panels; write sweep.json."""
    _OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    rows = [
        _run_panel(variant, ctx_len, topology)
        for ctx_len in _CTX_POINTS
        for variant in _VARIANTS
    ]
    sweep = {
        "version": 1,
        "variants": list(_VARIANTS),
        "ctx_points": list(_CTX_POINTS),
        "rows": rows,
    }
    _SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
    print(f"  gqa-prefill-cb: {len(rows)} rows -> {_SWEEP_JSON}")
    return len(rows)
