"""milestone-gqa-decode-long-ctx-4cases: long-context decode 4-cases study.

Per GQA_full_deck.pptx slides 11-17: 4 KV-cache sharding strategies on
the LLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs) at long
context (S_kv = 128K, T_q = 1):

    Case 1  Cube-SP / PE-TP  → KV split by S_kv across cubes; PEs TP on batch
    Case 2  Cube-Repl / PE-TP → full KV per cube; PEs TP on batch
    Case 3  Cube-Repl / PE-SP → full KV per cube; PEs SP on S_kv (intra-cube AR)
    Case 4  Cube-SP / PE-SP  → KV split 64-way; 2-phase AR on (m,ℓ,O)  ★ optimal

Each case is a separate panel. The bench drives all panels in one
invocation and writes per-panel op_log_summary to sweep.json so the
comparative analysis (latency, GEMM/MAC util, comm volume) can be
generated from a single sweep.

Deviation from slide 13: slide prescribes AllReduce (every rank has
the answer); the kernel does reduce-to-root (only the lrab center
cube has it) per ADR-0060 §4. Treated as the kernbench Case-4 baseline.

Gated by ``GQA_DECODE_LONG_CTX_4CASES_RUN=1``.
"""
from __future__ import annotations

import json
import os
from pathlib import Path

from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_repl_pe_sp import (
    gqa_attention_decode_long_ctx_cube_repl_pe_sp_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_repl_pe_tp import (
    gqa_attention_decode_long_ctx_cube_repl_pe_tp_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import (
    gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp import (
    gqa_attention_decode_long_ctx_cube_sp_pe_tp_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead import (
    gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp import (
    gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel,
)
from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import (
    _ccl_cfg,
    _summarize_op_log,
)
from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
from kernbench.policy.placement.dp import DPPolicy

# File is at benches/gqa_helpers/long_ctx/ — go up 2 parents to reach
# benches/, then into 1H_milestone_output/gqa/long_ctx/.
_OUTPUT_DIR = (
    Path(__file__).resolve().parents[2]
    / "1H_milestone_output" / "gqa" / "long_ctx"
)
_SWEEP_JSON = _OUTPUT_DIR / "sweep_decode.json"


# ── Panel registry ───────────────────────────────────────────────────


_PANELS = (
    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp",         # Case 6 ★ optimal (S_kv-SP both axes; was "Case 4" in slide-17 numbering)
    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp",       # Case 1 (no sharding; 40 GB / PE @ 1M)
    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp",       # Case 3 (intra-cube AR on (m,ℓ,O); 5 GB / PE)
    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp",         # Case 2 (inter-cube AR on (m,ℓ,O); 5 GB / PE)
    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp_dhead",   # Case 4 (Cube-SP × PE-TP-on-d_head — d_head-TP partial-score AR intra-cube)
    "single_kv_group_decode_long_ctx_gqa_cube_tp_dhead_pe_sp",   # Case 5 (Cube-TP-on-d_head × PE-SP — d_head-TP partial-score AR inter-cube)
)

# Each entry: (kind, panel-specific params).
# LLaMA-3.1-70B single-KV-head group target:
#   1 KV head, h_q = 8 (G = 8 group), d_head = 128
#   8 cubes (head-parallel group), 8 PEs/cube
#   S_kv = 128K (long-context decode), T_q = 1 (one new token per pass)
_PANEL_DISPATCH: dict[str, tuple[str, dict]] = {
    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp": ("decode_long_ctx_cube_sp_pe_sp", {
        # Case 4: KV split 64-way (Cube-SP × PE-SP), 2-level reduce.
        # The sub_w=4/sub_h=2 lrab center-root geometry (root cube 6) is
        # baked into the Case-4 wrapper kernel.
        "C": 8, "P": 8,
        "T_q": 1, "S_kv": 8_192,
        "d_head": 128, "h_q": 8, "h_kv": 1,
    }),
    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp": ("decode_long_ctx_cube_repl_pe_tp", {
        # Case 2: K, V replicated everywhere (8× memory waste); PEs TP
        # on batch. For B=1 only one rank works (slide-11 PE-TP waste).
        # No inter-rank communication.
        "C": 8, "P": 8,
        "T_q": 1, "S_kv": 8_192,
        "d_head": 128, "h_q": 8, "h_kv": 1,
    }),
    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp": ("decode_long_ctx_cube_repl_pe_sp", {
        # Case 3: K, V replicated per cube (8× memory); PEs SP on S_kv
        # within each cube. Intra-CUBE 8-way reduce; no inter-CUBE comm
        # (every cube ends with full answer; designated writer = cube 0).
        "C": 8, "P": 8,
        "T_q": 1, "S_kv": 8_192,
        "d_head": 128, "h_q": 8, "h_kv": 1,
    }),
    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp": ("decode_long_ctx_cube_sp_pe_tp", {
        # Case 1: K, V split across cubes (S_local = S_kv/C per cube);
        # PEs TP on batch — at B=1 only PE 0 of each cube works (PE-TP
        # waste). Inter-CUBE lrab AR (root = cube 6); no intra-CUBE comm.
        "C": 8, "P": 8,
        "T_q": 1, "S_kv": 8_192,
        "d_head": 128, "h_q": 8, "h_kv": 1,
    }),
    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp_dhead": ("decode_long_ctx_cube_sp_pe_tp_dhead", {
        # Case 4 (new numbering): Cube-SP on S_kv × PE-TP on d_head.
        # Partial-score AR happens intra-cube (NoC); (m,ℓ,O) merge happens
        # inter-cube (UCIe). Memory-equal to Case 6; comm dominated by
        # partial-score AR which scales with S_kv.
        "C": 8, "P": 8,
        "T_q": 1, "S_kv": 8_192,
        "d_head": 128, "h_q": 8, "h_kv": 1,
    }),
    "single_kv_group_decode_long_ctx_gqa_cube_tp_dhead_pe_sp": ("decode_long_ctx_cube_tp_dhead_pe_sp", {
        # Case 5 (new numbering): Cube-TP on d_head × PE-SP on S_kv.
        # Mirror of Case 4: partial-score AR happens INTER-cube (UCIe);
        # (m,ℓ,O) merge happens INTRA-cube (NoC). Memory-equal to Case 6.
        "C": 8, "P": 8,
        "T_q": 1, "S_kv": 8_192,
        "d_head": 128, "h_q": 8, "h_kv": 1,
    }),
}


