"""KDA forward (chunk form) — custom Triton kernel implementation.

Reimplements the FLA chunk-parallel KDA forward without calling the FLA library.

The op:
  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 here)
  beta: (B, T, H)    bf16
  output: o (B, T, H, V) bf16

Algorithm overview:
  For each (B, H), split T into chunks of BT=64.
  Per-chunk:
    1. cumsum g within chunk
    2. build A_qk [BT, BT] (lower-tri w/ diag) using q, k, g
    3. build A_kk [BT, BT] via the WY-style delta-rule factorization
    4. precompute w = A @ (exp(g) * k), u = A @ v
  Inter-chunk recurrent pass with state S [K, V]:
    For each chunk in order:
      v_i = u_i - w_i @ S
      o_i = (q_i * exp(g_i)) @ S + A_qk @ v_i
      S = S * exp(g_i[-1]) + (k_i * exp(g_i[-1] - g_i)).T @ v_i
"""
from __future__ import annotations

import torch
import torch.nn as nn
import triton
import triton.language as tl

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


# ============================================================
# Kernel 1: in-chunk cumsum of g
# g_raw: (B, T, H, K) fp32, g_cum: same shape
# Per program: one chunk, one (B, H) -> output (BT, K) block.
# ============================================================
@triton.jit
def _g_cumsum_kernel(
    g_ptr, gc_ptr,
    H, T,
    K: tl.constexpr, BT: tl.constexpr,
):
    pid_t = tl.program_id(0)
    pid_bh = tl.program_id(1)
    pid_b = pid_bh // H
    pid_h = pid_bh % H

    base = pid_b * T * H * K + pid_h * K
    offs_t = pid_t * BT + tl.arange(0, BT)
    offs_k = tl.arange(0, K)
    m_t = (pid_t * BT + offs_t) < T

    ptrs = g_ptr + base + offs_t[:, None] * H * K + offs_k[None, :]
    g = tl.load(ptrs, mask=m_t[:, None], other=0.0)
    gc = tl.cumsum(g, axis=0)
    out_ptrs = gc_ptr + base + offs_t[:, None] * H * K + offs_k[None, :]
    tl.store(out_ptrs, gc, mask=m_t[:, None])


