"""GQA decode kernel — Case 4 (Cube-SP × PE-SP). ★ slide-11 optimal.

Per GQA_full_deck.pptx slide 11:
  - K, V split 64-way (cube=row_wise, pe=row_wise); each rank owns
    ``S_local = S_kv / (C·P)``.
  - Local attention via per-rank S_kv-axis tile sweep (ADR-0063 §A.2)
    keeps scratch bounded by ``TILE_S_KV``.
  - Two-level reduce on the partial ``(m, ℓ, O)``:
      • intra-CUBE 8-way (row chain along intra_W + col bridge along
        intra_N) over the 2×4 PE grid;
      • inter-CUBE 2-phase lrab over the 4×2 CUBE sub-mesh (ADR-0060
        §4.2 lrab-adapted center-root reduce): Phase 1 row reduce
        converges at root_col=2; Phase 2 col reduce on root_col
        converges at root_row=1; root cube id = 6.
  - PE 0 of CUBE 6 stores ``O``.

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

Tensor layout:
  Q : (T_q, h_q · d_head) replicated; loaded as (G·T_q, d_head).
  K : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank
      loads its (d_head, S_local) slice via byte-conserving reshape
      of (S_local, h_kv·d_head) — correct for zero / symmetric
      inputs (ADR-0060 §3 reshape-not-transpose caveat).
  V : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank
      loads its (S_local, d_head) slice.
  O : (T_q, h_q · d_head) — only PE 0 of CUBE 6 stores.

Topology / SFR:
  - Requires ``configure_sfr_intercube_multisip`` (provides disjoint
    ``intra_*`` and ``E/W/N/S`` namespaces).
  - Intra-CUBE PEs arranged as a 2×4 grid (no wrap).
  - Inter-CUBE: 4×2 CUBE sub-mesh starting at origin (0, 0).
"""
from __future__ import annotations


TILE_S_KV = 1024  # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).

# lrab geometry for the C=8 single-KV-head group (4×2 cube sub-mesh).
_SUB_W = 4
_SUB_H = 2
_ROOT_COL = _SUB_W // 2   # 2
_ROOT_ROW = _SUB_H // 2   # 1
_ROOT_CUBE = _ROOT_ROW * _SUB_W + _ROOT_COL  # 6