# ── Per-panel runner ─────────────────────────────────────────────────


def _run_decode_panel_long_ctx_cube_repl_pe_tp(
    ctx, *, panel: str, C: int, P: int,
    T_q: int, S_kv: int,
    d_head: int, h_q: int, h_kv: int,
) -> None:
    """Case 2 runner: K, V replicated everywhere; B=1 single-rank work.

    DPPolicy models the cluster-wide memory waste — every rank holds
    full K, V in its HBM region. Only PE 0 of CUBE 0 computes (the
    kernel early-returns on every other rank), so only one rank reads
    from its HBM copy and writes the output.
    """
    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
    dp_repl = DPPolicy(cube="replicate", pe="replicate",
                       num_cubes=C, num_pes=P)
    q = ctx.zeros((T_q, h_q * d_head),
                  dtype="f16", dp=dp_repl, name=f"{panel}_q")
    k = ctx.zeros((S_kv, h_kv * d_head),
                  dtype="f16", dp=dp_repl, name=f"{panel}_k")
    v = ctx.zeros((S_kv, h_kv * d_head),
                  dtype="f16", dp=dp_repl, name=f"{panel}_v")
    o = ctx.empty((T_q, h_q * d_head),
                  dtype="f16", dp=dp_repl, name=f"{panel}_o")
    ctx.launch(
        panel, gqa_attention_decode_long_ctx_cube_repl_pe_tp_kernel,
        q, k, v, o,
        T_q, S_kv, h_q, h_kv, d_head, C, P,
        _auto_dim_remap=False,
    )


def _run_decode_panel_long_ctx_cube_repl_pe_sp(
    ctx, *, panel: str, C: int, P: int,
    T_q: int, S_kv: int,
    d_head: int, h_q: int, h_kv: int,
) -> None:
    """Case 3 runner: K, V replicated per cube; PEs SP on S_kv within cube.

    DPPolicy models the cluster-wide 8× memory waste — every cube
    holds full K, V in its HBM region, then splits the S_kv axis
    row_wise across its 8 PEs. The kernel does an intra-CUBE 8-way
    reduce on (m, ℓ, O); only cube 0's PE 0 writes the output.
    """
    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="replicate", pe="row_wise",
                     num_cubes=C, num_pes=P)
    q = ctx.zeros((T_q, h_q * d_head),
                  dtype="f16", dp=dp_full, name=f"{panel}_q")
    k = ctx.zeros((S_kv, h_kv * d_head),
                  dtype="f16", dp=dp_kv, name=f"{panel}_k")
    v = ctx.zeros((S_kv, h_kv * d_head),
                  dtype="f16", dp=dp_kv, name=f"{panel}_v")
    o = ctx.empty((T_q, h_q * d_head),
                  dtype="f16", dp=dp_full, name=f"{panel}_o")
    ctx.launch(
        panel, gqa_attention_decode_long_ctx_cube_repl_pe_sp_kernel,
        q, k, v, o,
        T_q, S_kv, h_q, h_kv, d_head, C, P,
        _auto_dim_remap=False,
    )


