"""Fused Triton scan for the PSSA selective diagonal state-space recurrence.

    h_t = exp(delta_t * A) * h_{t-1} + delta_t * b_t * x_t
    y_t = <h_t, c_t>

`A` is (D, DS) negative, `delta` is (B, T, D) positive, `b`/`c` are (B, T, DS).
The whole (forward + backward) runs as one kernel launch per chunk, with the
per-token state `h_t` materialised so the backward pass is a single reverse
sweep.  Carries `h0` across chunks and returns `hT`.
"""
import torch
import triton
import triton.language as tl

__all__ = ["ScanFn", "scan_reference", "selftest"]


@triton.jit
def scan_fwd(x_ptr, d_ptr, bp_ptr, cp_ptr, A_ptr, h0_ptr, y_ptr, hT_ptr, H_ptr,
             T: tl.constexpr, D: tl.constexpr, DS: tl.constexpr, BD: tl.constexpr):
    b = tl.program_id(0)
    pd = tl.program_id(1)
    od = pd * BD + tl.arange(0, BD)
    os_ = tl.arange(0, DS)
    A = tl.load(A_ptr + od[:, None] * DS + os_[None, :])
    h = tl.load(h0_ptr + b * D * DS + od[:, None] * DS + os_[None, :])
    for t in range(T):
        x = tl.load(x_ptr + b * T * D + t * D + od).to(tl.float32)
        dl = tl.load(d_ptr + b * T * D + t * D + od).to(tl.float32)
        bp = tl.load(bp_ptr + b * T * DS + t * DS + os_).to(tl.float32)
        cp = tl.load(cp_ptr + b * T * DS + t * DS + os_).to(tl.float32)
        h = tl.exp(dl[:, None] * A) * h + dl[:, None] * bp[None, :] * x[:, None]
        tl.store(y_ptr + b * T * D + t * D + od, tl.sum(h * cp[None, :], axis=1))
        tl.store(H_ptr + b * T * D * DS + t * D * DS + od[:, None] * DS + os_[None, :], h)
    tl.store(hT_ptr + b * D * DS + od[:, None] * DS + os_[None, :], h)


@triton.jit
def scan_bwd(dy_ptr, x_ptr, d_ptr, bp_ptr, cp_ptr, A_ptr, H_ptr, h0_ptr,
             gh0_ptr, gx_ptr, gd_ptr, gbp_ptr, gcp_ptr, gA_ptr,
             T: tl.constexpr, D: tl.constexpr, DS: tl.constexpr, BD: tl.constexpr):
    b = tl.program_id(0)
    pd = tl.program_id(1)
    od = pd * BD + tl.arange(0, BD)
    os_ = tl.arange(0, DS)
    A = tl.load(A_ptr + od[:, None] * DS + os_[None, :])
    g = tl.zeros((BD, DS), tl.float32)
    gA = tl.zeros((BD, DS), tl.float32)
    for t in range(T - 1, -1, -1):
        dy = tl.load(dy_ptr + b * T * D + t * D + od).to(tl.float32)
        x = tl.load(x_ptr + b * T * D + t * D + od).to(tl.float32)
        dl = tl.load(d_ptr + b * T * D + t * D + od).to(tl.float32)
        bp = tl.load(bp_ptr + b * T * DS + t * DS + os_).to(tl.float32)
        cp = tl.load(cp_ptr + b * T * DS + t * DS + os_).to(tl.float32)
        h = tl.load(H_ptr + b * T * D * DS + t * D * DS + od[:, None] * DS + os_[None, :])
        if t < T - 1:
            dln = tl.load(d_ptr + b * T * D + (t + 1) * D + od).to(tl.float32)
            g = dy[:, None] * cp[None, :] + tl.exp(dln[:, None] * A) * g
        else:
            g = dy[:, None] * cp[None, :]
        if t > 0:
            hp = tl.load(H_ptr + b * T * D * DS + (t - 1) * D * DS + od[:, None] * DS + os_[None, :])
        else:
            hp = tl.load(h0_ptr + b * D * DS + od[:, None] * DS + os_[None, :])
        a = tl.exp(dl[:, None] * A)
        ga = g * hp
        gd = tl.sum(ga * a * A, axis=1) + tl.sum(g * bp[None, :] * x[:, None], axis=1)
        gbp = tl.sum(g * dl[:, None] * x[:, None], axis=0)
        gcp = tl.sum(dy[:, None] * h, axis=0)
        gx = tl.sum(g * dl[:, None] * bp[None, :], axis=1)
        gA += ga * a * dl[:, None]
        if t == 0:
            tl.store(gh0_ptr + b * D * DS + od[:, None] * DS + os_[None, :], a * g)
        tl.store(gx_ptr + b * T * D + t * D + od, gx)
        tl.store(gd_ptr + b * T * D + t * D + od, gd)
        tl.atomic_add(gbp_ptr + b * T * DS + t * DS + os_, gbp)
        tl.atomic_add(gcp_ptr + b * T * DS + t * DS + os_, gcp)
    tl.atomic_add(gA_ptr + od[:, None] * DS + os_[None, :], gA)


