"""Kimi Delta Attention (KDA) forward, chunk form — batched PyTorch.

Reproduces fla/ops/kda chunk-parallel KDA forward without calling any FLA op.
Mirrors reference.py exactly; only the two structurally-sequential loops remain
(the intra-chunk triangular solve over BT=64 and the inter-chunk KxV state
recurrence over NT chunks), both batched over (B, H) so the heavy GEMMs hit
cuBLAS tensor cores on SM90.

Interface (per reference.py): Model, get_inputs, get_init_inputs.
"""
from __future__ import annotations

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

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


def _kda_fwd(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
    scale: float,
    chunk_size: int = 64,
) -> torch.Tensor:
    dtype = v.dtype
    B, T, H, K = q.shape
    V = v.shape[-1]
    BT = chunk_size
    NT = T // BT
    q, k, v, g, beta = (x.to(torch.float32) for x in (q, k, v, g, beta))
    q = q * scale

    q = rearrange(q, "b (n c) h d -> b h n c d", c=BT)
    k = rearrange(k, "b (n c) h d -> b h n c d", c=BT)
    v = rearrange(v, "b (n c) h d -> b h n c d", c=BT)
    g = rearrange(g, "b (n c) h d -> b h n c d", c=BT)
    beta = rearrange(beta, "b (n c) h -> b h n c", c=BT)

    g = g.cumsum(-2)

    # Gated keys. A[c, i] = <kg_c, kgi_i>; Aqk[c, j] = <qg_c, kgi_j>.
    kg = k * g.exp()
    kgi = k * (-g).exp()
    qg = q * g.exp()

    # ---- Intra-chunk A: strict-lower K-K interaction (incl. diag masked) ----
    A = torch.bmm(kg.reshape(-1, BT, K), kgi.reshape(-1, BT, K).transpose(-1, -2))
    A = A.view(B, H, NT, BT, BT)
    A = A * beta[:, :, :, :, None]
    mask = torch.triu(torch.ones(BT, BT, device=q.device, dtype=torch.bool), diagonal=0)
    A = -A.masked_fill(mask, 0)
    for i in range(1, BT):
        A[..., i, :i] += (A[..., i, :, None] * A[..., :, :i]).sum(-2)
    A = (A + torch.eye(BT, device=q.device)) * beta[:, :, :, None, :]

    # w = A @ (exp(g)*k), u = A @ v
    ek = g.exp() * k
    w = torch.bmm(A.reshape(-1, BT, BT), ek.reshape(-1, BT, K)).view(B, H, NT, BT, K)
    u = torch.bmm(A.reshape(-1, BT, BT), v.reshape(-1, BT, V)).view(B, H, NT, BT, V)

    # Aqk: inter-chunk q-k interaction, strict-upper masked.
    Aqk = torch.bmm(qg.reshape(-1, BT, K), kgi.reshape(-1, BT, K).transpose(-1, -2))
    Aqk = Aqk.view(B, H, NT, BT, BT)
    mask_strict = torch.triu(torch.ones(BT, BT, device=q.device, dtype=torch.bool), diagonal=1)
    Aqk = Aqk.masked_fill(mask_strict, 0)

    # ---- Inter-chunk recurrence over the KxV state S ----
    S = q.new_zeros(B, H, K, V)
    o = torch.zeros_like(v)
    for i in range(NT):
        q_i, k_i, u_i, g_i, w_i = q[:, :, i], k[:, :, i], u[:, :, i], g[:, :, i], w[:, :, i]
        aqk_i = Aqk[:, :, i]
        v_i = u_i - w_i @ S
        o[:, :, i] = qg[:, :, i] @ S + aqk_i @ v_i
        S = S * g_i[:, :, -1].exp()[..., None]
        S = S + ((g_i[:, :, -1:] - g_i).exp() * k_i).transpose(-1, -2) @ v_i

    o = rearrange(o, "b h n c d -> b (n c) h d")
    return o.to(dtype)


class Model(nn.Module):
    """KDA forward (chunk form). No learned parameters; all inputs are activations."""

    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: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        g: torch.Tensor,
        beta: torch.Tensor,
    ) -> torch.Tensor:
        return _kda_fwd(q, k, v, g, beta, scale=self.scale, chunk_size=self.chunk_size)


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]