def _run_decode_panel_long_ctx_cube_sp_pe_tp(
    ctx, *, panel: str, C: int, P: int,
    T_q: int, S_kv: int,
    d_head: int, h_q: int, h_kv: int,
) -> None:
    """Case 1 runner: K, V split across cubes; PE-TP single-rank work at B=1.

    DPPolicy: K, V are cube=row_wise (S_local = S_kv/C per cube),
    pe=replicate (full S_local within each cube). At B=1, the kernel
    runs only on PE 0 of every cube; PEs 1-7 early-return. PE 0 per
    cube does local attention on S_local, then participates in the
    inter-CUBE lrab AR; PE 0 of cube 6 (lrab center) writes O.
    """
    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="replicate",
                     num_cubes=C, num_pes=P)
    q = ctx.zeros((T_q, h_q * d_head),
                  dtype="f16", dp=dp_full, name=f"{panel}_q")
    k = ctx.zeros((S_kv, h_kv * d_head),
                  dtype="f16", dp=dp_kv, name=f"{panel}_k")
    v = ctx.zeros((S_kv, h_kv * d_head),
                  dtype="f16", dp=dp_kv, name=f"{panel}_v")
    o = ctx.empty((T_q, h_q * d_head),
                  dtype="f16", dp=dp_full, name=f"{panel}_o")
    ctx.launch(
        panel, gqa_attention_decode_long_ctx_cube_sp_pe_tp_kernel,
        q, k, v, o,
        T_q, S_kv, h_q, h_kv, d_head, C, P,
        _auto_dim_remap=False,
    )


def _run_decode_panel_long_ctx_cube_sp_pe_sp(
    ctx, *, panel: str, C: int, P: int,
    T_q: int, S_kv: int,
    d_head: int, h_q: int, h_kv: int,
) -> None:
    """Case 4 runner: K, V split 64-way (Cube-SP × PE-SP).

    DPPolicy: K, V are cube=row_wise, pe=row_wise — each rank owns
    S_local = S_kv/(C·P). The wrapper kernel bakes in the lrab
    sub_w=4 / sub_h=2 geometry (root cube 6, ADR-0060 §4.2).
    """
    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, h_q * d_head),
                  dtype="f16", dp=dp_full, name=f"{panel}_q")
    k = ctx.zeros((S_kv, h_kv * d_head),
                  dtype="f16", dp=dp_kv, name=f"{panel}_k")
    v = ctx.zeros((S_kv, h_kv * d_head),
                  dtype="f16", dp=dp_kv, name=f"{panel}_v")
    o = ctx.empty((T_q, h_q * d_head),
                  dtype="f16", dp=dp_full, name=f"{panel}_o")
    ctx.launch(
        panel, gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel,
        q, k, v, o,
        T_q, S_kv, h_q, h_kv, d_head, C, P,
        _auto_dim_remap=False,
    )


def _run_decode_panel_long_ctx_cube_sp_pe_tp_dhead(
    ctx, *, panel: str, C: int, P: int,
    T_q: int, S_kv: int,
    d_head: int, h_q: int, h_kv: int,
) -> None:
    """Case 4 (new) runner: Cube-SP × PE-TP-on-d_head.

    Q is cube-replicated, PE-sharded column-wise (d_head/P per PE).
    K, V are cube=row_wise on S_kv, pe=column_wise on d_head — each PE
    owns (S_kv/C, d_head/P). The kernel does an intra-cube AR-SUM of
    partial scores then an inter-cube online-softmax merge of
    (m, ℓ, O_shard).
    """
    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
    dp_q = DPPolicy(cube="replicate", pe="column_wise",
                    num_cubes=C, num_pes=P)
    dp_kv = DPPolicy(cube="row_wise", pe="column_wise",
                     num_cubes=C, num_pes=P)
    q = ctx.zeros((T_q, h_q * d_head),
                  dtype="f16", dp=dp_q, name=f"{panel}_q")
    k = ctx.zeros((S_kv, h_kv * d_head),
                  dtype="f16", dp=dp_kv, name=f"{panel}_k")
    v = ctx.zeros((S_kv, h_kv * d_head),
                  dtype="f16", dp=dp_kv, name=f"{panel}_v")
    o = ctx.empty((T_q, h_q * d_head),
                  dtype="f16", dp=dp_q, name=f"{panel}_o")
    ctx.launch(
        panel, gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel,
        q, k, v, o,
        T_q, S_kv, h_q, h_kv, d_head, C, P,
        _auto_dim_remap=False,
    )


