"""Triton circuit kernels adapted from Qwenputer by Divyansh Agrawal.
MIT license: see third_party/QWENPUTER_LICENSE.
Changes: explicit per-row active masks to freeze halted batch rows.
"""
import triton
import triton.language as tl

@triton.jit
def _ffn_kernel(x_ptr, y_ptr, D, U, rpa, ixa, wa, rpb, ixb, wb, coef, outd, active_ptr,
                BLOCK: tl.constexpr, KA: tl.constexpr, KB: tl.constexpr):
    pid = tl.program_id(0)
    bat = tl.program_id(1)
    if tl.load(active_ptr + bat) == 0:
        return
    u = pid * BLOCK + tl.arange(0, BLOCK)
    m = u < U
    xb = x_ptr + bat * D
    a0 = tl.load(rpa + u, mask=m, other=0)
    na = tl.load(rpa + u + 1, mask=m, other=0) - a0
    ka = tl.arange(0, KA)
    ma = m[:, None] & (ka[None, :] < na[:, None])
    ia = tl.load(ixa + a0[:, None] + ka[None, :], mask=ma, other=0)
    av = tl.sum(tl.load(wa + a0[:, None] + ka[None, :], mask=ma, other=0.0)
                * tl.load(xb + ia, mask=ma, other=0.0), axis=1)
    b0 = tl.load(rpb + u, mask=m, other=0)
    nb = tl.load(rpb + u + 1, mask=m, other=0) - b0
    kb = tl.arange(0, KB)
    mb = m[:, None] & (kb[None, :] < nb[:, None])
    ib = tl.load(ixb + b0[:, None] + kb[None, :], mask=mb, other=0)
    bv = tl.sum(tl.load(wb + b0[:, None] + kb[None, :], mask=mb, other=0.0)
                * tl.load(xb + ib, mask=mb, other=0.0), axis=1)
    val = tl.load(coef + u, mask=m, other=0.0) * av * tl.maximum(bv, 0.0)
    od = tl.load(outd + u, mask=m, other=0)
    tl.atomic_add(y_ptr + bat * D + od, val, mask=m)


@triton.jit
def _lag_load(xb, hb, i, m, D):
    """The lagged residual: index lag * D + d (lag 0: x, lags 1..3: the history)."""
    lo = i < D
    return tl.where(lo, tl.load(xb + i, mask=m & lo, other=0.0),
                    tl.load(hb + i - D, mask=m & (i >= D), other=0.0))


@triton.jit
def _rows(xb, hb, D, rp, ix, w, base, ch, KP: tl.constexpr):
    """relu of 128 sparse rows (row base + ch) over the lagged residual."""
    r0 = tl.load(rp + base + ch)
    n = tl.load(rp + base + ch + 1) - r0
    kk = tl.arange(0, KP)
    mk = kk[None, :] < n[:, None]
    i = tl.load(ix + r0[:, None] + kk[None, :], mask=mk, other=0)
    acc = tl.sum(tl.load(w + r0[:, None] + kk[None, :], mask=mk, other=0.0)
                 * _lag_load(xb, hb, i, mk, D), axis=1)
    return tl.maximum(acc, 0.0)


@triton.jit
def _scalar(xb, hb, D, rp, ix, w, vh):
    lane = tl.arange(0, 32)
    s0 = tl.load(rp + vh)
    n = tl.load(rp + vh + 1) - s0
    mk = lane < n
    i = tl.load(ix + s0 + lane, mask=mk, other=0)
    return tl.sum(tl.load(w + s0 + lane, mask=mk, other=0.0) * _lag_load(xb, hb, i, mk, D), axis=0)


