"""Mesh-native 2D row-then-col AllReduce-mlo attention — decode (ADR-0059 extension).

Each cube holds the full Q (replicated) and 1/(mesh_rows * mesh_cols) of
KV (sequence-sharded across the 2D cube sub-mesh). The kernel decomposes
the AllReduce-mlo into a two-stage reduction:

  Stage 1 — Row reduce (E/W edges, ``mesh_cols - 1`` steps)
     Bidirectional ring within each row. After this stage every cube in
     row ``r`` holds the partial ``(m, ℓ, o)`` over the ``mesh_cols`` KV
     chunks in row ``r``.

  Stage 2 — Col reduce (N/S edges, ``mesh_rows - 1`` steps)
     Bidirectional ring within each column. After this stage every cube
     holds the partial over all ``mesh_rows × mesh_cols`` KV chunks —
     the AllReduce result.

The online-softmax mlo merge is associative, so row-then-col partitioning
of the reduction is mathematically equivalent to a 1D ring AllReduce-mlo
over all ``mesh_rows × mesh_cols`` cubes. The 2D form takes
``(mesh_cols - 1) + (mesh_rows - 1)`` steps instead of
``mesh_rows × mesh_cols - 1`` (e.g. 4 vs 7 at 2×4; 6 vs 15 at 4×4).

Designed to run on hardware wired by
``configure_sfr_intercube_multisip``, which installs both E/W and N/S
intra-SIP cube-mesh edges (``sfr_config.py:135-143``). The 1D
``_attention_mesh_mlo.py`` remains for the single_user PE-ring case;
this 2D variant supersedes it for multi_user_decode where the per-KV-group
cube count crosses a row boundary in the 4×4 cube mesh.

``mesh_rows = 1`` is supported as a degenerate row-only case so the
validation config (``_N_RANKS_MULTI_USER = 4`` → ``(1, 4)``) reduces to
the 1D ring's step count without behavioral change.
"""
from __future__ import annotations

from kernbench.common.pe_commands import TensorHandle


def _view(handle: TensorHandle, new_shape: tuple[int, ...]) -> TensorHandle:
    """Reshape — metadata only, no command emitted (cf. ``tl.trans``)."""
    return TensorHandle(
        id=handle.id,
        addr=handle.addr,
        shape=new_shape,
        dtype=handle.dtype,
        nbytes=handle.nbytes,
        data=handle.data,
        space=handle.space,
        pinned=handle.pinned,
    )


def _bidir_allreduce_mlo(
    m: TensorHandle,
    ell: TensorHandle,
    o: TensorHandle,
    rank: int,
    n_ranks: int,
    dir_pos: str,
    dir_neg: str,
    *,
    tl,
) -> tuple[TensorHandle, TensorHandle, TensorHandle]:
    """One bidirectional AllReduce-mlo ring along ``(dir_pos, dir_neg)``.

    Mirrors the 1D ``_attention_mesh_mlo.py`` algorithm but parameterized
    on direction labels so the 2D kernel can call it once with ``("E", "W")``
    for the row reduce and once with ``("S", "N")`` for the col reduce.
    Forwards the received triplets in subsequent steps so chunk ``c_i``
    reaches rank ``j`` at step ``|i - j|``.

    Returns the running ``(m, ℓ, o)`` after ``n_ranks - 1`` steps. Degenerate
    cases (``n_ranks <= 1``) are no-ops — the for-loop body simply does not
    execute.
    """
    has_pos = rank < n_ranks - 1
    has_neg = rank > 0

    to_send_pos_m: TensorHandle | None = m
    to_send_pos_ell: TensorHandle | None = ell
    to_send_pos_o: TensorHandle | None = o
    to_send_neg_m: TensorHandle | None = m
    to_send_neg_ell: TensorHandle | None = ell
    to_send_neg_o: TensorHandle | None = o

    for step in range(1, n_ranks):
        if has_pos and to_send_pos_m is not None:
            tl.send(dir=dir_pos, src=to_send_pos_m)
            tl.send(dir=dir_pos, src=to_send_pos_ell)
            tl.send(dir=dir_pos, src=to_send_pos_o)
        if has_neg and to_send_neg_m is not None:
            tl.send(dir=dir_neg, src=to_send_neg_m)
            tl.send(dir=dir_neg, src=to_send_neg_ell)
            tl.send(dir=dir_neg, src=to_send_neg_o)

        m_from_neg: TensorHandle | None = None
        ell_from_neg: TensorHandle | None = None
        o_from_neg: TensorHandle | None = None
        if has_neg and (rank - step) >= 0:
            m_from_neg = tl.recv(dir=dir_neg, shape=m.shape, dtype="f16")
            ell_from_neg = tl.recv(dir=dir_neg, shape=ell.shape, dtype="f16")
            o_from_neg = tl.recv(dir=dir_neg, shape=o.shape, dtype="f16")
            m_combined = tl.maximum(m, m_from_neg)
            scale_old = tl.exp(m - m_combined)
            scale_new = tl.exp(m_from_neg - m_combined)
            ell = ell * scale_old + ell_from_neg * scale_new
            o = o * scale_old + o_from_neg * scale_new
            m = m_combined

        m_from_pos: TensorHandle | None = None
        ell_from_pos: TensorHandle | None = None
        o_from_pos: TensorHandle | None = None
        if has_pos and (rank + step) < n_ranks:
            m_from_pos = tl.recv(dir=dir_pos, shape=m.shape, dtype="f16")
            ell_from_pos = tl.recv(dir=dir_pos, shape=ell.shape, dtype="f16")
            o_from_pos = tl.recv(dir=dir_pos, shape=o.shape, dtype="f16")
            m_combined = tl.maximum(m, m_from_pos)
            scale_old = tl.exp(m - m_combined)
            scale_new = tl.exp(m_from_pos - m_combined)
            ell = ell * scale_old + ell_from_pos * scale_new
            o = o * scale_old + o_from_pos * scale_new
            m = m_combined

        to_send_pos_m = m_from_neg
        to_send_pos_ell = ell_from_neg
        to_send_pos_o = o_from_neg
        to_send_neg_m = m_from_pos
        to_send_neg_ell = ell_from_pos
        to_send_neg_o = o_from_pos

    return m, ell, o


