"""ldsa — fused Local Dense Synthesizer Attention window op as a noarch Triton kernel.

Given synthesized per-position window logits `a_logits [B, H, W, T]` and values
`v [B, H, Dh, T]` with a window `(left, right)`, `W = left + right + 1`:

    a[b,h,k,t]   = softmax over k of a_logits[b,h,:,t]        # softmax across the window
    out[b,h,d,t] = sum_k a[b,h,k,t] * v[b,h,d, t-left+k]      # weighted sum; zero outside [0,T)

i.e. local attention whose weights are SYNTHESIZED per frame (an MLP over the input) rather
than from query-key dot products (Synthesizer, arXiv 2005.00743). The kernel fuses the softmax
+ windowed weighted-sum into one launch (fp32 accumulation); it does NOT include the weight/
value/output projections — the caller supplies `a_logits` and `v`.

    from kernels import get_kernel
    k = get_kernel("futo-org/ldsa", version=1)
    out = k.ldsa_local_attention(a_logits, v, left, right)   # [B,H,Dh,T]; contiguous, any float

Inference kernel (no backward): the Triton path runs only on CUDA under `torch.no_grad()`;
grad-enabled / CPU calls take the autograd-safe `eager_ldsa` reference (also exported).

Reference: M. Xu, S. Li, X.-L. Zhang, "Transformer-based End-to-End Speech Recognition with
Local Dense Synthesizer Attention," ICASSP 2021 (arXiv:2010.12155;
github.com/mlxu995/multihead-LDSA). Differences from that reference are in the model card.
"""

from __future__ import annotations

import torch
import torch.nn.functional as F

# Hard dependency: this module IS the Triton kernel path (imported only when GPU kernels are
# opted into). A missing triton must fail LOUDLY here, not silently degrade to eager (which
# hides a broken serving setup). The eager reference stays for the legitimate RUNTIME fallbacks
# (CPU / grad-enabled), which are correctness, not error-hiding.
import triton
import triton.language as tl


def eager_ldsa(a_logits: torch.Tensor, v: torch.Tensor, left: int, right: int) -> torch.Tensor:
    """Autograd-safe reference — bit-matches LocalSynthAttention's softmax+pad+window loop.
    a_logits [B,H,W,T], v [B,H,Dh,T] -> out [B,H,Dh,T]."""
    a = a_logits.softmax(dim=2)
    t = v.shape[3]
    vp = F.pad(v, (left, right))
    out = a[:, :, 0:1, :] * vp[:, :, :, 0:t]
    for k in range(1, a.shape[2]):
        out = out + a[:, :, k : k + 1, :] * vp[:, :, :, k : k + t]
    return out


def _configs():
    return [triton.Config({"BLOCK_T": bt}, num_warps=nw) for bt in (64, 128, 256) for nw in (4, 8)]


@triton.autotune(configs=_configs(), key=["T", "Dh", "W"])
@triton.jit
def _ldsa_kernel(
    a_ptr,
    v_ptr,
    out_ptr,
    T,
    Dh,
    left,
    stride_a_bh,
    stride_v_bh,
    W: tl.constexpr,
    BLOCK_T: tl.constexpr,
    BLOCK_DH: tl.constexpr,
    BLOCK_W: tl.constexpr,
):
    bh = tl.program_id(0)
    offs_t = tl.program_id(1) * BLOCK_T + tl.arange(0, BLOCK_T)  # output frames
    mt = offs_t < T
    a_base = a_ptr + bh * stride_a_bh

    # softmax stats over the window (axis 0), one max/sum per output column
    offs_w = tl.arange(0, BLOCK_W)
    a_tile = tl.load(
        a_base + offs_w[:, None] * T + offs_t[None, :],
        mask=(offs_w < W)[:, None] & mt[None, :],
        other=float("-inf"),
    ).to(tl.float32)  # [BLOCK_W, BLOCK_T]
    m = tl.max(a_tile, axis=0)  # [BLOCK_T]
    z = tl.sum(tl.exp(a_tile - m[None, :]), axis=0)  # [BLOCK_T]

    # accumulate the windowed weighted sum into a [Dh, BLOCK_T] tile
    offs_dh = tl.arange(0, BLOCK_DH)
    md = offs_dh < Dh
    v_base = v_ptr + bh * stride_v_bh
    acc = tl.zeros((BLOCK_DH, BLOCK_T), dtype=tl.float32)
    for k in tl.static_range(W):
        a_k = tl.load(a_base + k * T + offs_t, mask=mt, other=float("-inf")).to(tl.float32)
        wk = tl.exp(a_k - m) / z  # softmax weight for window pos k, [BLOCK_T]
        src = offs_t - left + k  # source frame per output column
        vk = tl.load(
            v_base + offs_dh[:, None] * T + src[None, :],
            mask=md[:, None] & (src >= 0)[None, :] & (src < T)[None, :] & mt[None, :],
            other=0.0,
        ).to(tl.float32)  # [BLOCK_DH, BLOCK_T]
        acc += wk[None, :] * vk
    tl.store(
        out_ptr + bh * stride_v_bh + offs_dh[:, None] * T + offs_t[None, :],
        acc,
        mask=md[:, None] & mt[None, :],
    )


def _triton_ldsa(a_logits: torch.Tensor, v: torch.Tensor, left: int) -> torch.Tensor:
    b, h, w, t = a_logits.shape
    dh = v.shape[2]
    a_c, v_c = a_logits.contiguous(), v.contiguous()
    out = torch.empty_like(v_c)
    grid = lambda meta: (b * h, triton.cdiv(t, meta["BLOCK_T"]))  # noqa: E731
    _ldsa_kernel[grid](
        a_c,
        v_c,
        out,
        t,
        dh,
        left,
        w * t,
        dh * t,
        W=w,
        BLOCK_DH=triton.next_power_of_2(dh),
        BLOCK_W=triton.next_power_of_2(w),
    )
    return out


def ldsa_local_attention(
    a_logits: torch.Tensor, v: torch.Tensor, left: int, right: int
) -> torch.Tensor:
    """Fused LDSA window op; Triton on CUDA under no_grad, else the autograd-safe eager
    reference (training, CPU). a_logits [B,H,W,T], v [B,H,Dh,T] -> [B,H,Dh,T]."""
    if v.is_cuda and not torch.is_grad_enabled():
        return _triton_ldsa(a_logits, v, left)
    return eager_ldsa(a_logits, v, left, right)


__all__ = ["ldsa_local_attention", "eager_ldsa"]