"""GQA composite kernel across models — per-KV-group multi-model comparison.

The topology per model uses C = h_q (cubes per KV group = query heads
per KV group in the model). P = 8 PEs / cube is fixed to the physical
layout. So Gemma-2 27B (G=2) runs on a 2-cube slice; Command R+ (G=12)
runs on 12 cubes; LLaMA-3-70B (G=8) matches the Case-6 baseline
exactly.

Companion to ``gqa_decode_long_ctx_models.py``. That bench sweeps a
fixed set of topology sizes (N ∈ {16, 32, 64}) across six GQA models
and showed that latency is model-agnostic when only G varies.

This bench takes a different slice: it fixes the topology to match
each model's G ratio — one cube per Q head in the KV group:

    C = G  (cubes per KV group = h_q of the KV group)
    P = 8  (fixed, physical PEs per cube)
    N = C · P = 8·G

Six models, one topology per model, one context length:

    Gemma 2 27B      G=2   → C=2   → N=16
    LLaMA-3 8B       G=4   → C=4   → N=32
    Qwen 2.5 7B      G=7   → C=7   → N=56  (2 KV groups per SIP, 2 cubes idle)
    LLaMA-3 70B      G=8   → C=8   → N=64
    Qwen 2.5 72B     G=8   → C=8   → N=64
    Command R+       G=12  → C=12  → N=96  (1 KV group per SIP, 4 cubes idle)

Physical SIP is 4×4 = 16 cubes (topology.yaml sip.cube_mesh: {w:4, h:4}),
so all six configurations fit in one SIP.

Total: 6 engine runs at S_kv = 131 072 (~40 min wall-clock).
"""
from __future__ import annotations

import json
from pathlib import Path

from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite import (  # noqa: E501
    gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel,
)
from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import _ccl_cfg
from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
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_decode_models.json"


# ── Model registry (per-KV-group shapes) ─────────────────────────────


# Kernel operates on ONE KV group at a time. h_q / h_kv here are the
# per-group counts (not the model's totals); G = h_q / h_kv.
# Full-model context (full_h_q, full_h_kv, hidden, layers) is metadata
# for annotation only — it does not enter the kernel launch.
_MODELS = {
    "llama3-8b": dict(
        family="Meta",     params_b=8,
        h_q=4,  h_kv=1, d_head=128,
        full_h_q=32,  full_h_kv=8,  hidden=4096,  layers=32,
    ),
    "llama3-70b": dict(
        family="Meta",     params_b=70,
        h_q=8,  h_kv=1, d_head=128,
        full_h_q=64,  full_h_kv=8,  hidden=8192,  layers=80,
    ),
    "qwen2.5-7b": dict(
        family="Alibaba",  params_b=7,
        h_q=7,  h_kv=1, d_head=128,
        full_h_q=28,  full_h_kv=4,  hidden=3584,  layers=28,
    ),
    "qwen2.5-72b": dict(
        family="Alibaba",  params_b=72,
        h_q=8,  h_kv=1, d_head=128,
        full_h_q=64,  full_h_kv=8,  hidden=8192,  layers=80,
    ),
    "gemma2-27b": dict(
        family="Google",   params_b=27,
        h_q=2,  h_kv=1, d_head=128,
        full_h_q=32,  full_h_kv=16, hidden=4608,  layers=46,
    ),
    "command-r-plus": dict(
        family="Cohere",   params_b=104,
        h_q=12, h_kv=1, d_head=128,
        full_h_q=96,  full_h_kv=8,  hidden=12288, layers=64,
    ),
}
_MODEL_ORDER = tuple(_MODELS.keys())

_S_KV = 131_072   # 128 K context
_T_Q = 1
_P = 8            # fixed physical PEs per cube


def _c_for(model_key: str) -> int:
    m = _MODELS[model_key]
    return m["h_q"] // m["h_kv"]


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 _op_breakdown_ns(op_log) -> dict[str, float]:
    """Sum occupancy per op_kind across all components (not
    critical-path — ops overlap on independent components). First-order
    view of where time is spent inside the engine."""
    totals = {"matmul": 0.0, "comm": 0.0, "other": 0.0}
    for r in op_log:
        dur = r.t_end - r.t_start
        if r.op_kind in ("gemm", "math"):
            totals["matmul"] += dur
        elif r.op_kind == "memory":
            totals["comm"] += dur
        else:
            totals["other"] += dur
    return totals


