"""Kimi Delta Attention forward (chunk form) — custom Triton kernels for H100.

Chunk-parallel KDA forward, written from scratch:

  Stage 1 (intra, one program per 64-token chunk per (b,h)):
    - in-chunk cumulative sum of the per-channel log-decay g (log2 space)
    - Aqk[c,j] = scale * q_c . (exp2(g_c - g_j) * k_j)   (c >= j)
    - Akk[c,j] = beta_c * k_c . (exp2(g_c - g_j) * k_j)  (c >  j)
    - M = (I + Akk)^{-1} via a Neumann-doubling product of the nilpotent part
    - w = M @ (beta * exp2(g) * k), u = M @ (beta * v)
    - kg = k * exp2(g_last - g), qg = scale * q * exp2(g), dl = exp2(g_last)

  Stage 2 (recurrent, one program per (b,h,v-block), sequential over chunks):
    - v_new = u - w @ S ; store S (entering state, bf16) and v_new
    - S = S * dl + kg^T @ v_new

  Stage 3 (output, fully parallel over chunks):
    - o = qg @ S + tril(Aqk) @ v_new
"""
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"]

RCP_LN2 = tl.constexpr(1.4426950408889634)  # 1 / ln(2)


@triton.jit
def _kda_fwd_intra_kernel(
    q, k, v, g, beta,
    qg, Aqk, w, u, kg, dl,
    scale,
    T: tl.constexpr, H: tl.constexpr,
    K: tl.constexpr, V: tl.constexpr,
    BT: tl.constexpr,
):
    i_t, i_bh = tl.program_id(0), tl.program_id(1)
    i_b, i_h = i_bh // H, i_bh % H
    bos = i_b * T
    NT = T // BT

    q += (bos * H + i_h) * K
    k += (bos * H + i_h) * K
    g += (bos * H + i_h) * K
    v += (bos * H + i_h) * V
    beta += bos * H + i_h
    qg += (bos * H + i_h) * K
    Aqk += (bos * H + i_h) * BT
    w += (bos * H + i_h) * K
    u += (bos * H + i_h) * V
    kg += (bos * H + i_h) * K
    dl += ((i_b * NT + i_t) * H + i_h) * K

    row0 = i_t * BT
    o_c = tl.arange(0, BT)
    m_strict = o_c[:, None] > o_c[None, :]
    m_incl = o_c[:, None] >= o_c[None, :]
    eye = (o_c[:, None] == o_c[None, :]).to(tl.float32)

    p_b = tl.make_block_ptr(beta, (T,), (H,), (row0,), (BT,), (0,))
    b_beta = tl.load(p_b).to(tl.float32)

    # ---- loop over K halves: cumsum + decay factors + A accumulation ----
    # Per-channel rebasing to the mid row keeps exp2 args small; the base
    # cancels inside every q/k dot product.
    b_Aqk = tl.zeros([BT, BT], dtype=tl.float32)
    b_Akk = tl.zeros([BT, BT], dtype=tl.float32)
    xw0 = tl.zeros([BT, 64], dtype=tl.bfloat16)
    xw1 = tl.zeros([BT, 64], dtype=tl.bfloat16)
    o_k = tl.arange(0, 64)
    for i_k in tl.static_range(K // 64):
        p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (row0, i_k * 64), (BT, 64), (1, 0))
        b_gc = tl.cumsum(tl.load(p_g).to(tl.float32) * RCP_LN2, 0)

        base = tl.sum(tl.where((o_c == BT // 2 - 1)[:, None], b_gc, 0.0), 0)
        gl = tl.sum(tl.where((o_c == BT - 1)[:, None], b_gc, 0.0), 0)
        tl.store(dl + i_k * 64 + o_k, tl.exp2(gl))

        p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (row0, i_k * 64), (BT, 64), (1, 0))
        b_k = tl.load(p_k)
        p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (row0, i_k * 64), (BT, 64), (1, 0))
        b_q = tl.load(p_q)

        b_e = tl.exp2(b_gc - base[None, :])
        kd = (b_k * b_e).to(tl.bfloat16)
        qd = (b_q * b_e * scale).to(tl.bfloat16)
        b_qg = (b_q * (b_e * tl.exp2(base)[None, :]) * scale).to(tl.bfloat16)
        p_qg = tl.make_block_ptr(qg, (T, K), (H * K, 1), (row0, i_k * 64), (BT, 64), (1, 0))
        tl.store(p_qg, b_qg)

        b_e = tl.exp2(base[None, :] - b_gc)
        ks = (b_k * b_e).to(tl.bfloat16)
        b_kg = (b_k * (b_e * tl.exp2(gl - base)[None, :])).to(tl.bfloat16)
        p_kg = tl.make_block_ptr(kg, (T, K), (H * K, 1), (row0, i_k * 64), (BT, 64), (1, 0))
        tl.store(p_kg, b_kg)

        kst = tl.trans(ks)
        b_Aqk = tl.dot(qd, kst, b_Aqk)
        b_Akk = tl.dot(kd, kst, b_Akk)

        b_xw = kd * (tl.exp2(base)[None, :] * b_beta[:, None]).to(tl.bfloat16)
        if i_k == 0:
            xw0 = b_xw
        else:
            xw1 = b_xw

    p_aqk = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (row0, 0), (BT, BT), (1, 0))
    tl.store(p_aqk, tl.where(m_incl, b_Aqk, 0.0).to(tl.bfloat16))

    # ---- M = (I + B)^{-1}, B = strict lower tril of beta-scaled Akk ----
    # With C = -B (nilpotent, C^64 = 0):
    #   M = sum_n C^n = (I+C)(I+C^2)(I+C^4)(I+C^8)(I+C^16)(I+C^32)
    C = tl.where(m_strict, -b_Akk * b_beta[:, None], 0.0)
    X = eye + C
    P = tl.dot(C, C)      # C^2
    X = X + tl.dot(X, P)
    P = tl.dot(P, P)      # C^4
    X = X + tl.dot(X, P)
    P = tl.dot(P, P)      # C^8
    X = X + tl.dot(X, P)
    P = tl.dot(P, P)      # C^16
    X = X + tl.dot(X, P)
    P = tl.dot(P, P)      # C^32
    X = X + tl.dot(X, P)
    Mb = X.to(tl.bfloat16)

    # ---- w = M @ (beta*exp2(g)*k), u = M @ (beta*v) ----
    b_w = tl.dot(Mb, xw0)
    p_w = tl.make_block_ptr(w, (T, K), (H * K, 1), (row0, 0), (BT, 64), (1, 0))
    tl.store(p_w, b_w.to(tl.bfloat16))
    b_w = tl.dot(Mb, xw1)
    p_w = tl.make_block_ptr(w, (T, K), (H * K, 1), (row0, 64), (BT, 64), (1, 0))
    tl.store(p_w, b_w.to(tl.bfloat16))

    p_v = tl.make_block_ptr(v, (T, V), (H * V, 1), (row0, 0), (BT, V), (1, 0))
    xu = (tl.load(p_v) * b_beta[:, None].to(tl.bfloat16))
    b_u = tl.dot(Mb, xu)
    p_u = tl.make_block_ptr(u, (T, V), (H * V, 1), (row0, 0), (BT, V), (1, 0))
    tl.store(p_u, b_u.to(tl.bfloat16))


