"""KDA forward (chunk form) — Triton kernels for SM90.

3-kernel pipeline:
  K1 "kk":  per (chunk, b*h): g cumsum in-register (tril-ones tf32 dot, log2
     units), decayed qk/kk dots (fp16 MMA, fp32 acc). Emits Aqk (masked,
     scale folded), qg, kg (transposed layout), ke (beta folded), Akk
     (negated strictly-lower), e^{gc_last}.
  K2 "solve_wu": per (chunk, b*h): load Akk, invert (I-A) in-register
     (Neumann doubling, fp16 MMA), w = Ainv @ ke, u = Ainv @ (v*beta).
  K3 "scan_out" (fused): per (b*h, V-block): sequential over chunks,
     S = K x BV fp32 state in registers:
       o = qg @ S + Aqk @ (u - w @ S)      (written out directly)
       S = S * 2^gc_last + kgT @ (u - w @ S)

Host-side: per-shape cached workspace + cached CompiledKernel handles with
direct `.run` launches to keep per-call overhead ~10 us.
"""
from __future__ import annotations

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

LOG2E = tl.constexpr(1.4426950408889634)


@triton.jit
def _kda_kk_kernel(
    q, k, g, beta,
    Aqk, Akk, qg, kg, ke, gcl,
    scale,
    T,
    H: tl.constexpr, K: tl.constexpr,
    BT: tl.constexpr, BK: tl.constexpr, NK: tl.constexpr,
):
    i_t, i_bh = tl.program_id(0), tl.program_id(1)
    i_b = i_bh // H
    i_h = i_bh % H

    t0 = i_t * BT
    o_t = t0 + tl.arange(0, BT)
    o_r = tl.arange(0, BT)

    # lower-tri ones for matrix cumsum (tf32)
    m_qk = o_r[:, None] >= o_r[None, :]
    m_low = o_r[:, None] > o_r[None, :]
    b_L = tl.where(m_qk, 1.0, 0.0)

    b_beta = tl.load(beta + (i_b * T + o_t) * H + i_h).to(tl.float32)

    b_Aqk = tl.zeros([BT, BT], dtype=tl.float32)
    b_A = tl.zeros([BT, BT], dtype=tl.float32)

    for i_k in tl.static_range(NK):
        o_k = i_k * BK + tl.arange(0, BK)
        pg = g + ((i_b * T + o_t[:, None]) * H + i_h) * K + o_k[None, :]
        b_g = tl.load(pg)
        # cumsum along time via triangular-ones matmul (tf32 MMA)
        b_gc = tl.dot(b_L, b_g, input_precision="tf32") * LOG2E

        pq = q + ((i_b * T + o_t[:, None]) * H + i_h) * K + o_k[None, :]
        pk = k + ((i_b * T + o_t[:, None]) * H + i_h) * K + o_k[None, :]
        b_q = tl.load(pq).to(tl.float32)
        b_k = tl.load(pk).to(tl.float32)

        b_e = tl.exp2(b_gc)
        b_ei = tl.exp2(-b_gc)

        b_qe = (b_q * b_e * scale).to(tl.float16)
        b_ke = (b_k * b_e).to(tl.float16)
        b_kei = (b_k * b_ei).to(tl.float16)

        b_Aqk += tl.dot(b_qe, tl.trans(b_kei))
        b_A += tl.dot(b_ke, tl.trans(b_kei))

        pqg = qg + ((i_b * T + o_t[:, None]) * H + i_h) * K + o_k[None, :]
        tl.store(pqg, b_qe)
        pke = ke + ((i_b * T + o_t[:, None]) * H + i_h) * K + o_k[None, :]
        tl.store(pke, (b_ke.to(tl.float32) * b_beta[:, None]).to(tl.float16))

        b_gcl = tl.exp2(tl.sum(tl.where((o_r == BT - 1)[:, None], b_gc, 0.0), axis=0))  # (BK,)
        NTdim = tl.num_programs(0)
        pgcl = gcl + ((i_b * NTdim + i_t) * H + i_h) * K + o_k
        tl.store(pgcl, b_gcl)

        b_kg = (b_kei.to(tl.float32) * b_gcl[None, :]).to(tl.float16)
        pkg = kg + (((i_b * NTdim + i_t) * H + i_h) * K + o_k[:, None]) * BT + o_r[None, :]
        tl.store(pkg, tl.trans(b_kg))

    b_Aqk = tl.where(m_qk, b_Aqk, 0.0)
    pAqk = Aqk + ((i_b * T + o_t[:, None]) * H + i_h) * BT + o_r[None, :]
    tl.store(pAqk, b_Aqk.to(tl.float16))

    b_A = -tl.where(m_low, b_A * b_beta[:, None], 0.0)
    pAkk = Akk + ((i_b * T + o_t[:, None]) * H + i_h) * BT + o_r[None, :]
    tl.store(pAkk, b_A.to(tl.float16))