@triton.jit
def _head_kernel(x_ptr, h_ptr, S_ptr, y_ptr, D, V,
                 qrp, qix, qw, krp, kix, kw, vrp, vix, vw, zrp, zix, zw,
                 arp, aix, aw, brp, bix, bw, bconst, odim, ow, khead, nrows,
                 active_ptr, KQ: tl.constexpr, KK: tl.constexpr, KV: tl.constexpr, KZ: tl.constexpr,
                 QHOT: tl.constexpr):
    vh = tl.program_id(0)
    bat = tl.program_id(1)
    if tl.load(active_ptr + bat) == 0:
        return
    ch = tl.arange(0, 128)
    xb = x_ptr + bat * D
    hb = h_ptr + bat * 3 * D
    h = tl.load(khead + vh)
    k = _rows(xb, hb, D, krp, kix, kw, h * 128, ch, KK)
    v = _rows(xb, hb, D, vrp, vix, vw, vh * 128, ch, KV)
    z = _rows(xb, hb, D, zrp, zix, zw, vh * 128, ch, KZ)
    q = _rows(xb, hb, D, qrp, qix, qw, h * 128, ch, KQ)
    bv = _scalar(xb, hb, D, brp, bix, bw, vh) + tl.load(bconst + vh)
    Sb = S_ptr + (bat * V + vh) * 16384
    nr = tl.load(nrows + vh)  # > 0: the head decays, in its first nr rows
    if nr > 0:
        av = _scalar(xb, hb, D, arp, aix, aw, vh)
        fac = tl.where(av == 0.0, 1.0, tl.where(av == 110.0, 0.5, 0.0))
        if fac != 1.0:
            for r in range(nr):
                tl.store(Sb + r * 128 + ch, tl.load(Sb + r * 128 + ch) * fac)
    kmax = tl.max(k, axis=0)
    krow = tl.argmax(k, axis=0)
    if (kmax > 0.0) & (bv == 1.0):
        tl.store(Sb + krow * 128 + ch, v)
    tl.debug_barrier()
    # the read: the sum of the rows whose query is set (multi-hot queries have equal values)
    o = tl.zeros([128], dtype=tl.float32)
    for _ in range(QHOT):
        qm = tl.max(q, axis=0)
        qi = tl.argmax(q, axis=0)
        if qm > 0.0:
            o += tl.load(Sb + qi * 128 + ch)
        q = tl.where(ch == qi, 0.0, q)
    ref = tl.sum(tl.where(ch < 2, o * o, 0.0), axis=0)
    norm = tl.math.sqrt_rn(ref / 128.0)
    out = tl.where(ref > 0.0, tl.math.div_rn(o, tl.where(ref > 0.0, norm, 1.0)), 0.0)
    wv = tl.load(ow + vh * 128 + ch)
    dd = tl.load(odim + vh * 128 + ch)
    tl.atomic_add(y_ptr + bat * D + dd, out * z * wv, mask=wv != 0.0)


@triton.jit
def _apply_kernel(x_ptr, y_ptr, h_ptr, D, active_ptr, HIST: tl.constexpr, BLOCK: tl.constexpr):
    """x += y, y = 0; with HIST, the lag history (h: [3, D] per row, h[0] = lag 1) moves on
    first, the input x becoming lag 1."""
    pid = tl.program_id(0)
    bat = tl.program_id(1)
    if tl.load(active_ptr + bat) == 0:
        return
    i = pid * BLOCK + tl.arange(0, BLOCK)
    m = i < D
    xo = x_ptr + bat * D + i
    x = tl.load(xo, mask=m, other=0.0)
    if HIST:
        hb = h_ptr + bat * 3 * D + i
        h1 = tl.load(hb + D, mask=m, other=0.0)
        h0 = tl.load(hb, mask=m, other=0.0)
        tl.store(hb + 2 * D, h1, mask=m)
        tl.store(hb + D, h0, mask=m)
        tl.store(hb, x, mask=m)
    yo = y_ptr + bat * D + i
    tl.store(xo, x + tl.load(yo, mask=m, other=0.0), mask=m)
    tl.store(yo, tl.zeros([BLOCK], dtype=tl.float32), mask=m)