# ============================================================
# Kernel 2: intra-chunk (cumsum, A_qk, A_kk, w, u)
# Per program: one chunk, one (B, H).
# Strategy: load q, k, v, g, beta once. Cumsum g. Compute A_kk_basic
# (using matmuls over K), do the WY completion via doubling, then add I
# and column-scale by beta. Compute A_qk. Then compute u and w.
# ============================================================
@triton.jit
def _intra_kernel(
    q_ptr, k_ptr, v_ptr, g_ptr, beta_ptr,
    Aqk_ptr, w_ptr, u_ptr,
    H, T, V,
    scale,
    K: tl.constexpr, BT: tl.constexpr, BV: tl.constexpr,
):
    pid_t = tl.program_id(0)
    pid_bh = tl.program_id(1)
    pid_b = pid_bh // H
    pid_h = pid_bh % H

    base_qk = pid_b * T * H * K + pid_h * K
    base_v = pid_b * T * H * V + pid_h * V
    base_g = pid_b * T * H * K + pid_h * K
    base_beta = pid_b * T * H + pid_h
    base_A = pid_b * T * H * BT + pid_h * BT
    base_w = pid_b * T * H * K + pid_h * K
    base_u = pid_b * T * H * V + pid_h * V

    offs_t = pid_t * BT + tl.arange(0, BT)
    m_t = offs_t < T
    o_i = tl.arange(0, BT)  # within-chunk row/col index

    # Load g for the chunk and cumsum
    g_ptrs = g_ptr + base_g + offs_t[:, None] * H * K + tl.arange(0, K)[None, :]
    g = tl.load(g_ptrs, mask=m_t[:, None], other=0.0)  # (BT, K) fp32
    g = tl.cumsum(g, axis=0)  # (BT, K) fp32

    # Load beta for the chunk
    b = tl.load(beta_ptr + base_beta + offs_t * H, mask=m_t, other=0.0).to(tl.float32)  # (BT,)

    # ===== Build A_kk_strict_lower (BT, BT) =====
    # A_kk_sl[i, j] = beta[i] * (k[j] * exp(g[j] - g[i])).dot(k[i])  for j < i
    # We compute via a single matmul over K.
    #   kg_neg = k * exp(-g)  (BT, K)
    #   kg_pos = k * exp(g)   (BT, K)
    #   A_kk_basic = kg_neg @ kg_pos.T  (BT, BT)
    # Since K is constexpr (128), the matmul is over the full K dimension.

    b_Akk = tl.zeros([BT, BT], dtype=tl.float32)
    K_offs = tl.arange(0, K)
    k_tile = tl.load(k_ptr + base_qk + offs_t[:, None] * H * K + K_offs[None, :],
                     mask=m_t[:, None], other=0.0).to(tl.float32)
    g_tile = tl.load(g_ptr + base_g + offs_t[:, None] * H * K + K_offs[None, :],
                     mask=m_t[:, None], other=0.0)
    kg_neg = k_tile * tl.exp(-g_tile)
    kg_pos = k_tile * tl.exp(g_tile)
    b_Akk = tl.dot(kg_neg.to(tl.bfloat16), tl.trans(kg_pos).to(tl.bfloat16),
                   out_dtype=tl.float32)

    # Apply row-wise beta scaling
    b_Akk = b_Akk * b[:, None]
    # Mask to strict-lower triangular and negate
    mask_sl = o_i[:, None] > o_i[None, :]
    mask_sl = mask_sl * m_t[:, None] * m_t[None, :]
    b_Akk = tl.where(mask_sl, -b_Akk, 0.0)

    # ===== WY completion via doubling =====
    # We want T = A + A^2 + ... + A^(BT-1) for the strict-lower part.
    # Doubling: T_{2m} = T_m + A^{2m} @ (I + T_m) = T_m + A^{2m} + A^{2m} @ T_m
    # where A^{2m} = (A^m) @ (A^m).
    # We precompute powers A^2, A^4, A^8, ..., A^32 (5 squarings).
    P1 = tl.dot(b_Akk, b_Akk, out_dtype=tl.float32, allow_tf32=False)
    P1 = tl.where(mask_sl, P1, 0.0)
    P2 = tl.dot(P1, P1, out_dtype=tl.float32, allow_tf32=False)
    P2 = tl.where(mask_sl, P2, 0.0)
    P3 = tl.dot(P2, P2, out_dtype=tl.float32, allow_tf32=False)
    P3 = tl.where(mask_sl, P3, 0.0)
    P4 = tl.dot(P3, P3, out_dtype=tl.float32, allow_tf32=False)
    P4 = tl.where(mask_sl, P4, 0.0)
    P5 = tl.dot(P4, P4, out_dtype=tl.float32, allow_tf32=False)
    P5 = tl.where(mask_sl, P5, 0.0)

    # T = b_Akk (which is A)
    T_acc = b_Akk
    # T = T + A^2 + A^2 @ T
    T_acc = T_acc + P1 + tl.dot(P1, T_acc, out_dtype=tl.float32, allow_tf32=False)
    T_acc = tl.where(mask_sl, T_acc, 0.0)
    # T = T + A^4 + A^4 @ T
    T_acc = T_acc + P2 + tl.dot(P2, T_acc, out_dtype=tl.float32, allow_tf32=False)
    T_acc = tl.where(mask_sl, T_acc, 0.0)
    # T = T + A^8 + A^8 @ T
    T_acc = T_acc + P3 + tl.dot(P3, T_acc, out_dtype=tl.float32, allow_tf32=False)
    T_acc = tl.where(mask_sl, T_acc, 0.0)
    # T = T + A^16 + A^16 @ T
    T_acc = T_acc + P4 + tl.dot(P4, T_acc, out_dtype=tl.float32, allow_tf32=False)
    T_acc = tl.where(mask_sl, T_acc, 0.0)
    # T = T + A^32 + A^32 @ T
    T_acc = T_acc + P5 + tl.dot(P5, T_acc, out_dtype=tl.float32, allow_tf32=False)
    T_acc = tl.where(mask_sl, T_acc, 0.0)
    # A^64 = 0, so T = T (final T is the geometric sum)

    b_Akk = T_acc  # T = A + A^2 + ... + A^(BT-1)

    # Add identity: A + I (so diagonal is now 1)
    I = (o_i[:, None] == o_i[None, :]).to(tl.float32)
    I = I * m_t[:, None].to(tl.float32) * m_t[None, :].to(tl.float32)
    b_Akk = b_Akk + I

    # Column-scale by beta
    b_Akk = b_Akk * b[None, :]

    # ===== Compute A_qk =====
    # A_qk[c, j] = (q[c] * exp(g[c] - g[j])).dot(k[j])  for c >= j
    # = (q * exp(g))[c, :].dot((k * exp(-g))[j, :])
    # So A_qk = (q * exp(g)) @ (k * exp(-g)).T
    b_Aqk = tl.zeros([BT, BT], dtype=tl.float32)
    q_tile = tl.load(q_ptr + base_qk + offs_t[:, None] * H * K + K_offs[None, :],
                     mask=m_t[:, None], other=0.0).to(tl.float32) * scale
    qg = q_tile * tl.exp(g_tile)
    kg = k_tile * tl.exp(-g_tile)
    b_Aqk = tl.dot(qg.to(tl.bfloat16), tl.trans(kg).to(tl.bfloat16),
                   out_dtype=tl.float32)

    # Mask to lower-tri w/ diag
    mask_qk = o_i[:, None] >= o_i[None, :]
    mask_qk = mask_qk * m_t[:, None] * m_t[None, :]
    b_Aqk = tl.where(mask_qk, b_Aqk, 0.0)

    # Store A_qk as (B, T, H, BT) in fp32
    # (Each chunk's A_qk is a (BT, BT) block stored at T[pid_t*BT : pid_t*BT+BT, h, :])
    Aqk_offs = offs_t[:, None] * H * BT + o_i[None, :]
    tl.store(Aqk_ptr + base_A + Aqk_offs, b_Aqk, mask=m_t[:, None])

    # ===== Compute u = A_kk @ v (over V dim) =====
    for vv in range(0, V, BV):
        vv_offs = vv + tl.arange(0, BV)
        m_vv = vv_offs < V
        v = tl.load(v_ptr + base_v + offs_t[:, None] * H * V + vv_offs[None, :],
                    mask=m_t[:, None] & m_vv[None, :], other=0.0).to(tl.float32)
        b_u = tl.dot(b_Akk.to(tl.bfloat16), v.to(tl.bfloat16), out_dtype=tl.float32)
        tl.store(u_ptr + base_u + offs_t[:, None] * H * V + vv_offs[None, :],
                 b_u.to(tl.bfloat16), mask=m_t[:, None] & m_vv[None, :])

    # ===== Compute w = A_kk @ (k * exp(g)) =====
    # k and g are already loaded above
    kg = k_tile * tl.exp(g_tile)
    b_w = tl.dot(b_Akk.to(tl.bfloat16), kg.to(tl.bfloat16), out_dtype=tl.float32)
    tl.store(w_ptr + base_w + offs_t[:, None] * H * K + K_offs[None, :],
             b_w.to(tl.bfloat16), mask=m_t[:, None])