@triton.jit
def _kda_fwd_h_kernel(
    kg, u, w, dl,
    h, v_new,
    T: tl.constexpr, H: tl.constexpr,
    K: tl.constexpr, V: tl.constexpr,
    BT: tl.constexpr, BV: tl.constexpr,
):
    i_v, i_bh = tl.program_id(0), tl.program_id(1)
    i_b, i_h = i_bh // H, i_bh % H
    bos = i_b * T
    NT = T // BT

    kg += (bos * H + i_h) * K
    u += (bos * H + i_h) * V
    w += (bos * H + i_h) * K
    dl += (i_b * NT * H + i_h) * K
    v_new += (bos * H + i_h) * V
    h += (i_b * NT * H + i_h) * K * V

    b_h1 = tl.zeros([64, BV], dtype=tl.float32)
    b_h2 = tl.zeros([64, BV], dtype=tl.float32)

    o_k = tl.arange(0, 64)
    for i_t in range(NT):
        # store entering state
        p_h1 = tl.make_block_ptr(h + i_t * H * K * V, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0))
        p_h2 = tl.make_block_ptr(h + i_t * H * K * V, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0))
        b_h1b = b_h1.to(tl.bfloat16)
        b_h2b = b_h2.to(tl.bfloat16)
        tl.store(p_h1, b_h1b)
        tl.store(p_h2, b_h2b)

        # v_new = u - w @ S
        p_w1 = tl.make_block_ptr(w, (T, K), (H * K, 1), (i_t * BT, 0), (BT, 64), (1, 0))
        p_w2 = tl.make_block_ptr(w, (T, K), (H * K, 1), (i_t * BT, 64), (BT, 64), (1, 0))
        b_vn = tl.dot(tl.load(p_w1), b_h1b) + tl.dot(tl.load(p_w2), b_h2b)
        p_u = tl.make_block_ptr(u, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
        b_vn = tl.load(p_u) - b_vn
        p_vn = tl.make_block_ptr(v_new, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
        tl.store(p_vn, b_vn.to(p_vn.dtype.element_ty))

        # decay state to chunk end, add new writes
        b_dl1 = tl.load(dl + i_t * H * K + o_k)
        b_dl2 = tl.load(dl + i_t * H * K + 64 + o_k)
        b_h1 *= b_dl1[:, None]
        b_h2 *= b_dl2[:, None]

        b_vnb = b_vn.to(tl.bfloat16)
        p_kg1 = tl.make_block_ptr(kg, (K, T), (1, H * K), (0, i_t * BT), (64, BT), (0, 1))
        p_kg2 = tl.make_block_ptr(kg, (K, T), (1, H * K), (64, i_t * BT), (64, BT), (0, 1))
        b_h1 = tl.dot(tl.load(p_kg1), b_vnb, b_h1)
        b_h2 = tl.dot(tl.load(p_kg2), b_vnb, b_h2)


@triton.jit
def _kda_fwd_o_kernel(
    qg, Aqk, v_new, h,
    o,
    T: tl.constexpr, H: tl.constexpr,
    K: tl.constexpr, V: tl.constexpr,
    BT: tl.constexpr, BV: tl.constexpr,
):
    i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
    i_b, i_h = i_bh // H, i_bh % H
    bos = i_b * T
    NT = T // BT

    qg += (bos * H + i_h) * K
    Aqk += (bos * H + i_h) * BT
    v_new += (bos * H + i_h) * V
    o += (bos * H + i_h) * V
    h += ((i_b * NT + i_t) * H + i_h) * K * V

    b_o = tl.zeros([BT, BV], dtype=tl.float32)
    for i_k in range(K // 64):
        p_qg = tl.make_block_ptr(qg, (T, K), (H * K, 1), (i_t * BT, i_k * 64), (BT, 64), (1, 0))
        b_qg = tl.load(p_qg)
        p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * 64, i_v * BV), (64, BV), (1, 0))
        b_h = tl.load(p_h)
        b_o += tl.dot(b_qg, b_h)

    m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :]
    p_A = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0))
    b_A = tl.load(p_A)
    b_A = tl.where(m_s, b_A, 0.0).to(tl.bfloat16)
    p_vn = tl.make_block_ptr(v_new, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
    b_vn = tl.load(p_vn)
    b_o += tl.dot(b_A, b_vn)

    p_o = tl.make_block_ptr(o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
    tl.store(p_o, b_o.to(p_o.dtype.element_ty))


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

    qg = torch.empty(B, T, H, K, dtype=torch.bfloat16, device=q.device)
    Aqk = 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)
    kg = torch.empty(B, T, H, K, dtype=torch.bfloat16, device=q.device)
    dl = torch.empty(B, NT, H, K, dtype=torch.float32, device=q.device)
    h = torch.empty(B, NT, H, K, V, dtype=torch.bfloat16, device=q.device)
    v_new = torch.empty(B, T, H, V, dtype=torch.bfloat16, device=q.device)
    o = torch.empty(B, T, H, V, dtype=torch.bfloat16, device=q.device)

    _kda_fwd_intra_kernel[(NT, B * H)](
        q, k, v, g, beta,
        qg, Aqk, w, u, kg, dl,
        scale,
        T=T, H=H, K=K, V=V, BT=BT,
        num_warps=4, num_stages=1,
    )

    BV_h = 32
    _kda_fwd_h_kernel[(V // BV_h, B * H)](
        kg, u, w, dl,
        h, v_new,
        T=T, H=H, K=K, V=V, BT=BT, BV=BV_h,
        num_warps=4, num_stages=3,
    )

    BV_o = 64
    _kda_fwd_o_kernel[(V // BV_o, NT, B * H)](
        qg, Aqk, v_new, h,
        o,
        T=T, H=H, K=K, V=V, BT=BT, BV=BV_o,
        num_warps=4, num_stages=3,
    )
    return o


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


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