def _run_panel_fn(model_key: str, C: int, P: int, S_kv: int):
    m = _MODELS[model_key]
    panel = f"decode_models_{model_key}_c{C}p{P}_s{S_kv}"

    def _bench_fn(ctx):
        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
        dp_full = DPPolicy(cube="replicate", pe="replicate",
                           num_cubes=C, num_pes=P)
        dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
                         num_cubes=C, num_pes=P)
        q = ctx.zeros((_T_Q, m["h_q"] * m["d_head"]),
                      dtype="f16", dp=dp_full, name=f"{panel}_q")
        k = ctx.zeros((S_kv, m["h_kv"] * m["d_head"]),
                      dtype="f16", dp=dp_kv, name=f"{panel}_k")
        v = ctx.zeros((S_kv, m["h_kv"] * m["d_head"]),
                      dtype="f16", dp=dp_kv, name=f"{panel}_v")
        o = ctx.empty((_T_Q, m["h_q"] * m["d_head"]),
                      dtype="f16", dp=dp_full, name=f"{panel}_o")
        ctx.launch(
            panel,
            gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel,
            q, k, v, o,
            _T_Q, S_kv, m["h_q"], m["h_kv"], m["d_head"], C, P,
            _auto_dim_remap=False,
        )

    return _bench_fn


def _emit_dispatch(model_key: str, C: int, P: int,
                   S_kv: int) -> tuple[int, float]:
    from kernbench.common.pe_commands import PeCpuOverheadCmd
    from kernbench.common.pe_cost_model import DEFAULT_PE_COST_MODEL
    from kernbench.triton_emu.tl_context import TLContext, run_kernel

    m = _MODELS[model_key]
    cube_id = min(6, C - 1)
    tl = TLContext(
        pe_id=0, num_programs=P, cost_model=DEFAULT_PE_COST_MODEL,
        cube_id=cube_id, num_cubes=C, scratch_base=1 << 61,
        scratch_size=1 << 20,
    )
    run_kernel(
        gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel, tl,
        0x1000, 0x2000, 0x3000, 0x4000,
        _T_Q, S_kv, m["h_q"], m["h_kv"], m["d_head"], C, P,
    )
    cmds = [c for c in tl.commands if isinstance(c, PeCpuOverheadCmd)]
    return len(cmds), sum(c.cycles for c in cmds)


def _engine_run(model_key: str, C: int, P: int,
                S_kv: int, topology: str):
    """Run the engine sim; return (latency_ns, op_kind_breakdown)."""
    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(model_key, C, P, S_kv),
        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-decode-models {model_key} C={C} P={P} S={S_kv} "
            f"failed: {result.completion}"
        )
    op_log = result.engine.op_log
    latency = _end_to_end_ns(op_log)
    breakdown = _op_breakdown_ns(op_log)
    return latency, breakdown


def run_sweep(topology: str = "topology.yaml") -> int:
    """One row per model at its per-KV-group topology (C = h_q). Writes
    sweep_decode_models.json."""
    _OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    rows = []
    for model_key in _MODEL_ORDER:
        m = _MODELS[model_key]
        C = _c_for(model_key)
        N = C * _P
        n_cmds, cycles = _emit_dispatch(model_key, C, _P, _S_KV)
        latency, breakdown = _engine_run(model_key, C, _P, _S_KV, topology)
        rows.append({
            "model": model_key,
            "family": m["family"],
            "params_b": m["params_b"],
            "h_q": m["h_q"],
            "h_kv": m["h_kv"],
            "G": C,
            "d_head": m["d_head"],
            "full_h_q": m["full_h_q"],
            "full_h_kv": m["full_h_kv"],
            "hidden": m["hidden"],
            "layers": m["layers"],
            "C": C, "P": _P, "N": N,
            "T_q": _T_Q,
            "S_kv": _S_KV,
            "S_local": _S_KV // N,
            "variant": "composite",
            "pe_cpu_cmd_count": n_cmds,
            "pe_cpu_dispatch_cycles": cycles,
            "latency_ns": latency,
            # Op-kind occupancy (summed across all components — not
            # critical-path; first-order view of where engine time goes)
            "matmul_ns": breakdown["matmul"],
            "comm_ns": breakdown["comm"],
            "other_ns": breakdown["other"],
        })
        print(
            f"  {model_key:<18}  G={C:>2}  C={C:>2} P={_P}  N={N:>3}  "
            f"latency={latency/1e3:>7.1f} µs  cmds={n_cmds:>4}  "
            f"matmul={breakdown['matmul']/1e3:>7.1f}  "
            f"comm={breakdown['comm']/1e3:>7.1f}"
        )
    sweep = {
        "version": 1,
        "variant": "composite",
        "S_kv": _S_KV,
        "P": _P,
        "note": "C = h_q per model (cubes per KV group = query heads per KV group)",
        "models": list(_MODEL_ORDER),
        "rows": rows,
    }
    _SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
    print(f"  gqa-decode-models: {len(rows)} rows -> {_SWEEP_JSON}")
    return len(rows)


if __name__ == "__main__":
    run_sweep()