# ============================================================
# Kernel 3: state pass (per (B, H), serially through chunks)
# For each chunk:
#   v_new = u - w @ S_prev
#   Store S_prev to h[i_t]
#   S_new = S_prev * exp(g[-1]) + (k * exp(g[-1] - g)).T @ v_new
#   Store v_new
# ============================================================
@triton.jit
def _state_pass_kernel(
    u_ptr, w_ptr, k_ptr, gc_ptr,
    v_new_ptr, h_ptr,
    H, T, NT,
    K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr,
):
    pid_bh = tl.program_id(0)
    pid_b = pid_bh // H
    pid_h = pid_bh % H

    base_v = pid_b * T * H * V + pid_h * V
    base_qk = pid_b * T * H * K + pid_h * K
    base_g = pid_b * T * H * K + pid_h * K
    base_h = pid_b * NT * H * K * V + pid_h * K * V

    # S in registers (K, V) as fp32
    S = tl.zeros([K, V], dtype=tl.float32)
    K_offs = tl.arange(0, K)
    V_offs = tl.arange(0, V)

    for i_t in range(NT):
        offs_t = i_t * BT + tl.arange(0, BT)
        m_t = offs_t < T

        # Store S_prev to h[i_t] (BEFORE update)
        # h[i_t] has shape (B, NT, H, K, V), stored as fp32
        tl.store(h_ptr + base_h + i_t * H * K * V + K_offs[:, None] * V + V_offs[None, :],
                 S)

        # Load w [BT, K], u [BT, V], k [BT, K], g [BT, K]
        w = tl.load(w_ptr + base_qk + offs_t[:, None] * H * K + K_offs[None, :],
                    mask=m_t[:, None], other=0.0).to(tl.float32)
        u = tl.load(u_ptr + base_v + offs_t[:, None] * H * V + V_offs[None, :],
                    mask=m_t[:, None], other=0.0).to(tl.float32)
        k = tl.load(k_ptr + base_qk + offs_t[:, None] * H * K + K_offs[None, :],
                    mask=m_t[:, None], other=0.0).to(tl.float32)
        g = tl.load(gc_ptr + base_g + offs_t[:, None] * H * K + K_offs[None, :],
                    mask=m_t[:, None], other=0.0)

        # v_new = u - w @ S
        v_new = u - tl.dot(w.to(tl.bfloat16), S.to(tl.bfloat16), out_dtype=tl.float32)
        # Store v_new
        tl.store(v_new_ptr + base_v + offs_t[:, None] * H * V + V_offs[None, :],
                 v_new.to(tl.bfloat16), mask=m_t[:, None])

        # S_new = S * exp(g[-1]) + (k * exp(g[-1] - g)).T @ v_new
        # Load last g of the chunk
        last_idx = i_t * BT + BT - 1
        g_last = tl.load(gc_ptr + base_g + last_idx * H * K + K_offs).to(tl.float32)  # (K,)
        decay = tl.exp(g_last)  # (K,)
        S = S * decay[None, :]
        # kg = k * exp(g[-1] - g[c])  (BT, K)
        kg = k * tl.exp(g_last[None, :] - g)  # (BT, K)
        # S += kg.T @ v_new  (K, BT) @ (BT, V) = (K, V)
        S = S + tl.dot(tl.trans(kg).to(tl.bfloat16), v_new.to(tl.bfloat16), out_dtype=tl.float32)