def attention_mesh_mlo_2d_kernel(
    q_ptr: int,
    k_ptr: int,
    v_ptr: int,
    o_ptr: int,
    S_q: int,
    S_kv_per_rank: int,
    h_q: int,
    h_kv: int,
    d_head: int,
    mesh_rows: int,
    mesh_cols: int,
    rank_axis: int = 0,
    cube_start: int = 0,
    *,
    tl,
) -> None:
    """2D row-then-col AllReduce-mlo decode kernel — see module docstring.

    ``rank_axis`` selects which program-id dimension carries the cube
    rank (matches the 1D kernel convention):

      0 — single_user_* (TL/BL): rank == tl.program_id(axis=0) (PE id).
          Not used at headline scale — single_user uses the 1D intra-cube
          PE ring (``_attention_mesh_mlo``). Kept here so the signature
          mirrors the 1D kernel.

      1 — multi_user_* (TR/BR): rank == tl.program_id(axis=1) (cube id).
          KV is split @ cubes inter-cube; the ring runs over the
          ``mesh_rows × mesh_cols`` cubes of one KV-group. The kernel
          gates ``pe_id != 0`` to return early — same v1 simplification
          as ``_attention_mesh_mlo`` (validation B=1).

    ``cube_start`` matches the value passed to ``DPPolicy.cube_start`` for
    the launch's tensor placement. kernbench's ``tl.program_id(axis=1)``
    returns the physical cube id (ADR-0022), so when the launch is
    offset within the SIP (e.g. cube_start=8 placing the second 2×4
    KV-group on cubes 8..15), the kernel must subtract ``cube_start``
    to recover the launch-local rank for ring arithmetic. Default 0
    preserves the cube_start=0 launches unchanged.
    """
    # For multi_user (rank_axis=1) only PE 0 in each cube runs the ring.
    if rank_axis != 0 and tl.program_id(axis=0) != 0:
        return

    rank = tl.program_id(axis=rank_axis)
    if rank_axis != 0:
        rank = rank - cube_start
    my_row = rank // mesh_cols
    my_col = rank % mesh_cols

    # Q is replicated on every cube — loaded once.
    Q = tl.load(q_ptr, shape=(S_q, h_q * d_head), dtype="f16")

    # Local KV chunk (sequence-sharded across the 2D sub-mesh).
    K = tl.load(k_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")
    V = tl.load(v_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")

    # ── One-shot local partial attention ──────────────────────────
    K_2d_T = _view(K, (h_q * d_head, S_kv_per_rank))
    V_2d = _view(V, (S_kv_per_rank, h_q * d_head))
    scores = tl.dot(Q, K_2d_T)
    m = tl.max(scores, axis=-1)
    P = tl.softmax(scores, axis=-1)
    scores_centered = scores - m
    exp_scores = tl.exp(scores_centered)
    ell = tl.sum(exp_scores, axis=-1)
    o = tl.dot(P, V_2d)

    # ── Stage 1: row AllReduce (E/W, mesh_cols - 1 steps) ─────────
    m, ell, o = _bidir_allreduce_mlo(
        m, ell, o, my_col, mesh_cols, "E", "W", tl=tl,
    )

    # ── Stage 2: col AllReduce (N/S, mesh_rows - 1 steps) ─────────
    # ``dir_pos="S"`` matches the SFR convention: ``S`` goes to higher
    # row (configure_sfr_intercube_multisip:140).
    m, ell, o = _bidir_allreduce_mlo(
        m, ell, o, my_row, mesh_rows, "S", "N", tl=tl,
    )

    # Final normalize: O := o / ℓ.
    O_final = o / ell
    tl.store(o_ptr, O_final)
