"""GQA decode kernel — Case 1 (Cube-SP × PE-TP).

Per GQA_full_deck.pptx slide 11:
  - K, V split S_kv-wise across the 8 cubes (cube=row_wise);
    replicated within each cube (pe=replicate). Each active rank
    owns S_local = S_kv / C.
  - PEs nominally split on the batch dim (PE-TP). At B=1 (the
    slide-11 single-user default) only PE 0 of each cube has work;
    PEs 1-7 of every cube idle (PE-TP-at-B=1 waste).
  - NO intra-CUBE communication (only PE 0 holds a valid partial).
  - Inter-CUBE 8-way reduce on (m, ℓ, O) via the lrab-adapted
    center-root pattern (ADR-0060 §4.2): Phase 1 row reduce
    converges at root_col; Phase 2 col reduce on root_col converges
    at root_row. For sub_w=4, sub_h=2: root_col=2, root_row=1,
    root_cube=6 (lrab geometric center).

Tensor layout:
  Q : (T_q, h_q · d_head) replicated on every rank.
  K : (S_kv, h_kv · d_head) with cube=row_wise, pe=replicate —
      each cube's PE 0 owns the full (S_local, h_kv·d_head) shard.
  V : same as K.
  O : (T_q, h_q · d_head) — only PE 0 of CUBE 6 stores.

Topology / SFR:
  - Requires ``configure_sfr_intercube_multisip`` for the E/W/N/S
    inter-CUBE lanes used by the lrab reduce.
"""
from __future__ import annotations


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


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_tp_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-1 decode: PE 0 per cube does attention; inter-CUBE lrab AR only."""
    # B=1 single-user decode + PE-TP: only PE 0 per cube has work.
    pe_id = tl.program_id(axis=0)
    cube_id = tl.program_id(axis=1)
    if pe_id != 0:
        return

    # Cube-SP geometry: 4×2 sub-mesh, lrab center-root cube = 6.
    sub_w = 4
    sub_h = C // sub_w
    if sub_w < 2 or sub_h < 2 or sub_w * sub_h != C:
        raise ValueError(
            f"Case 1 requires C decomposable as sub_w=4, sub_h>=2; got C={C}"
        )
    root_col = sub_w // 2
    root_row = sub_h // 2
    root_cube = root_row * sub_w + root_col

    G = h_q // h_kv
    S_local = S_kv // C  # cube=row_wise, pe=replicate ⇒ S_local per cube
    KV_ROW_BYTES = d_head * 2  # f16

    # ── 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

    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)

    # ── Inter-CUBE lrab-adapted center-root reduce (ADR-0060 §4.2) ──
    row = cube_id // sub_w
    col = cube_id % sub_w

    # Phase 1: row reduce — converge at col == root_col.
    if col == 0 and root_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:
        if root_col > 0:
            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)
        if sub_w - 1 > root_col:
            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 and sub_w - 1 > root_col:
        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 and root_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:
            if root_row > 0:
                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 cube_id == root_cube:
        O_final = O_local / l_local
        tl.store(o_ptr, O_final)
