"""GQA decode kernel — Case 4 (Cube-SP × PE-TP on d_head).

Per the per-PE memory/comm summary chart  (paper_plot_gqa_4cases_summary.py):
this is the d_head-TP variant — KV is sharded 64-way (cube=row_wise on
S_kv, pe=column_wise on d_head), so per-PE KV memory matches Case 6 ★
(640 MB / PE at 1M context across 80 layers), but the attention
algorithm needs an extra mid-layer partial-score AllReduce because
each PE only owns d_head/P dimensions of K/V.

Algorithm (single layer, decode T_q=1, B=1):
  1. Each PE loads its own d_head/P slice of Q, K, V (Q replicated
     in HBM; PE-load via column-wise DPPolicy carves out the slice).
  2. Per S_kv tile, each PE computes a PARTIAL score tile:
        partial_score = Q_shard · K_shard^T  shape: (G·T_q, S_local)
  3. AR-SUM these partial scores INTRA-CUBE (across the P PEs of a
     cube) so every PE has the full Q · K^T for the cube's S_local
     tokens. Online-softmax merge folds tiles into a running state.
  4. After all tiles, each PE has (m_i, ℓ_i, O_i_shard) for cube i,
     where O_i_shard is the cube's partial attention output on this
     PE's d_head/P slice.
  5. INTER-CUBE online-softmax merge of (m, ℓ, O_shard) across C
     cubes (lrab center-root reduce). After this, every PE has the
     full-sequence attention output for its d_head/P slice.
  6. Final normalise + store: the P PEs of the root cube each write
     their (G·T_q, d_head/P) slice of the output.

Tensor layout:
  Q : (T_q, h_q · d_head)  cube=replicate, pe=column_wise on d_head
      → per-PE shape (T_q, h_q · d_head/P), byte-reshaped to
        (G·T_q, d_head/P) for the per-PE GEMM input.
  K : (S_kv, h_kv · d_head) cube=row_wise on S_kv, pe=column_wise
      on d_head → per-PE shape (S_local, d_head/P), byte-reshaped to
      (d_head/P, S_local) for the K^T-form GEMM.
  V : same layout as K.
  O : (T_q, h_q · d_head) cube=replicate, pe=column_wise on d_head
      → each writer PE owns its d_head/P slice.

Topology / SFR:
  - Requires ``configure_sfr_intercube_multisip`` (provides disjoint
    ``intra_*`` and ``E/W/N/S`` namespaces).
  - Intra-CUBE PEs in a 2×4 grid; AR-SUM of partial scores uses a
    reduce-to-PE0-then-broadcast pattern over intra_E/W + intra_N/S.
  - Inter-CUBE lrab center-root reduce on (m, ℓ, O_shard) over the
    4×2 cube sub-mesh (root cube = 6).

Naming note: this kernel exemplifies the chart's "PE-TP" label
meaning **d_head TP** (NOT the batch-axis TP of the original slide-17
Case 1). The chart-side Case label is "Case 4  Cube-SP × PE-TP".

ADR-0065 P5 scope note: this kernel is meant for op_log / latency
measurement (matches the milestone-1h-gqa sweep mode). Numeric
data-mode parity is a separate follow-up.
"""
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.

    O is the per-PE d_head-slice partial; merge is element-wise so it
    composes naturally even with O being a d_head shard.
    """
    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 _intra_cube_allreduce_sum(tensor, *, tl,
                              pe_col: int, pe_row: int,
                              pe_cols_used: int, pe_rows_used: int):
    """Sum-AllReduce a tensor across the 8 PEs of one cube (2×4 grid).

    Reduce-to-PE0 (row chain west + col bridge north) followed by
    broadcast-from-PE0 (col bridge south + row chain east). Each phase
    operates in place on ``tensor`` via ``tl.copy_to``.

    Used here to combine PARTIAL Q·K^T score tiles across the P=8 PEs
    of a cube — each PE has only d_head/P dims, so each contributes a
    partial dot product; the full score = elementwise sum across PEs.
    """
    # ── Phase 1a: row reduce (west) — accumulate at pe_col=0 ────────
    if pe_cols_used > 1:
        if pe_col == pe_cols_used - 1:
            tl.send(dir="intra_W", src=tensor)
        elif 0 < pe_col < pe_cols_used - 1:
            with tl.scratch_scope():
                other = tl.recv(dir="intra_E", shape=tensor.shape, dtype="f16")
                tl.copy_to(tensor, tensor + other)
            tl.send(dir="intra_W", src=tensor)
        elif pe_col == 0:
            with tl.scratch_scope():
                other = tl.recv(dir="intra_E", shape=tensor.shape, dtype="f16")
                tl.copy_to(tensor, tensor + other)

    # ── Phase 1b: col reduce on pe_col=0 (north) — accumulate at row=0 ─
    if pe_col == 0 and pe_rows_used > 1:
        if pe_row == pe_rows_used - 1:
            tl.send(dir="intra_N", src=tensor)
        elif pe_row == 0:
            with tl.scratch_scope():
                other = tl.recv(dir="intra_S", shape=tensor.shape, dtype="f16")
                tl.copy_to(tensor, tensor + other)

    # ── Phase 2a: col broadcast (south) — row=0 → row=1 on pe_col=0 ─
    if pe_col == 0 and pe_rows_used > 1:
        if pe_row == 0:
            tl.send(dir="intra_S", src=tensor)
        elif pe_row > 0:
            with tl.scratch_scope():
                received = tl.recv(dir="intra_N", shape=tensor.shape, dtype="f16")
                tl.copy_to(tensor, received)

    # ── Phase 2b: row broadcast (east) — pe_col=0 → all on each row ─
    if pe_cols_used > 1:
        if pe_col == 0:
            tl.send(dir="intra_E", src=tensor)
        elif 0 < pe_col < pe_cols_used - 1:
            with tl.scratch_scope():
                received = tl.recv(dir="intra_W", shape=tensor.shape, dtype="f16")
                tl.copy_to(tensor, received)
            tl.send(dir="intra_E", src=tensor)
        elif pe_col == pe_cols_used - 1:
            with tl.scratch_scope():
                received = tl.recv(dir="intra_W", shape=tensor.shape, dtype="f16")
                tl.copy_to(tensor, received)


def gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_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-TP-on-d_head, lrab center-root at cube 6."""
    G = h_q // h_kv
    if d_head % P != 0:
        raise ValueError(
            f"Case 4 requires d_head={d_head} divisible by P={P}"
        )
    if S_kv % C != 0:
        raise ValueError(
            f"Case 4 requires S_kv={S_kv} divisible by C={C}"
        )

    d_head_local = d_head // P   # 128 / 8 = 16 in the headline setup
    S_local = S_kv // C
    pe_id = tl.program_id(axis=0)
    cube_id = tl.program_id(axis=1)

    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

    KV_ROW_BYTES = d_head_local * 2  # bytes per K/V row at this PE's d_head slice

    # ── Load this PE's d_head slice of Q (byte-conserving reshape) ──
    # Q in HBM is (T_q, h_q · d_head); the column_wise pe shard gives
    # each PE (T_q, h_q · d_head/P). Reshape to (G·T_q, d_head/P) for the
    # GEMM: byte-conserving since G = h_q (h_kv = 1).
    Q = tl.load(q_ptr, shape=(G * T_q, d_head_local), dtype="f16")

    # ── Local attention (S_kv-axis tile sweep) with intra-cube AR sum ─
    n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV

    # Tile 0: establish persistent (m, ℓ, O_shard). O_shard is the
    # per-PE d_head/P slice of the partial attention output.
    tile_s0 = min(TILE_S_KV, S_local)
    K_T = tl.load(k_ptr, shape=(d_head_local, tile_s0), dtype="f16")
    V = tl.load(v_ptr, shape=(tile_s0, d_head_local), dtype="f16")
    partial_scores0 = tl.dot(Q, K_T)
    # AR-SUM partial scores intra-cube → every PE in the cube has the
    # full Q · K^T_tile0 scores.
    _intra_cube_allreduce_sum(
        partial_scores0, tl=tl,
        pe_col=pe_col, pe_row=pe_row,
        pe_cols_used=pe_cols_used, pe_rows_used=pe_rows_used,
    )
    m_local = tl.max(partial_scores0, axis=-1)
    exp_scores = tl.exp(partial_scores0 - m_local)
    l_local = tl.sum(exp_scores, axis=-1)
    # Partial output sharded on d_head: (G·T_q, d_head/P) per PE.
    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_local, tile_s), dtype="f16")
            V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
                          shape=(tile_s, d_head_local), dtype="f16")
            partial_scores_t = tl.dot(Q, K_T_t)
            _intra_cube_allreduce_sum(
                partial_scores_t, tl=tl,
                pe_col=pe_col, pe_row=pe_row,
                pe_cols_used=pe_cols_used, pe_rows_used=pe_rows_used,
            )
            m_tile = tl.max(partial_scores_t, axis=-1)
            exp_scores_t = tl.exp(partial_scores_t - m_tile)
            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)

    # ── Inter-CUBE lrab-adapted center-root reduce on (m, ℓ, O_shard) ─
    # Every PE participates (each contributes its own d_head/P slice
    # of O_shard); the merge is element-wise so d_head-sharding is
    # preserved through the AR.
    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 cube: each P-PE writes its d_head slice) ─
    if cube_id == _ROOT_CUBE:
        O_final = O_local / l_local
        tl.store(o_ptr, O_final)
