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

Prefill counterpart to milestone-gqa-decode-long-ctx-4cases. Per
GQA_full_deck.pptx slides 11-17: 4 KV-cache sharding strategies on
the LLaMA-3.1-70B single-KV-head group (1 KV head, 8 Q heads,
d_head=128) on 8 cubes × 8 PEs.

Bench size T_q=S_kv=512 (equal — prompt length matches K/V length
for prefill). Cases 2 and 3 use Q-axis tiling (``TILE_Q``) so the
per-PE scratch is bounded by ``TILE_Q × TILE_S_KV`` regardless of
T_q — much larger sizes are possible at the cost of bench
wall-clock (Case 2's single-cube serial path dominates).

    Case 1  Cube-SP / PE-TP  → KV split S_kv across cubes (Ring),
                               Q/O split T_q across all 64 ranks
    Case 2  Cube-Repl / PE-TP → full KV per cube; only CUBE 0's P PEs
                               split T_q; CUBEs 1..7 are pure memory waste
    Case 3  Cube-Repl / PE-SP → full KV per cube; PEs SP on S_kv
                               (intra-cube AR); 8× redundant compute
    Case 4  Cube-SP / PE-SP  → KV split 64-way + Q split T_q across cubes;
                               Ring KV + intra-cube AR per cube  ★ optimal

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

Distinction from decode: decode T_q=1 makes PE-TP cases wasteful (1
active PE). Prefill T_q≫1 makes PE-TP useful (full parallelism). The
narrative differs — the figures emphasise Case 2/3 memory waste and
Case 3 compute redundancy, not PE-TP-at-B=1 wasted ranks.

Per-panel try/except in ``run()`` writes whichever cases succeed —
the bench tolerates per-panel scratch / SFR failures so sweep.json
always lands with at least the successful rows + a ``failures`` list.

Gated by ``GQA_PREFILL_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_prefill_long_ctx_cube_repl_pe_sp import (
    gqa_attention_prefill_long_ctx_cube_repl_pe_sp_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_prefill_long_ctx_cube_repl_pe_tp import (
    gqa_attention_prefill_long_ctx_cube_repl_pe_tp_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_prefill_long_ctx_cube_sp_pe_sp import (
    gqa_attention_prefill_long_ctx_cube_sp_pe_sp_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_prefill_long_ctx_cube_sp_pe_tp import (
    gqa_attention_prefill_long_ctx_cube_sp_pe_tp_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_ring
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_prefill.json"


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


_PANELS = (
    "single_kv_group_prefill_long_ctx_gqa_cube_sp_pe_sp",      # Case 4 ★ optimal
    "single_kv_group_prefill_long_ctx_gqa_cube_repl_pe_tp",    # Case 2
    "single_kv_group_prefill_long_ctx_gqa_cube_repl_pe_sp",    # Case 3
    "single_kv_group_prefill_long_ctx_gqa_cube_sp_pe_tp",      # Case 1
)

# 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
#   T_q = S_kv = 4096 (long-context prefill)
_PANEL_DISPATCH: dict[str, tuple[str, dict]] = {
    "single_kv_group_prefill_long_ctx_gqa_cube_sp_pe_sp": ("prefill_long_ctx_cube_sp_pe_sp", {
        # Case 4: KV split 64-way + Q split T_q across cubes; Ring KV
        # + intra-CUBE AR per cube. PE 0 of each cube writes its
        # T_q/C disjoint output rows.
        "C": 8, "P": 8,
        "T_q": 512, "S_kv": 512,
        "d_head": 128, "h_q": 8, "h_kv": 1,
    }),
    "single_kv_group_prefill_long_ctx_gqa_cube_repl_pe_tp": ("prefill_long_ctx_cube_repl_pe_tp", {
        # Case 2: K, V replicated everywhere (8× memory waste); within
        # CUBE 0 the P PEs split T_q disjointly. CUBEs 1..7 idle.
        # No inter-rank communication.
        "C": 8, "P": 8,
        "T_q": 512, "S_kv": 512,
        "d_head": 128, "h_q": 8, "h_kv": 1,
    }),
    "single_kv_group_prefill_long_ctx_gqa_cube_repl_pe_sp": ("prefill_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 does redundant T_q × S_kv compute; designated
        # writer = cube 0).
        "C": 8, "P": 8,
        "T_q": 512, "S_kv": 512,
        "d_head": 128, "h_q": 8, "h_kv": 1,
    }),
    "single_kv_group_prefill_long_ctx_gqa_cube_sp_pe_tp": ("prefill_long_ctx_cube_sp_pe_tp", {
        # Case 1: K, V split S_kv across cubes (S_local = S_kv/C per
        # cube); Q/O split T_q across all 64 ranks. Inter-CUBE snake
        # Ring KV (P parallel rings); no intra-CUBE comm.
        "C": 8, "P": 8,
        "T_q": 512, "S_kv": 512,
        "d_head": 128, "h_q": 8, "h_kv": 1,
    }),
}


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


def _ring_sfr(ctx):
    """Install snake Ring SFR (also provides intra_* lanes for AR cases).

    Snake submesh of shape (2, 4) over the default 4×4 cube mesh
    yields a Hamiltonian cycle through all 8 cubes
    (0→1→2→3→7→6→5→4→0), with intra_* lanes available for the
    Case 3/4 intra-CUBE reduce.
    """
    configure_sfr_intercube_ring(
        ctx.engine, ctx.spec, _ccl_cfg(),
        submesh_shape=(2, 4),
    )


def _run_prefill_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; CUBE 0 PE-TP on T_q.

    DPPolicy models the cluster-wide 8× memory waste — every rank
    holds full K, V in its HBM region. Within CUBE 0, P PEs split
    T_q (so P of 64 ranks actually compute). CUBEs 1..7 idle.
    """
    _ring_sfr(ctx)
    dp_repl = DPPolicy(cube="replicate", pe="replicate",
                       num_cubes=C, num_pes=P)
    dp_qo = 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_qo, 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_qo, name=f"{panel}_o")
    ctx.launch(
        panel, gqa_attention_prefill_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_prefill_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 intra-cube.

    DPPolicy models the cluster-wide 8× memory waste — every cube
    holds full K, V; intra-cube splits the S_kv axis row_wise across
    8 PEs. Every rank holds full Q (cube=replicate, pe=replicate).
    The kernel does an intra-CUBE 8-way reduce on (m, ℓ, O); only
    cube 0's PE 0 writes the output.
    """
    _ring_sfr(ctx)
    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_prefill_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_prefill_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 S_kv across cubes (Ring); Q/O split T_q 64-way.

    DPPolicy: K, V are cube=row_wise on S_kv (S_local = S_kv/C per
    cube), pe=replicate within cube. Q, O are cube=row_wise +
    pe=row_wise on T_q — every rank owns T_q/(C·P) disjoint Q rows.
    The kernel runs P parallel Rings through the snake submesh
    (each PE has its own IPCQ lane); no intra-CUBE reduce.
    """
    _ring_sfr(ctx)
    dp_qo = DPPolicy(cube="row_wise", pe="row_wise",
                     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_qo, 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_qo, name=f"{panel}_o")
    ctx.launch(
        panel, gqa_attention_prefill_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_prefill_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; Q split T_q across cubes.

    DPPolicy: K, V are cube=row_wise + pe=row_wise on S_kv — each
    rank owns S_local_pe = S_kv/(C·P). Q, O are cube=row_wise on
    T_q (T_q_cube = T_q/C per cube), pe=replicate within cube. Ring
    KV through the snake submesh (P parallel lanes per cube); after
    the ring, an intra-CUBE 8-way AR combines PE lanes so cube's
    PE 0 has the full attention for its T_q_cube rows. PE 0 of each
    cube writes its disjoint output rows (no inter-cube reduce).
    """
    _ring_sfr(ctx)
    dp_qo = DPPolicy(cube="row_wise", 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_qo, 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_qo, name=f"{panel}_o")
    ctx.launch(
        panel, gqa_attention_prefill_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 _make_bench_fn(panel: str):
    kind, params = _PANEL_DISPATCH[panel]

    def _bench_fn(ctx):
        if kind == "prefill_long_ctx_cube_sp_pe_sp":
            _run_prefill_panel_long_ctx_cube_sp_pe_sp(ctx, panel=panel, **params)
        elif kind == "prefill_long_ctx_cube_repl_pe_tp":
            _run_prefill_panel_long_ctx_cube_repl_pe_tp(ctx, panel=panel, **params)
        elif kind == "prefill_long_ctx_cube_repl_pe_sp":
            _run_prefill_panel_long_ctx_cube_repl_pe_sp(ctx, panel=panel, **params)
        elif kind == "prefill_long_ctx_cube_sp_pe_tp":
            _run_prefill_panel_long_ctx_cube_sp_pe_tp(ctx, panel=panel, **params)
        else:
            raise RuntimeError(
                f"milestone-gqa-prefill-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-prefill-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 prefill case panels; write sweep.json. Returns row count.

    Per-panel try/except: even with Q-axis tiling, very large T_q or
    unusual config combinations can hit the per-PE scratch budget. We
    keep going so sweep.json carries whichever panels succeed; failures
    land in a ``failures`` list for follow-up.
    """
    _OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    rows: list[dict] = []
    failures: list[dict] = []
    for panel in _PANELS:
        try:
            rows.append(_run_panel(panel, topology))
        except Exception as e:
            print(f"  panel {panel!r} FAILED: {e}")
            failures.append({"panel": panel, "error": str(e)})
    sweep = {
        "version": 1,
        "panels": list(_PANELS),
        "rows": rows,
        "failures": failures,
    }
    _SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
    print(f"  gqa-prefill-long-ctx-4cases: {len(rows)} rows -> {_SWEEP_JSON}")
    return len(rows)
