"""GQA decode kernel: composite + softmax_merge fused variant (3).

Identical mapping to ``_gqa_attention_decode_short.py``. Difference vs
the GEMM-only composite: per-tile softmax is folded into
the P·V composite via the ``softmax_merge`` prologue recipe, eliminating
the GEMM/MATH engine bubble that a flat primitives chain incurs.

Per-tile structure (ADR-0065 two-composite pattern, cf. decode_opt2):
  tile 0   primitives establish (m, ℓ, O)
  tile k>0  #1 Q·Kᵀ composite        → scores (pinned primary-out)
            #2 softmax_merge prologue → P (pinned, auto-binds to GEMM a)
               + P·V composite, ``out=O_local``
               + add epilogue folding the result into O_local

The intra-group chain reduce after the tile loop is unchanged.

Three-variant comparison:
  (1) without composite                    : ``_gqa_attention_decode_short.py``
  (2) with composite (GEMM-only, no fuse)  : ``…_composite.py``
  (3) with composite + softmax_merge fuse  : this file

Shard addressing, layouts, and caller contract are identical to the
first-level decode kernel (ADR-0011 D-VA1).
"""
from __future__ import annotations


TILE_S_KV = 1024


def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
                     h_q: int, h_kv: int, S_kv: int) -> None:
    """Validate composite-decode config — caller-side, sim-cost 0.

    Mirrors first-level decode ``_validate_config``.
    """
    if kv_per_cube not in (1, 2, 4, 8):
        raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
    if T_q != 1:
        raise ValueError(f"decode requires T_q == 1; got {T_q}")
    if P != 8:
        raise ValueError(f"2x4 mesh chain-reduce geometry requires P == 8; got {P}")
    if h_q % h_kv != 0:
        raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
    if C != h_kv // kv_per_cube:
        raise ValueError(
            f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
            f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
        )
    group_size = P // kv_per_cube
    if S_kv % (group_size * TILE_S_KV) != 0:
        raise ValueError(
            f"S_kv ({S_kv}) must be a multiple of group_size·TILE_S_KV "
            f"({group_size * TILE_S_KV}); each PE's sequence shard must "
            f"be a whole number of tiles"
        )


def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
    """Used only by the intra-group chain reduce below (recipe handles
    per-tile fold internally)."""
    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_short_composite_fused_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,
    kv_per_cube: int,
    *,
    tl,
) -> None:
    """Composite-GEMM decode — same mapping as first-level + tl.composite.

    Caller must invoke ``_validate_config(...)`` first.
    """
    group_size = P // kv_per_cube
    S_local = S_kv // group_size
    n_tiles_per_pe = S_local // TILE_S_KV

    G = h_q // h_kv
    pe_id = tl.program_id(axis=0)
    cube_id = tl.program_id(axis=1)
    pe_in_group = pe_id % group_size
    group_id_in_cube = pe_id // group_size

    KV_ROW_BYTES = d_head * 2
    K_TILE_BYTES = d_head * TILE_S_KV * 2
    Q_ROW_BYTES = G * d_head * 2
    K_HEAD_BYTES = S_kv * d_head * 2
    V_HEAD_BYTES = S_kv * KV_ROW_BYTES

    # Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
    q_base = q_ptr + cube_id * kv_per_cube * Q_ROW_BYTES + group_id_in_cube * Q_ROW_BYTES
    k_shard_base = (k_ptr
                    + cube_id * kv_per_cube * K_HEAD_BYTES
                    + group_id_in_cube * K_HEAD_BYTES
                    + pe_in_group * n_tiles_per_pe * K_TILE_BYTES)
    v_shard_base = (v_ptr
                    + cube_id * kv_per_cube * V_HEAD_BYTES
                    + group_id_in_cube * V_HEAD_BYTES
                    + pe_in_group * n_tiles_per_pe * TILE_S_KV * KV_ROW_BYTES)

    Q = tl.load(q_base, shape=(G, d_head), dtype="f16")

    # ── Tile 0: establish running (m, ℓ, O) with primitives ──────────
    # (Reference: decode_opt2 — running state is set up with tl.dot/MATH
    # primitives, not composite. Recipe-driven composite enters in tile 1+.)
    K_T = tl.load(k_shard_base, shape=(d_head, TILE_S_KV), dtype="f16")
    V = tl.load(v_shard_base, shape=(TILE_S_KV, 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)

    # ── Tiles 1..n_tiles_per_pe-1: two composites per tile ──────────
    # #1  Q·Kᵀ composite → scores (pinned primary-out, fed into #2).
    # #2  softmax_merge prologue + P·V GEMM + add epilogue, all in one
    #     composite: updates (m, ℓ) in place, computes P, runs P·V with
    #     pinned auto-bind, and folds the result into O_local.
    for tile_idx in range(1, n_tiles_per_pe):
        with tl.scratch_scope():
            K_T_t = tl.load(k_shard_base + tile_idx * K_TILE_BYTES,
                            shape=(d_head, TILE_S_KV), dtype="f16")
            V_t = tl.load(v_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
                          shape=(TILE_S_KV, d_head), dtype="f16")
            scores_t = tl.composite(op="gemm", a=Q, b=K_T_t)
            tl.composite(
                prologue=[{"op": "softmax_merge", "s": scores_t,
                           "m": m_local, "l": l_local, "O": O_local}],
                op="gemm", b=V_t, out=O_local,
                epilogue=[{"op": "add", "other": O_local}],
            )

    # ── Chain reduce ──
    group_cols = min(4, group_size)
    group_rows = (group_size + group_cols - 1) // group_cols
    pe_col_in_group = pe_in_group % group_cols
    pe_row_in_group = pe_in_group // group_cols

    if group_cols > 1:
        if pe_col_in_group < group_cols - 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_in_group > 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_in_group == 0 and group_rows > 1:
        if pe_row_in_group < group_rows - 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_in_group > 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)

    if pe_in_group == 0:
        O_final = O_local / l_local
        o_base = (o_ptr
                  + cube_id * kv_per_cube * Q_ROW_BYTES
                  + group_id_in_cube * Q_ROW_BYTES)
        tl.store(o_base, O_final)