class ScanFn(torch.autograd.Function):
    """Differentiable fused scan. Returns (y, hT); hT is not differentiated."""

    @staticmethod
    def forward(ctx, x, delta, bp, cp, A, h0, bd=16):
        B, T, D = x.shape
        DS = A.shape[1]
        y = torch.empty(B, T, D, device=x.device, dtype=torch.float32)
        hT = torch.empty(B, D, DS, device=x.device, dtype=torch.float32)
        H = torch.empty(B, T, D, DS, device=x.device, dtype=torch.float32)
        scan_fwd[(B, triton.cdiv(D, bd))](
            x, delta, bp, cp, A, h0, y, hT, H, T=T, D=D, DS=DS, BD=bd)
        ctx.save_for_backward(x, delta, bp, cp, A, H, h0)
        ctx.bd = bd
        ctx.mark_non_differentiable(hT)
        return y, hT

    @staticmethod
    def backward(ctx, dy, dhT):
        x, delta, bp, cp, A, H, h0 = ctx.saved_tensors
        B, T, D = x.shape
        DS = A.shape[1]
        bd = ctx.bd
        gx = torch.empty(B, T, D, device=x.device, dtype=torch.float32)
        gd = torch.empty_like(delta, dtype=torch.float32)
        gbp = torch.zeros_like(bp, dtype=torch.float32)
        gcp = torch.zeros_like(cp, dtype=torch.float32)
        gA = torch.zeros_like(A, dtype=torch.float32)
        gh0 = torch.zeros(B, D, DS, device=x.device, dtype=torch.float32)
        scan_bwd[(B, triton.cdiv(D, bd))](
            dy.contiguous(), x, delta, bp, cp, A, H, h0, gh0,
            gx, gd, gbp, gcp, gA, T=T, D=D, DS=DS, BD=bd)
        return gx.to(x.dtype), gd, gbp, gcp, gA, gh0.to(h0.dtype), None


def scan_reference(x, delta, bp, cp, A, h0):
    """Plain PyTorch loop, for tests. Returns (y, H, hT)."""
    B, T, D = x.shape
    h = h0
    ys, hs = [], []
    for t in range(T):
        h = torch.exp(delta[:, t, :, None] * A) * h + delta[:, t, :, None] * bp[:, t, None, :] * x[:, t, :, None]
        ys.append((h * cp[:, t, None, :]).sum(-1))
        hs.append(h)
    return torch.stack(ys, 1), torch.stack(hs, 1), h


def selftest(ds=16, d=64, b=2, t=7, verbose=True):
    torch.manual_seed(1)
    dev = "cuda"
    x = torch.randn(b, t, d, device=dev, requires_grad=True)
    delta = torch.nn.functional.softplus(torch.randn(b, t, d, device=dev)).requires_grad_(True)
    bp = torch.randn(b, t, ds, device=dev, requires_grad=True)
    cp = torch.randn(b, t, ds, device=dev, requires_grad=True)
    A = (-torch.nn.functional.softplus(torch.randn(d, ds, device=dev))).requires_grad_(True)
    h0 = (torch.randn(b, d, ds, device=dev) * 0.5).requires_grad_(True)
    yr, _, _ = scan_reference(x, delta, bp, cp, A, h0)
    y, hT = ScanFn.apply(x, delta, bp, cp, A, h0)
    gy = torch.randn_like(yr)
    gref = torch.autograd.grad(yr, (x, delta, bp, cp, A, h0), gy)
    gtri = torch.autograd.grad(y, (x, delta, bp, cp, A, h0), gy)
    out = {"fwd": ((y - yr).abs().max() / yr.abs().max()).item()}
    for nm, a, bb in zip(["dx", "dd", "dbp", "dcp", "dA", "dh0"], gtri, gref):
        out[nm] = ((a - bb).abs().max() / bb.abs().max()).item()
    if verbose:
        print("scan selftest rel errs: " + "  ".join(f"{k}={v:.1e}" for k, v in out.items()))
    return out


if __name__ == "__main__":
    selftest()