"""Local 2D window attention.

Every token attends to all tokens within a square window of Chebyshev radius ``r``
(i.e. ``|dy| <= r and |dx| <= r``) on its 2D grid.

Two numerically equivalent implementations are provided:

* a fused Triton kernel (CUDA only, forward pass only);
* PyTorch ``scaled_dot_product_attention`` with a boolean window mask (CPU, or when
  Triton is unavailable).

Both work on a packed batch: every sample is its own token grid of any size.
"""
import torch
import torch.nn.functional as F

try:  # Triton is optional; without it the SDPA path is used.
    import triton
    import triton.language as tl

    _HAS_TRITON = True
except Exception:  # pragma: no cover
    _HAS_TRITON = False

_LOG2E = 1.4426950408889634

# Global switch: "auto" (kernel when possible), "sdpa" (always masked SDPA).
BACKEND = "auto"


if _HAS_TRITON:

    @triton.jit
    def _window_attn_fwd(Q, K, V, O, DIMS, NH: tl.constexpr, NMAX, WY: tl.constexpr,
                         WX: tl.constexpr, R: tl.constexpr, D: tl.constexpr,
                         BR: tl.constexpr, qk_scale):
        # One program = one WY x WX tile of queries for one (batch, head).
        # Keys are streamed from the tile extended by R on every side ("halo").
        t, bh = tl.program_id(0), tl.program_id(1)
        b = bh // NH
        gyp = tl.load(DIMS + b * 4 + 0)
        gxp = tl.load(DIMS + b * 4 + 1)
        gyr = tl.load(DIMS + b * 4 + 2)
        gxr = tl.load(DIMS + b * 4 + 3)
        ntx = gxp // WX
        if t >= (gyp // WY) * ntx:
            return
        ty, tx = t // ntx, t % ntx
        base = bh * NMAX * D
        SX: tl.constexpr = WX + 2 * R
        SY: tl.constexpr = WY + 2 * R

        qi = tl.arange(0, WY * WX)
        qyr, qxr = qi // WX, qi % WX
        qy, qx = ty * WY + qyr, tx * WX + qxr
        d = tl.arange(0, D)
        qo = base + (qy * gxp + qx)[:, None] * D + d[None, :]
        q = tl.load(Q + qo)
        okq = (qy < gyr) & (qx < gxr)

        kxr = tl.arange(0, BR * SX) % SX
        okx = okq[:, None] & (tl.abs(qxr[:, None] - (kxr[None, :] - R)) <= R)
        kxg = tx * WX - R + kxr
        okg = (kxg >= 0) & (kxg < gxr)

        m = tl.full((WY * WX,), -1e30, tl.float32)
        l = tl.zeros((WY * WX,), tl.float32)
        acc = tl.zeros((WY * WX, D), tl.float32)

        for r0 in range(0, SY, BR):
            kyr = r0 + tl.arange(0, BR * SX) // SX
            kyg = ty * WY - R + kyr
            ok = okg & (kyg >= 0) & (kyg < gyr) & (kyr < SY)
            off = base + (kyg * gxp + kxg) * D
            k = tl.load(K + off[:, None] + d[None, :], mask=ok[:, None], other=0.0)
            v = tl.load(V + off[:, None] + d[None, :], mask=ok[:, None], other=0.0)

            s = tl.dot(q, tl.trans(k)) * qk_scale
            good = okx & (tl.abs(qyr[:, None] - (kyr[None, :] - R)) <= R) & ok[None, :]
            s = tl.where(good, s, -1e30)

            mn = tl.maximum(m, tl.max(s, 1))
            p = tl.where(good, tl.math.exp2(s - mn[:, None]), 0.0)
            a = tl.math.exp2(m - mn)
            l = l * a + tl.sum(p, 1)
            acc = acc * a[:, None] + tl.dot(p.to(v.dtype), v).to(tl.float32)
            m = mn

        lc = tl.maximum(l, 1e-30)
        tl.store(O + qo, (acc / lc[:, None]).to(O.dtype.element_ty))


def _tile_width(r):
    # The kernel needs a power-of-two key strip width (WX + 2R).
    for wx in (16, 8):
        sx = wx + 2 * r
        if sx & (sx - 1) == 0:
            return wx
    raise ValueError(f"window radius {r}: no supported tile width")


@torch._dynamo.disable
def packed_window_attention(q, k, v, r, dims, max_tiles):
    """Windowed attention for a packed batch of differently shaped token grids.

    Sample ``b`` is stored row-major in a padded ``gyp x gxp`` grid at the start of the
    sequence, with its real ``gyr x gxr`` tokens in the top-left corner; the rest is padding.
    dims: (B, 4) int32 ``[gyp, gxp, gyr, gxr]`` (gyp a multiple of 8, gxp of the tile width).
    q, k, v: (B, heads, Nmax, head_dim). Padding outputs are zero (kernel) or unused (SDPA).
    """
    B, nh, N, D = q.shape
    if BACKEND == "auto" and _HAS_TRITON and q.is_cuda and D in (16, 32, 64, 128):
        q, k, v = (t.to(torch.bfloat16).contiguous() for t in (q, k, v))
        o = torch.zeros_like(q)
        _window_attn_fwd[(max_tiles, B * nh)](
            q, k, v, o, dims.to(q.device, torch.int32).contiguous(), NH=nh, NMAX=N, WY=8,
            WX=_tile_width(r), R=r, D=D, BR=4, qk_scale=D ** -0.5 * _LOG2E, num_warps=8, num_stages=3)
        return o
    # reference path: per-sample block mask over the packed sequence
    n = torch.arange(N, device=q.device)
    d = dims.to(q.device).long()
    yy, xx = n[None] // d[:, 1:2], n[None] % d[:, 1:2]                      # (B, N)
    valid = (yy < d[:, 2:3]) & (xx < d[:, 3:4])
    win = ((yy[:, :, None] - yy[:, None]).abs() <= r) & ((xx[:, :, None] - xx[:, None]).abs() <= r)
    m = win & valid[:, None, :]
    m = m | (~valid[:, :, None] & torch.eye(N, dtype=torch.bool, device=q.device))
    return F.scaled_dot_product_attention(q.contiguous(), k.contiguous(), v.contiguous(),
                                          attn_mask=m[:, None])