"""Custom Triton kernel for Kimi Delta Attention (KDA) chunk forward.

Implements the chunk-parallel KDA forward from fla (Songlin Yang et al.) using
hand-written Triton kernels for the H100 PCIe (SM90 Hopper, HBM2e, ~2.0 TB/s).

No FLA / fla.ops imports -- this is a from-scratch implementation.

Algorithm (H == HV in this problem, so no GVA expansion is needed):
  1. In-chunk cumsum of g (per channel) -- torch op.
  2. For each chunk c, compute intra-chunk (in fp32, with numerical centering
     around the middle of the chunk to keep exp values bounded):
       Aqk[c] = lower-tri(q[i] * exp(g[i] - g[j]), k[j]) * scale      [BT, BT]
       Akk[c] = lower-tri(k[i] * exp(g[i] - g[j]), k[j]) * beta[i]   [BT, BT]
     then solve Akk_inv by forward substitution (Woodbury-style; in fp32
     for stability), and compute:
       w[c] = Akk_inv @ (exp(g) * k)   [BT, K]
       u[c] = Akk_inv @ v              [BT, V]
  3. Inter-chunk recurrence (one program per (B, H, V-tile)):
       v_new[c] = u[c] - w[c] @ h_prev
       h_{c+1}  = h_c * exp(g[last]) + (k[c] * (exp(g[last]) - exp(g[c]))) @ v_new[c]
     Stores h in (B, NT, H, K, V) and v_new in (B, T, H, V).
  4. Output:
       o[c] = (q * exp(g)) @ h_{c} + Aqk @ v_new[c]
"""
from __future__ import annotations

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

# ---------------------------------------------------------------------------
# Step 1: in-chunk cumsum of g (PyTorch -- small op)
# ---------------------------------------------------------------------------


def _chunk_local_cumsum_g(g: torch.Tensor, chunk_size: int) -> torch.Tensor:
    """g: (B, T, H, K) fp32. Returns in-chunk cumsum (per chunk, per channel)."""
    B, T, H, K = g.shape
    NT = T // chunk_size
    g4 = g.view(B, NT, chunk_size, H, K)
    return g4.cumsum(dim=2).view(B, T, H, K)


# ---------------------------------------------------------------------------
# Step 2: intra-chunk compute -- Aqk [BT, BT] + Akk_inv [BT, BT] (fp32)
#         + w [BT, K] + u [BT, V]
# ---------------------------------------------------------------------------


