"""Shared (m, ℓ, O) reduce for the Case-6 long-context decode kernels.

Extracted from the Cube-SP × PE-SP decode kernel so the four command-form
variants (primitive / primitive-tiled / composite / composite_extended)
share one identical reduce. The local attention differs per variant;
the reduce does not. Byte-equal behavior at C=8 is guarded by
``tests/attention/test_gqa_decode_long_ctx_composite.py``.

The reduce is two-level:
  • 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-adapted center-root reduce over a
    ``sub_w × sub_h`` CUBE sub-mesh (ADR-0060 §4.2): Phase 1 does a
    bidirectional row reduce converging at ``root_col = sub_w // 2``;
    Phase 2 does a bidirectional col reduce on ``root_col`` converging
    at ``root_row = sub_h // 2``.

The submesh dimensions ``(sub_w, sub_h)`` and root are **derived from
C** (number of cubes launched per KV group), not hardcoded, so the
same reduce works for any ``C ≥ 1``. Peer-existence guards on every
send/recv elide operations that would target un-launched neighbors
(needed for non-rectangular C like 7 or 12; safe no-ops when the
neighbor exists). For ``C = 8`` this reproduces the previous
hardcoded ``{sub_w=4, sub_h=2, root_col=2, root_row=1, root_cube=6}``
exactly.

Plain ``+`` in the tree is replaced by log-sum-exp ``_merge_running``.
"""
from __future__ import annotations


PE_GRID_COLS = 4


def _submesh_for(C):
    """Derive ``(sub_w, sub_h, root_col, root_row, root_cube)`` for a
    C-cube inter-cube reduce.

    Submesh is at most 4 columns wide (matching the SIP's 4×4 cube
    mesh); rows fill left-to-right, top-to-bottom. For the Case-6
    baseline C=8, this returns ``(4, 2, 2, 1, 6)`` — identical to
    the previously hardcoded values.
    """
    sub_w = min(C, 4)
    sub_h = (C + sub_w - 1) // sub_w
    root_col = sub_w // 2
    root_row = sub_h // 2
    root_cube = root_row * sub_w + root_col
    return sub_w, sub_h, root_col, root_row, root_cube


def root_cube_for(C):
    """Root cube id for a C-cube reduce. Callers use this to gate the
    final ``tl.store`` in the kernel epilogue."""
    _, _, _, _, root_cube = _submesh_for(C)
    return root_cube


# Backward-compat alias: preserves the Case-6 baseline value (6) for
# any external code that still imports the module-level constant.
_ROOT_CUBE = root_cube_for(8)


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 reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, *, tl):
    """Two-level (m, ℓ, O) reduce-to-root (PE 0 of ``root_cube_for(C)``).
    In place.

    Updates ``m_local`` / ``l_local`` / ``O_local`` in place via
    ``copy_to``; after this call only PE 0 of ``root_cube_for(C)``
    holds the fully-merged result.

    Peer-existence guards on every inter-cube send/recv ensure the
    reduce completes for any launched C (rectangular or not); a
    send/recv toward a non-launched neighbor is skipped. For C = 8
    every neighbor exists → guards are all True → byte-equal to the
    previously hardcoded Case-6 reduce.
    """
    sub_w, sub_h, root_col, root_row, _root_cube = _submesh_for(C)

    # ── Intra-CUBE reduce (unchanged — depends only on P) ─────────────
    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 ────────────────────
    # 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``. Peer-existence guards elide sends/recvs whose
    # neighbor cube was not launched.
    if pe_id == 0:
        row = cube_id // sub_w
        col = cube_id % sub_w

        # Peer existence within the launched set of C cubes.
        east_exists  = (col + 1 < sub_w) and (cube_id + 1 < C)
        west_exists  = col > 0
        south_exists = cube_id + sub_w < C
        north_exists = row > 0

        # Phase 1: row reduce — converge at col == root_col.
        if col == 0:
            if east_exists:
                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:
            if west_exists:
                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 east_exists:
                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 west_exists:
                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 east_exists:
                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:
            if east_exists:
                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)
            if west_exists:
                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:
            if west_exists:
                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:
                if south_exists:
                    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:
                if north_exists:
                    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 south_exists:
                    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 north_exists:
                    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 and south_exists:
                    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:
                if south_exists:
                    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)
                if north_exists:
                    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:
                if north_exists:
                    tl.send(dir="N", src=m_local)
                    tl.send(dir="N", src=l_local)
                    tl.send(dir="N", src=O_local)
