"""Kimi Delta Attention (KDA) forward, chunk form -- custom Triton recurrence.

The reference (reference.py) splits the work into:
  (a) intra-chunk quantities that depend only on one chunk's q,k,v,g,beta,
      most importantly the delta-rule inverse A = (I - M)^{-1} . beta_key
      with M = -(decayed k-k similarity) . beta_q (strictly lower triangular),
      and the decayed q-k similarity Aqk; and
  (b) the inter-chunk recurrent read/write of the hidden state S (K x V)
      that threads all chunks of a (batch, head) together.

We compute (a) exactly on the host in fp32 (the delta-rule inverse via a
batched lower-triangular solve -- stable and exact), and implement (b) as a
custom Triton kernel: one thread block per (batch, head) loops over the chunks
keeping S in registers and performing every GEMM in fp32 (Hopper fp32 MMA
with fp32 accumulation) so the sequential recurrence matches the reference's
fp32 trajectory to well within the 0.05 tolerance.
"""
from __future__ import annotations

import torch
import torch.nn as nn
from einops import rearrange

import triton
import triton.language as tl

OP_TYPE = "linear_attention"
SUPPORTED_PRECISIONS = ["bf16"]
HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]


@triton.jit
def kda_recur_kernel(
    A_ptr, w_ptr, u_ptr, Aqk_ptr, qg_ptr, gd_ptr, delta_ptr, o_ptr,
    B, H, NT,
    BT: tl.constexpr, K: tl.constexpr, V: tl.constexpr,
    stride_Ab, stride_Ah, stride_An, stride_Ac, stride_Ak,
    stride_wb, stride_wh, stride_wn, stride_wc, stride_wk,
    stride_ub, stride_uh, stride_un, stride_uc, stride_uv,
    stride_qb, stride_qh, stride_qn, stride_qc, stride_qk,
    stride_gb, stride_gh, stride_gn, stride_gk,
    stride_db, stride_dh, stride_dn, stride_dc, stride_dk,
    stride_ob, stride_oh, stride_on, stride_oc, stride_ov,
):
    pid = tl.program_id(0)
    h = pid % H
    b = pid // H

    Ab_ = A_ptr + b * stride_Ab + h * stride_Ah
    wb_ = w_ptr + b * stride_wb + h * stride_wh
    ub_ = u_ptr + b * stride_ub + h * stride_uh
    Aqb_ = Aqk_ptr + b * stride_qb + h * stride_qh
    qgb_ = qg_ptr + b * stride_qb + h * stride_qh
    gdb_ = gd_ptr + b * stride_gb + h * stride_gh
    delb_ = delta_ptr + b * stride_db + h * stride_dh
    ob_ = o_ptr + b * stride_ob + h * stride_oh

    offs_c = tl.arange(0, BT)
    offs_k = tl.arange(0, K)
    offs_v = tl.arange(0, V)

    S = tl.zeros((K, V), dtype=tl.float32)

    for n in range(NT):
        An_ = Ab_ + n * stride_An
        wn_ = wb_ + n * stride_wn
        un_ = ub_ + n * stride_un
        Aqn_ = Aqb_ + n * stride_qn
        qgn_ = qgb_ + n * stride_qn
        gdn_ = gdb_ + n * stride_gn
        deln_ = delb_ + n * stride_dn
        on_ = ob_ + n * stride_on

        A_final = tl.load(An_ + offs_c[:, None] * stride_Ac + offs_c[None, :] * stride_Ak)
        w = tl.load(wn_ + offs_c[:, None] * stride_wc + offs_k[None, :] * stride_wk)
        u = tl.load(un_ + offs_c[:, None] * stride_uc + offs_v[None, :] * stride_uv)
        Aqk = tl.load(Aqn_ + offs_c[:, None] * stride_Ac + offs_c[None, :] * stride_Ak)
        qg = tl.load(qgn_ + offs_c[:, None] * stride_qc + offs_k[None, :] * stride_qk)
        gdecay = tl.load(gdn_ + offs_k * stride_gk)
        delta = tl.load(deln_ + offs_c[:, None] * stride_dc + offs_k[None, :] * stride_dk)

        # ---- recurrent read ----
        v_i = u - tl.dot(w, S)                       # (BT,V)
        o_i = tl.dot(qg, S) + tl.dot(Aqk, v_i)       # (BT,V)
        tl.store(on_ + offs_c[:, None] * stride_oc + offs_v[None, :] * stride_ov, o_i)

        # ---- recurrent update of S ----
        S = S * tl.exp(gdecay)[:, None]
        S = S + tl.dot(tl.trans(delta), v_i)         # (K,V)