@triton.jit
def _kda_intra_wu_kernel(
    q_ptr, k_ptr, v_ptr, g_ptr, beta_ptr,
    A_ptr, w_ptr, u_ptr,
    scale,
    NT, T,
    H: tl.constexpr, HV: tl.constexpr,
    K: tl.constexpr, V: tl.constexpr,
    BT: tl.constexpr,
    CENTER: tl.constexpr,
):
    """One program per (NT, B*H) tile. Computes Aqk, Akk, solves Akk_inv, w, u.

    Layout: tensors are stored as (B, T, H, ...) row-major contiguous.
    """
    i_t = tl.program_id(0)
    i_bh = tl.program_id(1)
    i_b = i_bh // HV
    i_hv = i_bh % HV

    bos = i_b * T

    q = q_ptr + (bos * H + i_hv) * K
    k = k_ptr + (bos * H + i_hv) * K
    v = v_ptr + (bos * HV + i_hv) * V
    g = g_ptr + (bos * HV + i_hv) * K
    beta_p = beta_ptr + bos * HV + i_hv
    A = A_ptr + (bos * HV + i_hv) * BT
    w_o = w_ptr + (bos * HV + i_hv) * K
    u_o = u_ptr + (bos * HV + i_hv) * V

    o_t = i_t * BT + tl.arange(0, BT)
    o_k = tl.arange(0, K)
    o_v = tl.arange(0, V)
    o_bt = tl.arange(0, BT)
    m_t = o_t < T

    # ---- Load g for whole chunk (in fp32) and pick a "center" row to subtract
    g_offs = o_t[:, None] * (HV * K) + o_k[None, :]
    b_g = tl.load(g + g_offs, mask=m_t[:, None], other=0.0).to(tl.float32)  # [BT, K]
    # Center on row CENTER of the chunk: a single vector g_c [K]
    b_gc = tl.load(g + (i_t * BT + CENTER) * (HV * K) + o_k).to(tl.float32)  # [K]
    b_gm = b_g - b_gc[None, :]  # [BT, K], centered

    # ---- Load q, k, v ----
    q_offs = o_t[:, None] * (H * K) + o_k[None, :]
    b_q = tl.load(q + q_offs, mask=m_t[:, None], other=0.0).to(tl.float32)  # [BT, K]

    k_offs = o_t[:, None] * (H * K) + o_k[None, :]
    b_k = tl.load(k + k_offs, mask=m_t[:, None], other=0.0).to(tl.float32)  # [BT, K]

    v_offs = o_t[:, None] * (HV * V) + o_v[None, :]
    b_v = tl.load(v + v_offs, mask=m_t[:, None], other=0.0).to(tl.float32)  # [BT, V]

    b_beta = tl.load(beta_p + o_t * HV, mask=m_t, other=0.0).to(tl.float32)  # [BT]

    # ---- Build Aqk [BT, BT] and Akk [BT, BT] ----
    # For each pair (i, j), we need exp(g[i] - g[j]). We center on g_c (same for all),
    # so exp(g[i] - g[j]) = exp((g[i] - g_c) - (g[j] - g_c)) = exp(gm[i]) * exp(-gm[j])
    # where gm = b_gm. We do the matmul directly:
    #   A[i, j] = <q[i, :], k[j, :] * exp(g[i, :] - g[j, :])> * beta[i or 1]
    # = <q[i, :] * exp(gm[i, :]), k[j, :] * exp(-gm[j, :])> * beta[i or 1]
    exp_gi = tl.exp(b_gm)              # [BT, K]
    exp_neg_gj = tl.exp(-b_gm)         # [BT, K]
    b_qg = b_q * exp_gi                # [BT, K]
    b_kg = b_k * exp_neg_gj            # [BT, K]
    b_kb_g = (b_k * b_beta[:, None]) * exp_neg_gj  # [BT, K]

    # Single matmul for Aqk and Akk
    b_Aqk = tl.dot(b_qg.to(tl.bfloat16), tl.trans(b_kg).to(tl.bfloat16), out_dtype=tl.float32)
    b_Akk = tl.dot(b_kb_g.to(tl.bfloat16), tl.trans(b_kg).to(tl.bfloat16), out_dtype=tl.float32)

    # Mask strictly lower triangular (set diag and above to 0)
    m_strict = o_bt[:, None] > o_bt[None, :]
    b_Aqk = tl.where(m_strict, b_Aqk, 0.0)
    b_Akk = tl.where(m_strict, b_Akk, 0.0)

    # Forward substitution in fp32 to compute (I + L + L^2 + ... + L^63) where L is
    # the lower-strict matrix above. b_Akk already includes the beta[i] factor.
    m_diag = o_bt[:, None] == o_bt[None, :]
    b_inv = b_Akk
    cur = b_Akk
    for _ in range(1, BT):
        cur = tl.dot(cur, b_Akk, out_dtype=tl.float32, input_precision='ieee')
        b_inv = b_inv + cur
    b_inv = b_inv + m_diag.to(tl.float32)  # add I
    # Multiply by beta[j] (column factor)
    b_inv = b_inv * b_beta[None, :]

    # ---- Compute w = b_inv @ (exp(g) * k) and u = b_inv @ v ----
    b_kg2 = b_k * tl.exp(b_g)  # [BT, K]
    b_w = tl.dot(b_inv.to(tl.bfloat16), b_kg2.to(tl.bfloat16), out_dtype=tl.float32)
    b_u = tl.dot(b_inv.to(tl.bfloat16), b_v.to(tl.bfloat16), out_dtype=tl.float32)

    # ---- Store outputs ----
    # w [BT, K]
    w_offs = o_t[:, None] * (HV * K) + o_k[None, :]
    tl.store(w_o + w_offs, b_w.to(tl.bfloat16), mask=m_t[:, None])

    # u [BT, V]
    u_offs = o_t[:, None] * (HV * V) + o_v[None, :]
    tl.store(u_o + u_offs, b_u.to(tl.bfloat16), mask=m_t[:, None])

    # Aqk [BT, BT] (per chunk) - store the lower-strict part scaled.
    A_offs = o_t[:, None] * (HV * BT) + o_bt[None, :]
    A_mask = m_t[:, None]
    b_Aqk_scaled = b_Aqk * scale
    tl.store(A + A_offs, b_Aqk_scaled.to(tl.bfloat16), mask=A_mask)


# ---------------------------------------------------------------------------
# Step 3: inter-chunk recurrence (one program per (B, H, V-tile))
# ---------------------------------------------------------------------------