# ============================================================
# Kernel 4: output
# Per program: one (B, H, chunk) and one V tile.
# Computes o = (q * exp(g)) @ S_prev * scale + A_qk @ v_new
# ============================================================
@triton.jit
def _output_kernel(
    q_ptr, gc_ptr, Aqk_ptr, v_new_ptr, h_ptr, o_ptr,
    H, T, NT,
    scale,
    K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BV: tl.constexpr,
):
    pid_v = tl.program_id(0)
    pid_t = tl.program_id(1)
    pid_bh = tl.program_id(2)
    pid_b = pid_bh // H
    pid_h = pid_bh % H

    base_qk = pid_b * T * H * K + pid_h * K
    base_g = pid_b * T * H * K + pid_h * K
    base_v = pid_b * T * H * V + pid_h * V
    base_A = pid_b * T * H * BT + pid_h * BT
    base_h = pid_b * NT * H * K * V + pid_h * K * V
    base_o = pid_b * T * H * V + pid_h * V

    offs_t = pid_t * BT + tl.arange(0, BT)
    m_t = offs_t < T
    o_i = tl.arange(0, BT)
    v_offs = pid_v * BV + tl.arange(0, BV)
    m_v = v_offs < V
    K_offs = tl.arange(0, K)

    # Initialize accumulator
    b_o = tl.zeros([BT, BV], dtype=tl.float32)
    # Load q and g
    q = tl.load(q_ptr + base_qk + offs_t[:, None] * H * K + K_offs[None, :],
                mask=m_t[:, None], other=0.0).to(tl.float32) * scale
    g = tl.load(gc_ptr + base_g + offs_t[:, None] * H * K + K_offs[None, :],
                mask=m_t[:, None], other=0.0)
    qg = q * tl.exp(g)
    # Load S [K, BV] tile
    S_tile = tl.load(h_ptr + base_h + pid_t * H * K * V + K_offs[:, None] * V + v_offs[None, :],
                     mask=m_v[None, :], other=0.0)
    b_o = tl.dot(qg.to(tl.bfloat16), S_tile.to(tl.bfloat16), out_dtype=tl.float32)

    # Add A_qk @ v_new
    # A_qk is stored as (B, T, H, BT) fp32, where the (BT, BT) block for chunk pid_t
    # is at T[pid_t*BT : pid_t*BT+BT, h, :]
    Aqk = tl.load(Aqk_ptr + base_A + offs_t[:, None] * H * BT + o_i[None, :],
                  mask=m_t[:, None], other=0.0)
    v_new = tl.load(v_new_ptr + base_v + offs_t[:, None] * H * V + v_offs[None, :],
                    mask=m_t[:, None] & m_v[None, :], other=0.0).to(tl.float32)
    b_o = b_o + tl.dot(Aqk.to(tl.bfloat16), v_new.to(tl.bfloat16), out_dtype=tl.float32)

    # Store o
    tl.store(o_ptr + base_o + offs_t[:, None] * H * V + v_offs[None, :],
             b_o.to(tl.bfloat16), mask=m_t[:, None] & m_v[None, :])