@triton.jit
def _kda_solve_wu_kernel(
    Akk, ke, v, beta,
    w, u,
    T,
    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_bh // H
    i_h = i_bh % H

    t0 = i_t * BT
    o_t = t0 + tl.arange(0, BT)
    o_r = tl.arange(0, BT)

    pAkk = Akk + ((i_b * T + o_t[:, None]) * H + i_h) * BT + o_r[None, :]
    b_A16 = tl.load(pAkk)  # fp16

    # inverse of (I - A): Neumann product doubling with fp16 MMA
    m_I = (o_r[:, None] == o_r[None, :]).to(tl.float16)
    b_Y = b_A16 + m_I
    b_P = b_A16
    for _ in tl.static_range(4):
        b_P = tl.dot(b_P, b_P).to(tl.float16)
        b_Y = b_Y + tl.dot(b_Y, b_P).to(tl.float16)
    b_Ai = b_Y + tl.dot(b_Y, tl.dot(b_P, b_P).to(tl.float16)).to(tl.float16)

    o_k = tl.arange(0, K)
    pke = ke + ((i_b * T + o_t[:, None]) * H + i_h) * K + o_k[None, :]
    b_ke = tl.load(pke)
    b_w = tl.dot(b_Ai, b_ke)
    pw = w + ((i_b * T + o_t[:, None]) * H + i_h) * K + o_k[None, :]
    tl.store(pw, b_w.to(tl.float16))

    b_beta = tl.load(beta + (i_b * T + o_t) * H + i_h).to(tl.float32)
    o_v = tl.arange(0, V)
    pv = v + ((i_b * T + o_t[:, None]) * H + i_h) * V + o_v[None, :]
    b_v = tl.load(pv).to(tl.float32)
    b_vb = (b_v * b_beta[:, None]).to(tl.float16)
    b_u = tl.dot(b_Ai, b_vb)
    pu = u + ((i_b * T + o_t[:, None]) * H + i_h) * V + o_v[None, :]
    tl.store(pu, b_u.to(tl.float16))


@triton.jit
def _kda_scan_out_kernel(
    w, u, qg, kg, Aqk, gcl, o,
    T, NT,
    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_bh // H
    i_h = i_bh % H

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

    b_h = tl.zeros([K, BV], dtype=tl.float32)

    w_base = w + (i_b * T * H + i_h) * K
    u_base = u + (i_b * T * H + i_h) * V
    qg_base = qg + (i_b * T * H + i_h) * K
    kg_base = kg + (i_b * NT * H + i_h).to(tl.int64) * K * BT
    Aqk_base = Aqk + (i_b * T * H + i_h) * BT
    gcl_base = gcl + (i_b * NT * H + i_h) * K
    o_base = o + (i_b * T * H + i_h) * V

    for c in range(NT):
        t0 = c * BT
        # ---- recurrence-critical path first ----
        pw = w_base + (t0 + o_t[:, None]) * H * K + o_k[None, :]
        b_w = tl.load(pw)
        b_h16 = b_h.to(tl.float16)
        b_vn = tl.dot(b_w, b_h16)  # (BT, BV)

        pqg = qg_base + (t0 + o_t[:, None]) * H * K + o_k[None, :]
        b_qg = tl.load(pqg)
        b_o = tl.dot(b_qg, b_h16)

        pu = u_base + (t0 + o_t[:, None]) * H * V + o_v[None, :]
        b_u = tl.load(pu).to(tl.float32)
        b_vn = b_u - b_vn
        b_vn16 = b_vn.to(tl.float16)

        pgcl = gcl_base + c * H * K + o_k
        b_gcl = tl.load(pgcl)
        pkg = kg_base + (c * H * K + o_k[:, None]) * BT + o_r[None, :]
        b_kg = tl.load(pkg)
        b_h = b_h * b_gcl[:, None]
        b_h = b_h + tl.dot(b_kg, b_vn16)  # state update completes next-iter dependency

        # ---- output path (off the recurrence chain) ----
        pA = Aqk_base + (t0 + o_t[:, None]) * H * BT + o_r[None, :]
        b_A = tl.load(pA)
        b_o += tl.dot(b_A, b_vn16)
        po = o_base + (t0 + o_t[:, None]) * H * V + o_v[None, :]
        tl.store(po, b_o.to(tl.bfloat16))


# ---------------------------------------------------------------------------
# host-side fast path
# ---------------------------------------------------------------------------

def _cfg(B, T, H, K, V):
    # (kk: (BK, warps, stages), solve: (warps, stages), scan: (BV, warps, stages))
    if B * H <= 8:
        return (128, 4, 2), (4, 3), (16, 4, 3)
    return (128, 4, 2), (4, 3), (32, 4, 3)


class _ShapePlan:
    __slots__ = ("grid_kk", "grid_scan", "cfg", "kk", "solve", "scan", "ws")

    def __init__(self):
        self.ws = None
        self.kk = None
        self.solve = None
        self.scan = None


_PLANS: dict = {}


def _get_stream(device_index):
    return torch.cuda.current_stream(device_index).cuda_stream


def _t(dtype, shape, dev):
    return torch.empty(shape, dtype=dtype, device=dev)


def kda_chunk_fwd(q, k, v, g, beta, scale, chunk_size=64):
    B, T, H, K = q.shape
    V = v.shape[-1]
    BT = chunk_size
    NT = T // BT
    dev = q.device
    f16 = torch.float16
    f32 = torch.float32

    q = q.contiguous(); k = k.contiguous(); v = v.contiguous()
    g = g.contiguous(); beta = beta.contiguous()

    key = (B, T, H, K, V, dev.index)
    plan = _PLANS.get(key)
    if plan is None:
        plan = _ShapePlan()
        plan.cfg = _cfg(B, T, H, K, V)
        plan.grid_kk = (NT, B * H)
        (kk_bk, kk_w, kk_s), (sv_w, sv_s), (sc_bv, sc_w, sc_s) = plan.cfg
        plan.grid_scan = (V // sc_bv, B * H)

        Aqk = _t(f16, (B, T, H, BT), dev)
        Akk = _t(f16, (B, T, H, BT), dev)
        qg = _t(f16, (B, T, H, K), dev)
        kg = _t(f16, (B, T, H, K, ), dev)
        ke = _t(f16, (B, T, H, K), dev)
        gcl = _t(f32, (B, NT, H, K), dev)
        w = _t(f16, (B, T, H, K), dev)
        u = _t(f16, (B, T, H, V), dev)
        o = _t(torch.bfloat16, (B, T, H, V), dev)
        plan.ws = (Aqk, Akk, qg, kg, ke, gcl, w, u)

        plan.kk = _kda_kk_kernel[plan.grid_kk](
            q, k, g, beta, Aqk, Akk, qg, kg, ke, gcl, scale, T,
            H=H, K=K, BT=BT, BK=kk_bk, NK=K // kk_bk,
            num_warps=kk_w, num_stages=kk_s)
        plan.solve = _kda_solve_wu_kernel[plan.grid_kk](
            Akk, ke, v, beta, w, u, T,
            H=H, K=K, V=V, BT=BT, num_warps=sv_w, num_stages=sv_s)
        plan.scan = _kda_scan_out_kernel[plan.grid_scan](
            w, u, qg, kg, Aqk, gcl, o, T, NT,
            H=H, K=K, V=V, BT=BT, BV=sc_bv,
            num_warps=sc_w, num_stages=sc_s)
        plan.kk._init_handles()
        plan.solve._init_handles()
        plan.scan._init_handles()
        _PLANS[key] = plan
        # first call returns o via the JIT launches above
        return o

    Aqk, Akk, qg, kg, ke, gcl, w, u = plan.ws
    o = _t(torch.bfloat16, (B, T, H, V), dev)
    stream = _get_stream(dev.index)
    (kk_bk, kk_w, kk_s), (sv_w, sv_s), (sc_bv, sc_w, sc_s) = plan.cfg

    ck = plan.kk
    ck.run(plan.grid_kk[0], plan.grid_kk[1], 1, stream, ck.function,
           ck.packed_metadata, None, None, None,
           q, k, g, beta, Aqk, Akk, qg, kg, ke, gcl, scale, T,
           H, K, BT, kk_bk, K // kk_bk)
    ck = plan.solve
    ck.run(plan.grid_kk[0], plan.grid_kk[1], 1, stream, ck.function,
           ck.packed_metadata, None, None, None,
           Akk, ke, v, beta, w, u, T, H, K, V, BT)
    ck = plan.scan
    ck.run(plan.grid_scan[0], plan.grid_scan[1], 1, stream, ck.function,
           ck.packed_metadata, None, None, None,
           w, u, qg, kg, Aqk, gcl, o, T, NT, H, K, V, BT, sc_bv)
    return o


class Model(nn.Module):
    """KDA forward (chunk form). No learned parameters."""

    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, k, v, g, beta):
        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]