@triton.jit
def _kda_inter_h_kernel(
    g_ptr, k_ptr, w_ptr, u_ptr, v_new_ptr, h_ptr,
    NT, T,
    H: tl.constexpr, HV: tl.constexpr,
    K: tl.constexpr, V: tl.constexpr,
    BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr,
):
    """Inter-chunk recurrent state update, stores h in a (B, NT, HV, K, V) buffer.

    Grid: (cdiv(V, BV), B * HV).
    """
    i_v = tl.program_id(0)
    i_bh = tl.program_id(1)
    i_b = i_bh // HV
    i_hv = i_bh % HV

    bos = i_b * T

    g = g_ptr + (bos * HV + i_hv) * K
    k = k_ptr + (bos * H + i_hv) * K
    w_p = w_ptr + (bos * HV + i_hv) * K
    u_p = u_ptr + (bos * HV + i_hv) * V
    v_new = v_new_ptr + (bos * HV + i_hv) * V
    h = h_ptr + ((i_b * NT) * HV + i_hv) * K * V

    # h state: [BK, BV] in registers, in fp32.
    b_h = tl.zeros([BK, BV], dtype=tl.float32)

    o_v = i_v * BV + tl.arange(0, BV)
    o_k = tl.arange(0, BK)
    o_t = tl.arange(0, BT)

    for i_t in range(NT):
        # Save h state
        h_offs = i_t * HV * K * V + o_k[:, None] * V + o_v[None, :]
        tl.store(h + h_offs, b_h.to(tl.bfloat16))

        # Load w[c]: [BT, K]
        w_offs = (i_t * BT + o_t[:, None]) * (HV * K) + o_k[None, :]
        b_w = tl.load(w_p + w_offs).to(tl.float32)

        # v_new = u - w @ h
        u_offs = (i_t * BT + o_t[:, None]) * (HV * V) + o_v[None, :]
        b_u = tl.load(u_p + u_offs).to(tl.float32)
        # w [BT, K] @ h [K, BV] -> [BT, BV]
        b_wh = tl.dot(b_w.to(tl.bfloat16), b_h.to(tl.bfloat16), out_dtype=tl.float32)
        b_v_new = b_u - b_wh
        tl.store(v_new + u_offs, b_v_new.to(tl.bfloat16))

        # Update h: h = h * exp(g[last, :]) + (k * exp(g_last - g)) @ v_new
        g_offs = (i_t * BT + o_t[:, None]) * (HV * K) + o_k[None, :]
        b_g = tl.load(g + g_offs).to(tl.float32)
        last_t = i_t * BT + BT - 1
        g_last_offs = last_t * (HV * K) + o_k
        b_g_last = tl.load(g + g_last_offs).to(tl.float32)
        b_g_last_exp = tl.exp(b_g_last)  # [K]

        k_offs = (i_t * BT + o_t[:, None]) * (H * K) + o_k[None, :]
        b_k = tl.load(k + k_offs).to(tl.float32)
        b_kg = b_k * tl.exp(b_g_last[None, :] - b_g)  # [BT, K]

        b_h = b_h * b_g_last_exp[:, None]
        b_kg_v = tl.dot(b_kg.to(tl.bfloat16), b_v_new.to(tl.bfloat16), out_dtype=tl.float32)
        b_h = b_h + b_kg_v


# ---------------------------------------------------------------------------
# Step 4: output o = (q * exp(g)) @ h + Aqk @ v_new
# ---------------------------------------------------------------------------