# ============================================================
# Model
# ============================================================
class Model(nn.Module):
    """KDA forward (chunk form). Custom Triton kernel."""

    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:
        B, T, H, K, V = self.B, self.T, self.H, self.K, self.V
        BT = self.chunk_size
        NT = T // BT
        device = q.device

        o = torch.empty_like(v)
        g_cum = torch.empty_like(g)
        Aqk = torch.empty(B, T, H, BT, device=device, dtype=torch.float32)
        w = torch.empty_like(k)
        u = torch.empty_like(v)
        v_new = torch.empty_like(v)
        h = torch.empty(B, NT, H, K, V, device=device, dtype=torch.float32)

        BK = 128
        BV = 64

        # Kernel 1: in-chunk cumsum
        _g_cumsum_kernel[(NT, B * H)](
            g, g_cum,
            H, T,
            K=K, BT=BT,
            num_warps=4,
        )

        # Kernel 2: intra-chunk (A_qk, A_kk, w, u)
        _intra_kernel[(NT, B * H)](
            q, k, v, g_cum, beta,
            Aqk, w, u,
            H, T, V,
            self.scale,
            K=K, BT=BT, BV=BV,
            num_warps=4,
        )

        # Kernel 3: state pass
        _state_pass_kernel[(B * H,)](
            u, w, k, g_cum,
            v_new, h,
            H, T, NT,
            K=K, V=V, BT=BT,
            num_warps=4,
        )

        # Kernel 4: output
        _output_kernel[(triton.cdiv(V, BV), NT, B * H)](
            q, g_cum, Aqk, v_new, h, o,
            H, T, NT,
            self.scale,
            K=K, V=V, BT=BT, BV=BV,
            num_warps=4,
        )

        return o


# ============================================================
# Module-level shape shims
# ============================================================
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, K, 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]