def _run_decode_panel_long_ctx_cube_tp_dhead_pe_sp(
    ctx, *, panel: str, C: int, P: int,
    T_q: int, S_kv: int,
    d_head: int, h_q: int, h_kv: int,
) -> None:
    """Case 5 (new) runner: Cube-TP-on-d_head × PE-SP-on-S_kv.

    Q is cube-sharded column_wise (d_head/C per cube), pe-replicated
    within cube. K, V are cube=column_wise on d_head, pe=row_wise on
    S_kv — each PE owns (S_kv/P, d_head/C). The kernel does an
    inter-cube AR-SUM of partial scores then an intra-cube online-
    softmax merge of (m, ℓ, O) reduce-to-PE0.
    """
    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
    dp_q = DPPolicy(cube="column_wise", pe="replicate",
                    num_cubes=C, num_pes=P)
    dp_kv = DPPolicy(cube="column_wise", pe="row_wise",
                     num_cubes=C, num_pes=P)
    q = ctx.zeros((T_q, h_q * d_head),
                  dtype="f16", dp=dp_q, name=f"{panel}_q")
    k = ctx.zeros((S_kv, h_kv * d_head),
                  dtype="f16", dp=dp_kv, name=f"{panel}_k")
    v = ctx.zeros((S_kv, h_kv * d_head),
                  dtype="f16", dp=dp_kv, name=f"{panel}_v")
    o = ctx.empty((T_q, h_q * d_head),
                  dtype="f16", dp=dp_q, name=f"{panel}_o")
    ctx.launch(
        panel, gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel,
        q, k, v, o,
        T_q, S_kv, h_q, h_kv, d_head, C, P,
        _auto_dim_remap=False,
    )


def _make_bench_fn(panel: str):
    kind, params = _PANEL_DISPATCH[panel]

    def _bench_fn(ctx):
        if kind == "decode_long_ctx_cube_sp_pe_sp":
            _run_decode_panel_long_ctx_cube_sp_pe_sp(ctx, panel=panel, **params)
        elif kind == "decode_long_ctx_cube_repl_pe_tp":
            _run_decode_panel_long_ctx_cube_repl_pe_tp(ctx, panel=panel, **params)
        elif kind == "decode_long_ctx_cube_repl_pe_sp":
            _run_decode_panel_long_ctx_cube_repl_pe_sp(ctx, panel=panel, **params)
        elif kind == "decode_long_ctx_cube_sp_pe_tp":
            _run_decode_panel_long_ctx_cube_sp_pe_tp(ctx, panel=panel, **params)
        elif kind == "decode_long_ctx_cube_sp_pe_tp_dhead":
            _run_decode_panel_long_ctx_cube_sp_pe_tp_dhead(ctx, panel=panel, **params)
        elif kind == "decode_long_ctx_cube_tp_dhead_pe_sp":
            _run_decode_panel_long_ctx_cube_tp_dhead_pe_sp(ctx, panel=panel, **params)
        else:
            raise RuntimeError(
                f"milestone-gqa-decode-long-ctx-4cases panel {panel!r} has "
                f"unsupported kind={kind!r}."
            )
    return _bench_fn


# ── Panel metrics (sweep.json carries these for the comparative plot) ─


_ENGINE_SUFFIXES = (
    "pe_gemm", "pe_math", "pe_dma", "pe_fetch_store", "pe_ipcq", "pe_cpu",
)


def _end_to_end_ns(op_log) -> float:
    """End-to-end window: ``max(t_end) - min(t_start)`` over all records."""
    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_occupancy_ns(op_log) -> dict[str, float]:
    """Per-engine summed occupancy (component_id suffix match)."""
    return {
        eng: sum(
            r.t_end - r.t_start
            for r in op_log
            if r.component_id.endswith("." + eng)
        )
        for eng in _ENGINE_SUFFIXES
    }


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"milestone-gqa-decode-long-ctx-4cases panel {panel!r} failed: "
            f"{result.completion}"
        )
    kind, params = _PANEL_DISPATCH[panel]
    op_log = result.engine.op_log
    return {
        "panel": panel,
        "kind": kind,
        **params,
        "op_log_summary": _summarize_op_log(op_log),
        "latency_ns": _end_to_end_ns(op_log),
        "engine_occupancy_ns": _engine_occupancy_ns(op_log),
    }


# ── Sweep entry (called by the umbrella ``milestone_1h_gqa``) ────────


def run_sweep(topology: str = "topology.yaml") -> int:
    """Drive all decode case panels; write sweep.json. Returns row count."""
    _OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    rows = [_run_panel(panel, topology) for panel in _PANELS]
    sweep = {
        "version": 1,
        "panels": list(_PANELS),
        "rows": rows,
    }
    _SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
    print(f"  gqa-decode-long-ctx-4cases: {len(rows)} rows -> {_SWEEP_JSON}")
    return len(rows)