@triton.jit
def _kda_output_kernel(
    q_ptr, g_ptr, A_ptr, v_new_ptr, h_ptr, o_ptr,
    scale,
    NT, T,
    H: tl.constexpr, HV: tl.constexpr,
    K: tl.constexpr, V: tl.constexpr,
    BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr,
):
    """Output: o[c] = (q * exp(g)) @ h[c] + Aqk @ v_new[c].

    Grid: (cdiv(V, BV), NT, B * HV).
    """
    i_v = tl.program_id(0)
    i_t = tl.program_id(1)
    i_bh = tl.program_id(2)
    i_b = i_bh // HV
    i_hv = i_bh % HV

    bos = i_b * T
    o_v = i_v * BV + tl.arange(0, BV)
    o_t = i_t * BT + tl.arange(0, BT)
    m_t = o_t < T
    o_k = tl.arange(0, K)
    o_bt = tl.arange(0, BT)

    q = q_ptr + (bos * H + i_hv) * K
    g = g_ptr + (bos * HV + i_hv) * K
    A = A_ptr + (bos * HV + i_hv) * BT
    v_new = v_new_ptr + (bos * HV + i_hv) * V
    o = o_ptr + (bos * HV + i_hv) * V
    h = h_ptr + ((i_b * NT + i_t) * HV + i_hv) * K * V

    b_o = tl.zeros([BT, BV], dtype=tl.float32)

    # 1) o += Aqk @ v_new (intra-chunk part)
    A_offs = o_t[:, None] * (HV * BT) + o_bt[None, :]
    b_A = tl.load(A + A_offs, mask=m_t[:, None], other=0.0).to(tl.float32)
    v_offs = o_t[:, None] * (HV * V) + o_v[None, :]
    b_v = tl.load(v_new + v_offs, mask=m_t[:, None], other=0.0).to(tl.float32)
    m_ltri = o_bt[:, None] >= o_bt[None, :]
    b_A = tl.where(m_ltri, b_A, 0.0)
    b_o = b_o + tl.dot(b_A.to(tl.bfloat16), b_v.to(tl.bfloat16), out_dtype=tl.float32)

    # 2) o += (q * exp(g)) @ h
    g_offs = o_t[:, None] * (HV * K) + o_k[None, :]
    b_g = tl.load(g + g_offs, mask=m_t[:, None], other=0.0).to(tl.float32)
    q_offs = o_t[:, None] * (H * K) + o_k[None, :]
    b_q = tl.load(q + q_offs, mask=m_t[:, None], other=0.0).to(tl.float32)
    b_qg = b_q * tl.exp(b_g)
    h_offs = o_k[:, None] * V + o_v[None, :]
    b_h = tl.load(h + h_offs).to(tl.float32)
    b_o = b_o + tl.dot(b_qg.to(tl.bfloat16), b_h.to(tl.bfloat16), out_dtype=tl.float32)

    o_offs = o_t[:, None] * (HV * V) + o_v[None, :]
    tl.store(o + o_offs, b_o.to(tl.bfloat16), mask=m_t[:, None])


# ---------------------------------------------------------------------------
# Top-level wrapper
# ---------------------------------------------------------------------------


def kda_forward(
    q: torch.Tensor,    # (B, T, H, K) bf16
    k: torch.Tensor,    # (B, T, H, K) bf16
    v: torch.Tensor,    # (B, T, H, V) bf16
    g: torch.Tensor,    # (B, T, H, K) fp32
    beta: torch.Tensor, # (B, T, H) bf16
    scale: float,
    chunk_size: int = 64,
) -> torch.Tensor:
    B, T, H, K = q.shape
    V = v.shape[-1]
    BT = chunk_size
    NT = T // BT
    assert T % BT == 0
    assert H == v.shape[2]

    # Step 1: cumsum g within chunks
    g_cum = _chunk_local_cumsum_g(g.contiguous(), chunk_size).contiguous()  # (B, T, H, K)

    # Step 2: per-chunk Aqk + solve + w, u
    A = torch.empty((B, T, H, BT), dtype=torch.bfloat16, device=q.device)
    w = torch.empty((B, T, H, K), dtype=torch.bfloat16, device=q.device)
    u = torch.empty((B, T, H, V), dtype=torch.bfloat16, device=q.device)
    CENTER = (BT // 2) - 1  # middle row of the chunk
    grid_intra = (NT, B * H)
    _kda_intra_wu_kernel[grid_intra](
        q, k, v, g_cum, beta,
        A, w, u,
        scale,
        NT, T,
        H=H, HV=H,
        K=K, V=V,
        BT=BT, CENTER=CENTER,
        num_warps=4, num_stages=2,
    )

    # Step 3: inter-chunk recurrence; produces v_new and h
    v_new = torch.empty_like(u)
    h_buf = torch.empty((B, NT, H, K, V), dtype=torch.bfloat16, device=q.device)
    BV = 64 if V >= 64 else 32
    BK = min(64, K)
    grid_inter = (triton.cdiv(V, BV), B * H)
    _kda_inter_h_kernel[grid_inter](
        g_cum, k, w, u, v_new, h_buf,
        NT, T,
        H=H, HV=H,
        K=K, V=V,
        BT=BT, BK=BK, BV=BV,
        num_warps=4, num_stages=2,
    )

    # Step 4: output
    o = torch.empty((B, T, H, V), dtype=torch.bfloat16, device=q.device)
    grid_out = (triton.cdiv(V, BV), NT, B * H)
    _kda_output_kernel[grid_out](
        q, g_cum, A, v_new, h_buf, o,
        scale,
        NT, T,
        H=H, HV=H,
        K=K, V=V,
        BT=BT, BK=BK, BV=BV,
        num_warps=4, num_stages=2,
    )
    return o


# ---------------------------------------------------------------------------
# Module shim
# ---------------------------------------------------------------------------


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


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_forward(q, k, v, g, beta, self.scale, self.chunk_size)


# 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, 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]