def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
    """Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
    m_new = tl.maximum(m_local, m_other)
    scale_old = tl.exp(m_local - m_new)
    scale_new = tl.exp(m_other - m_new)
    l_new = l_local * scale_old + l_other * scale_new
    O_new = O_local * scale_old + O_other * scale_new
    return m_new, l_new, O_new


def gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel(
    q_ptr: int,
    k_ptr: int,
    v_ptr: int,
    o_ptr: int,
    T_q: int,
    S_kv: int,
    h_q: int,
    h_kv: int,
    d_head: int,
    C: int,
    P: int,
    *,
    tl,
) -> None:
    """Case-4 decode: Cube-SP × PE-SP, lrab center-root at cube 6."""
    G = h_q // h_kv
    n_ranks = C * P
    S_local = S_kv // n_ranks
    pe_id = tl.program_id(axis=0)
    cube_id = tl.program_id(axis=1)

    # ── Local attention (S_kv-axis tile sweep, ADR-0063 §A.2) ──
    Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
    n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
    KV_ROW_BYTES = d_head * 2  # f16

    # Tile 0: establishes persistent (m_local, l_local, O_local). Cannot
    # fold into Tiles 1..N loop (kernbench-only): persistent tensors
    # must live outside tl.scratch_scope or scope teardown discards
    # them; there's no scratch-backed (-inf, 0, 0) initializer.
    tile_s0 = min(TILE_S_KV, S_local)
    K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
    V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
    scores = tl.dot(Q, K_T)
    m_local = tl.max(scores, axis=-1)
    centered = scores - m_local
    exp_scores = tl.exp(centered)
    l_local = tl.sum(exp_scores, axis=-1)
    O_local = tl.dot(exp_scores, V)

    for tile_idx in range(1, n_tiles):
        tile_start = tile_idx * TILE_S_KV
        tile_s = min(TILE_S_KV, S_local - tile_start)
        with tl.scratch_scope():
            K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
                            shape=(d_head, tile_s), dtype="f16")
            V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
                          shape=(tile_s, d_head), dtype="f16")
            scores_t = tl.dot(Q, K_T_t)
            m_tile = tl.max(scores_t, axis=-1)
            centered_t = scores_t - m_tile
            exp_scores_t = tl.exp(centered_t)
            l_tile = tl.sum(exp_scores_t, axis=-1)
            O_tile = tl.dot(exp_scores_t, V_t)
            m_new, l_new, O_new = _merge_running(
                m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
            )
            tl.copy_to(m_local, m_new)
            tl.copy_to(l_local, l_new)
            tl.copy_to(O_local, O_new)

    # ── Intra-CUBE reduce: row chain (intra_W) + col bridge (intra_N) ──
    PE_GRID_COLS = 4
    pe_col = pe_id % PE_GRID_COLS
    pe_row = pe_id // PE_GRID_COLS
    pe_cols_used = min(PE_GRID_COLS, P)
    pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS

    if pe_cols_used > 1:
        if pe_col < pe_cols_used - 1:
            with tl.scratch_scope():
                m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16")
                l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16")
                O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16")
                m_new, l_new, O_new = _merge_running(
                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
                )
                tl.copy_to(m_local, m_new)
                tl.copy_to(l_local, l_new)
                tl.copy_to(O_local, O_new)
        if pe_col > 0:
            tl.send(dir="intra_W", src=m_local)
            tl.send(dir="intra_W", src=l_local)
            tl.send(dir="intra_W", src=O_local)

    if pe_col == 0 and pe_rows_used > 1:
        if pe_row < pe_rows_used - 1:
            with tl.scratch_scope():
                m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16")
                l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16")
                O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16")
                m_new, l_new, O_new = _merge_running(
                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
                )
                tl.copy_to(m_local, m_new)
                tl.copy_to(l_local, l_new)
                tl.copy_to(O_local, O_new)
        if pe_row > 0:
            tl.send(dir="intra_N", src=m_local)
            tl.send(dir="intra_N", src=l_local)
            tl.send(dir="intra_N", src=O_local)

    # ── Inter-CUBE lrab-adapted center-root reduce (ADR-0060 §4.2) ──
    # Only PE 0 of each CUBE participates. Adapts Phases 1-2 of
    # lrab_hierarchical_allreduce.py: bidirectional row reduce converges
    # at root_col; bidirectional col reduce on root_col converges at
    # root_row. Plain ``+`` replaced by log-sum-exp ``_merge_running``.
    if pe_id == 0:
        row = cube_id // _SUB_W
        col = cube_id % _SUB_W

        # Phase 1: row reduce — converge at col == _ROOT_COL.
        if col == 0:
            tl.send(dir="E", src=m_local)
            tl.send(dir="E", src=l_local)
            tl.send(dir="E", src=O_local)
        elif 0 < col < _ROOT_COL:
            with tl.scratch_scope():
                m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
                l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
                O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
                m_new, l_new, O_new = _merge_running(
                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
                )
                tl.copy_to(m_local, m_new)
                tl.copy_to(l_local, l_new)
                tl.copy_to(O_local, O_new)
            tl.send(dir="E", src=m_local)
            tl.send(dir="E", src=l_local)
            tl.send(dir="E", src=O_local)
        elif col == _ROOT_COL:
            with tl.scratch_scope():
                m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
                l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
                O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
                m_new, l_new, O_new = _merge_running(
                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
                )
                tl.copy_to(m_local, m_new)
                tl.copy_to(l_local, l_new)
                tl.copy_to(O_local, O_new)
            with tl.scratch_scope():
                m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
                l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
                O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
                m_new, l_new, O_new = _merge_running(
                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
                )
                tl.copy_to(m_local, m_new)
                tl.copy_to(l_local, l_new)
                tl.copy_to(O_local, O_new)
        elif _ROOT_COL < col < _SUB_W - 1:
            with tl.scratch_scope():
                m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
                l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
                O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
                m_new, l_new, O_new = _merge_running(
                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
                )
                tl.copy_to(m_local, m_new)
                tl.copy_to(l_local, l_new)
                tl.copy_to(O_local, O_new)
            tl.send(dir="W", src=m_local)
            tl.send(dir="W", src=l_local)
            tl.send(dir="W", src=O_local)
        elif col == _SUB_W - 1:
            tl.send(dir="W", src=m_local)
            tl.send(dir="W", src=l_local)
            tl.send(dir="W", src=O_local)

        # Phase 2: col reduce on col == _ROOT_COL — converge at row == _ROOT_ROW.
        if col == _ROOT_COL:
            if row == 0:
                tl.send(dir="S", src=m_local)
                tl.send(dir="S", src=l_local)
                tl.send(dir="S", src=O_local)
            elif 0 < row < _ROOT_ROW:
                with tl.scratch_scope():
                    m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
                    l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
                    O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
                    m_new, l_new, O_new = _merge_running(
                        m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
                    )
                    tl.copy_to(m_local, m_new)
                    tl.copy_to(l_local, l_new)
                    tl.copy_to(O_local, O_new)
                tl.send(dir="S", src=m_local)
                tl.send(dir="S", src=l_local)
                tl.send(dir="S", src=O_local)
            elif row == _ROOT_ROW:
                with tl.scratch_scope():
                    m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
                    l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
                    O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
                    m_new, l_new, O_new = _merge_running(
                        m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
                    )
                    tl.copy_to(m_local, m_new)
                    tl.copy_to(l_local, l_new)
                    tl.copy_to(O_local, O_new)
                if _SUB_H - 1 > _ROOT_ROW:
                    with tl.scratch_scope():
                        m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
                        l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
                        O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16")
                        m_new, l_new, O_new = _merge_running(
                            m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
                        )
                        tl.copy_to(m_local, m_new)
                        tl.copy_to(l_local, l_new)
                        tl.copy_to(O_local, O_new)
            elif _ROOT_ROW < row < _SUB_H - 1:
                with tl.scratch_scope():
                    m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
                    l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
                    O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16")
                    m_new, l_new, O_new = _merge_running(
                        m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
                    )
                    tl.copy_to(m_local, m_new)
                    tl.copy_to(l_local, l_new)
                    tl.copy_to(O_local, O_new)
                tl.send(dir="N", src=m_local)
                tl.send(dir="N", src=l_local)
                tl.send(dir="N", src=O_local)
            elif row == _SUB_H - 1 and _SUB_H - 1 > _ROOT_ROW:
                tl.send(dir="N", src=m_local)
                tl.send(dir="N", src=l_local)
                tl.send(dir="N", src=O_local)

    # ── Final normalise + store (root only: PE 0 of CUBE 6) ──
    if pe_id == 0 and cube_id == _ROOT_CUBE:
        O_final = O_local / l_local
        tl.store(o_ptr, O_final)
