"""Kimi Delta Attention (KDA) forward, chunk form -- custom vectorized kernel.

Faithful GPU-resident port of the reference chunk-parallel formulation
(fla naive.py, Songlin Yang et al.).  The computation is *identical* to the
reference (same operations, same order, all math in fp32, same accumulation
order), with two changes for speed:

  * building the intra-chunk A matrix and the per-chunk Aqk matrix gather their
    (row, col) outer product across K in a single batched reduction instead of
    64 Python einsum calls, and
  * the NT-chunk recurrent pass runs as batched matmuls over (B, H) so each
    chunk step is one kernel launch instead of per-head.

Numerics therefore match the oracle to well within the 0.05 atol/rtol bar,
and nothing calls fla's existing KDA kernel (problem.yaml constraint).

  q, k : (B, T, H, K)   bf16
  v    : (B, T, H, V)   bf16
  g    : (B, T, H, K)   fp32   (per-channel log-decay; in-chunk cumsum applied)
  beta : (B, T, H)      bf16

  o    : (B, T, H, V)   bf16
"""
from __future__ import annotations

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


def _kda_forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
                 g: torch.Tensor, beta: torch.Tensor, scale: float,
                 chunk_size: int = 64) -> torch.Tensor:
    """KDA forward.  Identical math to reference.py, run efficiently on GPU."""
    dtype = v.dtype
    B_, T_, H_, K_ = q.shape
    V_ = v.shape[-1]
    BT = chunk_size
    assert T_ % BT == 0
    NT = T_ // BT

    # The oracle casts everything to fp32 and scales only q -- replicate exactly.
    q = q.float() * scale
    k = k.float()
    v = v.float()
    g = g.float()
    beta = beta.float()

    # (B, T, H, K/V) -> (B, H, NT, BT, K/V).
    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)

    # In-chunk cumulative sum of the log-decay (the reference's only g tweak).
    g = g.cumsum(-2)

    # Static triangular masks on the chunk time axis (BT = 64).
    triu_incl = torch.triu(
        torch.ones(BT, BT, dtype=torch.bool, device=q.device), 0)   # row <= col
    triu_strict = torch.triu(
        torch.ones(BT, BT, dtype=torch.bool, device=q.device), 1)   # row < col

    # ------------------------------------------------------------------------
    # Intra-chunk transform A, shape (B, H, NT, row, col).  Reference builds
    #   A[row, col] = sum_d k[row,d] * exp(g[row,d]-g[col,d]) * k[col,d]
    # over (row > col) via a per-column loop; we fuse the (row, col) gather
    # across K into one reduction that yields identical values.
    g_diff = (g.unsqueeze(-2) - g.unsqueeze(-3)).exp()          # (B,H,NT,R,C,K)
    A = (k.unsqueeze(-2) * g_diff * k.unsqueeze(-3)).sum(-1)    # (B,H,NT,R,C)

    # Scale rows by beta: reference ``A * beta[..., None]``.
    A = A * beta.unsqueeze(-1)                # (B,H,NT,R,1) broadcasts over cols

    # Zero on & above the diagonal (incl. diagonal), then negate so the
    # strictly-lower terms are positive: ``-A.masked_fill(triu_incl, 0)``.
    A = -A.masked_fill(triu_incl, 0)

    # Reference inversion-style cumulative sum:
    #   for a in 1..BT-1: A[a,:a] += (A[a,:,None]*A[:,:a]).sum(-2)
    # Reproduce the loop (64 iters, each a cheap batched op) so the
    # accumulation order -- hence fp rounding -- matches the oracle.
    for a in range(1, BT):
        A[..., a, :a] = (A[..., a, :a]
                         + (A[..., a, :, None] * A[..., :, :a]).sum(-2))

    # Add identity and scale COLUMNS by beta:
    # reference ``(A + eye) * beta[..., None, :]`` -- beta multiplies cols.
    A = (A + torch.eye(BT, dtype=A.dtype, device=A.device)) * beta.unsqueeze(-2)

    # Precompute exp(g) and the chunk-wide w = A @ (exp(g)*k),  u = A @ v.
    gexp = g.exp()
    w = A @ (gexp * k)   # (B,H,NT,BT,K)
    u = A @ v             # (B,H,NT,BT,V)

    # ------------------------------------------------------------------------
    # Recurrent inter-chunk pass.  State S (B,H,K,V) is threaded through the
    # NT chunks; per-chunk quantities carry a leading (B,H,BT) so every op in
    # the loop is a single batched matmul over (B, H).
    S = q.new_zeros(B_, H_, K_, V_)       # fp32 state
    out = torch.zeros_like(v)             # (B, H, NT, BT, V)
    qgexp = q * gexp                      # reused term (q*exp(g))
    for i in range(NT):
        q_i = q[:, :, i]      # (B, H, BT, K)
        k_i = k[:, :, i]      # (B, H, BT, K)
        u_i = u[:, :, i]      # (B, H, BT, V)
        w_i = w[:, :, i]      # (B, H, BT, K)
        g_i = g[:, :, i]      # (B, H, BT, K)
        gexp_i = gexp[:, :, i]
        qgexp_i = qgexp[:, :, i]

        # Aqk[a, j] = sum_d q[a,d]*exp(g[a,d]-g[j,d])*k[j,d] for a >= j.
        # Reference zeros strictly-above-diagonal with triu(diagonal=1).
        gd = (g_i.unsqueeze(-2) - g_i.unsqueeze(-3)).exp()          # (B,H,R,C,K)
        Aqk = (q_i.unsqueeze(-2) * gd * k_i.unsqueeze(-3)).sum(-1)   # (B,H,R,C)
        Aqk = Aqk.masked_fill(triu_strict, 0)

        v_i = u_i - w_i @ S                       # (B, H, BT, V)
        out[:, :, i] = qgexp_i @ S + Aqk @ v_i   # (B, H, BT, V)

        # Decay state by exp(g at the chunk's last row), then add incoming.
        S = S * rearrange(g_i[:, :, -1].exp(), "b h k -> b h k 1")
        S = S + rearrange(
            (g_i[:, :, -1:] - g_i).exp() * k_i, "b h c k -> b h k c"
        ) @ v_i

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


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


class Model(nn.Module):
    """KDA forward (chunk form). No learned parameters; all inputs 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
        # No learned params; dummy buffer keeps state_dict well-defined & matches
        # reference (strict=True in check.py requires the same keys).
        self.register_buffer("_dummy", torch.zeros(1), persistent=False)

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


def get_inputs():
    """Return a list of activations for one forward call (bf16 q/k/v/beta, fp32 g)."""
    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
    # log-decay: small negative numbers so exp(g) sits in (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]