class Model(nn.Module):
    def __init__(self, B: int, T: int, H: int, K: int, V: int, chunk_size: int = 64):
        super().__init__()
        self.B, self.T, self.H, self.K, self.V = B, T, H, K, V
        self.chunk_size = chunk_size
        self.scale = float(K) ** -0.5
        self.register_buffer("_dummy", torch.zeros(1), persistent=False)

    def forward(self, q, k, v, g, beta):
        return _kda_chunk_fwd(q, k, v, g, beta, scale=self.scale, chunk_size=self.chunk_size)


def _kda_chunk_fwd(q, k, v, g, beta, scale, chunk_size=64):
    dtype = v.dtype
    B, T, H, K = q.shape
    V = v.shape[-1]
    BT = chunk_size
    NT = T // BT

    q = rearrange(q.to(torch.float32) * scale, "b (n c) h d -> b h n c d", c=BT)
    k = rearrange(k.to(torch.float32), "b (n c) h d -> b h n c d", c=BT)
    v = rearrange(v.to(torch.float32), "b (n c) h d -> b h n c d", c=BT)
    g = rearrange(g.to(torch.float32), "b (n c) h d -> b h n c d", c=BT)
    beta = rearrange(beta.to(torch.float32), "b (n c) h -> b h n c", c=BT)

    # ---- intra-chunk quantities (exact fp32 on host) ----
    g_cum = g.cumsum(-2)                       # (B,H,NT,BT,K)
    gexp = g_cum.exp()
    gdecay = g_cum[:, :, :, -1, :]             # (B,H,NT,K) = exp(g_last)
    L = k * torch.exp(-g_cum)                  # (B,H,NT,BT,K)
    R = k * gexp                              # (B,H,NT,BT,K)
    Akk = torch.matmul(L, R.transpose(-2, -1))                # (B,H,NT,BT,BT)
    M = -Akk * beta[:, :, :, :, None]          # beta over query axis
    mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0)
    M = M.masked_fill(mask, 0.0)              # strictly lower triangular
    # stable delta-rule inverse A = (I - M)^{-1} . beta_key
    I = torch.eye(BT, device=q.device)
    Ainv = torch.linalg.solve_triangular(I - M, I.expand_as(M), upper=False)  # (B,H,NT,BT,BT)
    A_final = Ainv * beta[:, :, :, None, :]    # beta over key axis

    w = torch.matmul(A_final, gexp * k)       # (B,H,NT,BT,K)
    u = torch.matmul(A_final, v)              # (B,H,NT,BT,V)
    Qq = q * gexp
    Aqk = torch.matmul(Qq, L.transpose(-2, -1))               # (B,H,NT,BT,BT)
    amask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=1)
    Aqk = Aqk.masked_fill(amask, 0.0)        # lower-triangular incl diag
    qg = Qq                                   # (B,H,NT,BT,K)
    delta = torch.exp(gdecay[:, :, :, None, :] - g_cum) * k              # (B,H,NT,BT,K)

    o = torch.empty(B, H, NT, BT, V, dtype=torch.float32, device=q.device)
    grid = (B * H,)
    kda_recur_kernel[grid](
        A_final, w, u, Aqk, qg, gdecay, delta, o,
        B, H, NT, BT, K, V,
        A_final.stride(0), A_final.stride(1), A_final.stride(2), A_final.stride(3), A_final.stride(4),
        w.stride(0), w.stride(1), w.stride(2), w.stride(3), w.stride(4),
        u.stride(0), u.stride(1), u.stride(2), u.stride(3), u.stride(4),
        qg.stride(0), qg.stride(1), qg.stride(2), qg.stride(3), qg.stride(4),
        gdecay.stride(0), gdecay.stride(1), gdecay.stride(2), gdecay.stride(3),
        delta.stride(0), delta.stride(1), delta.stride(2), delta.stride(3), delta.stride(4),
        o.stride(0), o.stride(1), o.stride(2), o.stride(3), o.stride(4),
        num_warps=4, num_stages=1,
    )
    torch.cuda.synchronize()
    o = rearrange(o, "b h n c d -> b (n c) h d")
    return o.to(dtype)


# Module-level shape shims (overridden by check.py / benchmark.py per shape).
B = 2
T = 1024
H = 8
K = 128
V = 128
CHUNK_SIZE = 64


def get_inputs():
    torch.manual_seed(0)
    q = torch.randn(B, T, H, K, dtype=torch.bfloat16) * 0.1
    k = torch.randn(B, T, H, K, dtype=torch.bfloat16) * 0.1
    v = torch.randn(B, T, H, V, dtype=torch.bfloat16) * 0.1
    g = (torch.randn(B, T, H, K, dtype=torch.float32) * 0.1 - 0.05)
    beta = torch.sigmoid(torch.randn(B, T, H, dtype=torch.bfloat16))
    return [q, k, v, g, beta]


def get_init_inputs():
    return [B, T, H, K, V, CHUNK_SIZE]