"""fp32-invariant rewrites of the Kev row graph that keep fp16 compute (GPU float16 activations or float16 storage,
an NPU) closer to the fp32 result. The final-kernel files are the form R64+sp+ec+dd+vs6 (Kev-0.8B) and, through
r15_form.py, R64+sp+ec+dd+vs6+in1+fn5 (Kev-4B).

A form is a '+'-joined spec applied to a kev_graph.load_text_model() model:

    import r13_kernel as R
    R.apply(model, "R64+sp+ec+dd+vs6")
    R.reset(model)

tokens
  loop | R64   the in-chunk inverse (I - A)^-1: the loop kernel's 63-step forward substitution, or recursive doubling
               (X <- I + A*M_1; X <- X + X (A*M_s) X for s = 2..32). Both copied from r12_kernel_product.py, which is
               not imported.
  sp           the softplus rewrite (r11_fp16_safe.apply("softplus"); the guarded GatedDeltaNet forward's F)
  ec           the exp clamp: every EXP argument in the kernel clamped at -80 (r11_fp16_safe.EXP_FLOOR)
  dd           direct decay sums: the kernel builds its decays from the cumulative sum g_c = cumsum(g) as differences,
               exp(g_c[i] - g_c[j]) (decay mask) and exp(g_c[-1] - g_c[i]) (state update). g_c sums up to 64 decays
               and grows large inside a chunk; where |g_c| >= 128, fp16 holds steps of 0.125 or more, so the
               difference of two such numbers carries an absolute error of that size exactly where it is small and
               matters. dd forms the window sums directly:
                   diff[i, j] = sum_{t=j+1..i} g_t = g @ W   (W constant 0/1 [c, c*c], W[t, i*c+j] = [j < t <= i])
                   suffix[i]  = sum_{t>i} g_t      = g @ S   (S constant 0/1 [c, c],   S[t, i] = [t > i])
               so every partial sum holds only the terms inside the window (the masked terms are exact zeros) and the
               absolute error scales with the window sum, not with g_c. exp(g_c[i]) (prefix decays) keeps the cumsum:
               there the value itself is the large number. fp32: the same sums in another order (rounding level).
  vs<k>        the kernel scales v by 2^k at its entry (the delta rule is linear in v, so the output is exactly 2^k
               larger) and the gated norm after it takes eps * 4^k: the norm's input moves up out of the fp16
               subnormal range by 2^k and its sum of squares by 4^k, an exact power-of-two scaling in fp32. The
               recurrent state the kernel carries is 2^k times larger too (a shared-state pair passes it unchanged
               between its two signatures). vs6 is the final files' value (fp16 headroom of the gated norm's sum of
               squares: r14_norm_range_probe.py).
  gn, gs<j>, nr, rs, fc<n>k<k>, bmc<chain>, sg, sga, sg2
               further rewrites tried for float16 compute (the gated norm pre-scale, a Newton step after RSQRT,
               two-stage reductions, FULLY_CONNECTED splits, sigmoid as EXP / DIV); not used by the final files.

Capture: CAPTURE["on"] and CAPTURE["layers"] select GatedDeltaNet layers whose kernel intermediates are appended to
CAPTURE["out"] as (f"L{layer:02d}_{name}", tensor) in execution order (a probe reads them). Off by default.

Nothing outside this file is modified (kev_qwen35_patch, r11_fp16_safe, r12_kernel_product are read only)."""
import torch

import kev_qwen35_patch as P

EXP_FLOOR = -80.0                     # = r11_fp16_safe.EXP_FLOOR
CAPTURE = {"on": False, "layers": (), "out": []}


def _cap(layer, name, t):
    if CAPTURE["on"] and layer in CAPTURE["layers"]:
        CAPTURE["out"].append((f"L{layer:02d}_{name}", t))


# ---- in-chunk inverse (copied from scripts/r12_kernel_product.py, unchanged) -------------------------------------------
def inverse_loop(attn, eye):
    """The shipped forward substitution (kev_qwen35_patch, same ops in the same order)."""
    c = attn.shape[-1]
    rows = [attn[..., 0, :]]
    for i in range(1, c):
        prev = torch.stack(rows, dim=-2)                    # [H, n, i, c]
        row_full = attn[..., i, :]                          # [H, n, c]
        row_lead = row_full[..., :i]                        # [H, n, i]
        sub = prev[..., :i]                                 # [H, n, i, i]
        upd = row_lead + (row_lead.unsqueeze(-1) * sub).sum(-2)
        rows.append(torch.cat([upd, row_full[..., i:]], dim=-1))
    return torch.stack(rows, dim=-2) + eye                # [H, n, c, c]


def level_masks(c, dtype, device):
    """M_s for s = 1, 2, 4, ..., c/2: 1.0 where (row, col) is in the lower-left s x s sub-block of a 2s diagonal block."""
    masks, s = [], 1
    while s < c:
        masks.append(torch.tensor([[1.0 if (i // (2 * s) == j // (2 * s) and (i // s) % 2 == 1 and (j // s) % 2 == 0)
                                    else 0.0 for j in range(c)] for i in range(c)], dtype=dtype, device=device))
        s *= 2
    return masks


def inverse_doubling(attn, eye):
    """(I - A)^-1 by recursive doubling of the block forward substitution (r12 module docstring)."""
    masks = level_masks(attn.shape[-1], attn.dtype, attn.device)
    x = eye + attn * masks[0]
    for m in masks[1:]:
        x = x + x @ (attn * m) @ x
    return x


INVERSES = {"loop": inverse_loop, "R64": inverse_doubling}


# ---- direct decay sums ------------------------------------------------------------------------------------------------
def window_matrix(c, dtype, device):
    """W [c, c*c]: W[t, i*c + j] = 1 if j < t <= i (python data -> a graph constant)."""
    return torch.tensor([[1.0 if (j < t <= i) else 0.0 for i in range(c) for j in range(c)] for t in range(c)],
                        dtype=dtype, device=device)


def suffix_matrix(c, dtype, device):
    """S [c, c]: S[t, i] = 1 if t > i."""
    return torch.tensor([[1.0 if t > i else 0.0 for i in range(c)] for t in range(c)], dtype=dtype, device=device)


def _cexp(x):
    return torch.clamp_min(x, EXP_FLOOR).exp()


def _texp(x):
    return x.exp()


def make_kernel(inverse="R64", expclamp=True, decay="cumsum", dtype=torch.float32, vscale_log2=0):
    """kev_qwen35_patch._rank4_chunk_gated_delta_rule (batch 1, rank <= 4) with the chosen inverse / exp / decay.
    dtype = the internal compute dtype (float32 in every export; float64 only for the selftest's reference)."""
    inv = INVERSES[inverse]
    ex = _cexp if expclamp else _texp
    assert decay in ("cumsum", "direct"), decay

    def kernel(query, key, value, g, beta, chunk_size=64, initial_state=None, output_final_state=False,
               use_qk_l2norm_in_kernel=False, _layer=-1, **kwargs):
        from transformers.models.qwen3_5 import modeling_qwen3_5
        cap = lambda name, t: _cap(_layer, name, t)
        initial_dtype = query.dtype
        cap("k_in_q", query)
        cap("k_in_k", key)
        cap("k_in_v", value)
        if vscale_log2:      # vs<k>: the delta rule is linear in v, so out(2^k v) = 2^k out(v) exactly (power of 2)
            value = value * float(2 ** vscale_log2)
        cap("k_in_g", g)
        cap("k_in_beta", beta)
        if use_qk_l2norm_in_kernel:
            query = modeling_qwen3_5.l2norm(query, dim=-1, eps=1e-6)
            key = modeling_qwen3_5.l2norm(key, dim=-1, eps=1e-6)
            cap("l2_q", query)
            cap("l2_k", key)
        q, k, v = [x.squeeze(0).transpose(0, 1).contiguous().to(dtype) for x in (query, key, value)]
        b_t, g_t = [x.squeeze(0).transpose(0, 1).contiguous().to(dtype) for x in (beta, g)]
        num_heads, seq_len, k_dim = k.shape
        v_dim = v.shape[-1]
        pad_size = (chunk_size - seq_len % chunk_size) % chunk_size
        if pad_size:   # tail pad as concat (the shipped kernel's GPU-safe form)
            zero_qk = torch.zeros(num_heads, pad_size, k_dim, dtype=q.dtype, device=q.device)
            zero_v = torch.zeros(num_heads, pad_size, v_dim, dtype=v.dtype, device=v.device)
            zero_t = torch.zeros(num_heads, pad_size, dtype=g_t.dtype, device=g_t.device)
            q = torch.cat([q, zero_qk], dim=1)
            k = torch.cat([k, zero_qk], dim=1)
            v = torch.cat([v, zero_v], dim=1)
            b_t = torch.cat([b_t, zero_t], dim=-1)
            g_t = torch.cat([g_t, zero_t], dim=-1)
        total = seq_len + pad_size
        n = total // chunk_size
        q = q * (q.shape[-1] ** -0.5)
        v_beta = v * b_t.unsqueeze(-1)
        k_beta = k * b_t.unsqueeze(-1)
        q, k, v, k_beta, v_beta = [x.reshape(num_heads, n, chunk_size, x.shape[-1]) for x in (q, k, v, k_beta, v_beta)]
        g_t = g_t.reshape(num_heads, n, chunk_size)
        eye, tril0_f, tril1_f = P._chunk_rule_constants(chunk_size, dtype, q.device)
        g_c = g_t.cumsum(dim=-1)
        cap("g_cumsum", g_c)
        if decay == "direct":
            W = window_matrix(chunk_size, dtype, q.device)
            S = suffix_matrix(chunk_size, dtype, q.device)
            diff = torch.matmul(g_t, W).reshape(num_heads, n, chunk_size, chunk_size)
            suffix = torch.matmul(g_t, S)                       # [H, n, c] = g_c[-1] - g_c[i]
            cap("decay_diff", diff)
            cap("decay_suffix", suffix)
            decay_mask = ex(diff) * tril0_f
        else:
            diff = (g_c.unsqueeze(-1) - g_c.unsqueeze(-2)) * tril0_f
            cap("decay_diff", diff)
            decay_mask = ex(diff).float() * tril0_f
            suffix = None
        cap("decay_mask", decay_mask)
        attn = -((k_beta @ k.transpose(-1, -2)) * decay_mask) * tril1_f
        cap("A_strict", attn)
        attn = inv(attn, eye)                                   # (I - A)^-1, [H, n, c, c]
        cap("T_solve", attn)
        v2 = attn @ v_beta
        cap("T_v_beta", v2)
        k_cumdecay = attn @ (k_beta * ex(g_c).unsqueeze(-1))
        cap("k_cumdecay", k_cumdecay)
        if initial_state is None:
            state = torch.zeros(num_heads, k_dim, v_dim, dtype=v2.dtype, device=v2.device)
        else:
            state = initial_state.reshape(-1, initial_state.shape[-2], initial_state.shape[-1]).to(v2)
        outs = []
        for i in range(n):
            q_i, k_i, v_i = q[:, i], k[:, i], v2[:, i]                       # [H, c, d]
            attn_i = q_i @ k_i.transpose(-1, -2) * decay_mask[:, i]
            v_prime = k_cumdecay[:, i] @ state
            v_new = v_i - v_prime
            attn_inter = (q_i * ex(g_c[:, i, :, None])) @ state
            o = attn_inter + attn_i @ v_new
            cap(f"chunk{i}_out", o)
            outs.append(o)
            if suffix is None:
                sdec = ex(g_c[:, i, -1, None] - g_c[:, i])
            else:
                sdec = ex(suffix[:, i])
            state = (state * ex(g_c[:, i, -1, None, None])
                     + (k_i * sdec[..., None]).transpose(-1, -2) @ v_new)
            cap(f"chunk{i}_state", state)
        last_state = state.unsqueeze(0) if output_final_state else None
        out = torch.stack(outs, dim=1)                        # [H, n, c, dv]
        out = out.reshape(num_heads, total, v_dim)[:, :seq_len]
        out = out.unsqueeze(0).transpose(1, 2).contiguous().to(initial_dtype)
        cap("k_out", out)
        return out, last_state

    kernel.__name__ = (f"r13_kernel_{inverse}" + ("_ec" if expclamp else "") + ("_dd" if decay == "direct" else "")
                       + (f"_vs{vscale_log2}" if vscale_log2 else ""))
    kernel.r13 = {"inverse": inverse, "expclamp": expclamp, "decay": decay, "vscale_log2": vscale_log2}
    return kernel


# ---- FC split ---------------------------------------------------------------------------------------------------------
class SplitLinear(torch.nn.Module):
    """A bias-free nn.Linear as N-slices of <= nmax outputs (CONCATENATION), each the sum of m K-slices of kslice inputs
    (pairwise ADD tree). Same math; in fp16 the Metal FULLY_CONNECTED accumulates over a shorter chain: its kernel for
    N <= 1024 splits K in 4, the one for N >= 1280 does not (a micro test on the real gate_proj input: the N-split
    and 4 K-slices reduce the relative RMS error several fold). The slices are separate constant weights (pre-sliced here), so the V2 recipe casts each one to fp16 as before."""

    def __init__(self, lin, nmax, kslice):
        super().__init__()
        W = lin.weight.data
        assert lin.bias is None
        N, Kd = W.shape
        m = max(1, Kd // kslice) if kslice else 1
        assert Kd % m == 0, (Kd, m)
        self.k = Kd // m
        self.m = m
        self.nslices = []
        self.parts = torch.nn.ModuleList()
        for j in range(0, N, nmax):
            row = torch.nn.ModuleList()
            for i in range(m):
                sub = torch.nn.Linear(self.k, min(nmax, N - j), bias=False)
                sub.weight.data = W[j:j + nmax, i * self.k:(i + 1) * self.k].clone()
                row.append(sub)
            self.parts.append(row)
        self.in_features, self.out_features = Kd, N
        self.orig = [lin]                      # a list: not a registered submodule (no duplicate weights in the export)

    def forward(self, x):
        ys = []
        for row in self.parts:
            ps = [sub(x[..., i * self.k:(i + 1) * self.k]) if self.m > 1 else sub(x) for i, sub in enumerate(row)]
            while len(ps) > 1:
                ps = [ps[i] + ps[i + 1] if i + 1 < len(ps) else ps[i] for i in range(0, len(ps), 2)]
            ys.append(ps[0])
        return torch.cat(ys, dim=-1) if len(ys) > 1 else ys[0]


class BmmLinear(torch.nn.Module):
    """A bias-free nn.Linear as one BATCH_MATMUL with a constant right operand + a SUM over the K-slices:
        y = sum_i x[:, K_i] W[:, K_i]^T,  x [1, L, K] -> [L, m, K/m] -> [m, L, K/m] @ Wb [m, K/m, N] -> sum over m
    = the K-split of SplitLinear in three to four ops instead of m FCs and m - 1 ADDs (the same error as the 16-FC
    split in a Metal fp16 micro test on the real gate_proj input). Wb is a buffer (a graph constant);
    r13_export.py casts it to fp16 with the V2 recipe's FLOAT_CASTING, which it registers for BATCH_MATMUL in its own
    process (the scope regex "BmmLinear" keeps the kernel's non-constant BATCH_MATMULs untouched)."""

    def __init__(self, lin, chain):
        super().__init__()
        W = lin.weight.data
        assert lin.bias is None
        N, Kd = W.shape
        self.m = max(1, Kd // chain)
        assert Kd % self.m == 0, (Kd, self.m)
        self.k = Kd // self.m
        self.register_buffer("Wb", W.t().reshape(self.m, self.k, N).clone().contiguous())
        self.in_features, self.out_features = Kd, N
        self.orig = [lin]

    def forward(self, x):
        b, L, Kd = x.shape
        assert b == 1
        xs = x.reshape(L, self.m, self.k).transpose(0, 1)          # [m, L, K/m]
        return torch.matmul(xs, self.Wb).sum(0, keepdim=True)        # [1, L, N]


def bmm_linears(model, chain):
    count = 0
    for layer in model.layers:
        for parent in list(layer.modules()):
            for name, child in list(parent.named_children()):
                if type(child) is torch.nn.Linear and child.bias is None:
                    setattr(parent, name, BmmLinear(child, chain))
                    count += 1
    return count


def split_linears(model, nmax, kslice):
    """Replace every bias-free nn.Linear of the text model's decoder layers by a SplitLinear; returns the count."""
    count = 0
    for layer in model.layers:
        for parent in list(layer.modules()):
            for name, child in list(parent.named_children()):
                if type(child) is torch.nn.Linear and child.bias is None:
                    setattr(parent, name, SplitLinear(child, nmax, kslice))
                    count += 1
    return count


def unsplit_linears(model):
    for layer in model.layers:
        for parent in list(layer.modules()):
            for name, child in list(parent.named_children()):
                if isinstance(child, (SplitLinear, BmmLinear)):
                    setattr(parent, name, child.orig[0])


# ---- gated norm / Newton rsqrt ----------------------------------------------------------------------------------------
GN = {"in_log2": 0, "eps_log4": 0, "newton": False, "rsplit": False}
NORM = {"newton": False, "rsplit": False}


def _two_stage(x2, op):
    """mean / sum over the last dim in two stages of <= 32 terms (rs): [*, D] -> [-1, D/c, c] -> op(-1) twice -> [*, 1].
    Same sum in another order (fp32: rounding level); an fp16 reduction then chains <= 32 terms per stage instead of D."""
    D = x2.shape[-1]
    c = 32 if D % 32 == 0 and D >= 512 else 16
    assert D % c == 0, D
    y = x2.reshape(-1, D // c, c)
    y = op(op(y, -1), -1)                     # rank 3 -> 2 -> 1
    return y.reshape(*x2.shape[:-1], 1)


def _mean(x2):
    return _two_stage(x2, lambda t, d: t.mean(d)) if NORM["rsplit"] else x2.mean(-1, keepdim=True)


def _sum(x2):
    return _two_stage(x2, lambda t, d: t.sum(d)) if NORM["rsplit"] else x2.sum(-1, keepdim=True)


def rsqrt_newton(v):
    """rsqrt(v) followed by one Newton step y <- y (1.5 - 0.5 v y^2) (v y first: no fp16 overflow for v up to 65504).
    fp32: the step moves an rsqrt that is already correctly rounded by <= 1 ulp. fp16 / HTP: an approximate RSQRT with
    relative error e leaves ~1.5 e^2."""
    y = torch.rsqrt(v)
    return y * (1.5 - 0.5 * ((v * y) * y))
_ORIG_GN = {}


def gated_norm_scaled(self, hidden_states, gate=None):
    """Qwen3_5RMSNormGated.forward for an input that the kernel already scaled by 2^v (vs<v>) and that is pre-scaled
    here by 2^j (gn<j>): x' = 2^(v+j) x, x' * rsqrt(mean(x'^2) + eps * 4^(v+j)) = x * rsqrt(mean(x^2) + eps) exactly up
    to the rsqrt rounding (powers of two; eps * 4^(v+j) is the fp32 eps scaled exactly). In fp16 the kernel output
    (a large share of its values below the fp16 normal range 6.1e-5, and almost all of its squares) leaves the
    subnormal range."""
    input_dtype = hidden_states.dtype
    h = hidden_states.to(torch.float32)
    if GN["in_log2"]:
        h = h * float(2 ** GN["in_log2"])
    variance = _mean(h.pow(2))
    eps = torch.tensor(self.variance_epsilon, dtype=torch.float32) * float(4 ** GN["eps_log4"])
    h = h * (rsqrt_newton(variance + eps) if NORM["newton"] else torch.rsqrt(variance + eps))
    h = self.weight * h.to(input_dtype)
    h = h * torch.nn.functional.silu(gate.to(torch.float32))
    return h.to(input_dtype)


_ORIG_NR = {}
_ORIG_SG = {}


def sigmoid_exp(x):
    """sg: sigmoid(x) = 1 / (1 + exp(max(-x, -80))) — EXP / ADD / DIV instead of LOGISTIC, for an NPU whose
    LOGISTIC is less accurate than its EXP (a layer-0 probe on the Galaxy S26 NPU, Qualcomm HTP). The clamp keeps the
    EXP argument out of the band where that NPU's EXP returned -inf, [-181.6, -87.6] (exp(-80) = 1.8e-35 -> the result
    is 1 in any precision); a negative x below -11 makes exp overflow fp16 to inf -> 1 / inf = 0 (the sigmoid's own
    limit). fp32: within an ulp of torch.sigmoid. Not used by the final files."""
    return 1.0 / (1.0 + torch.exp(torch.clamp_min(-x, EXP_FLOOR)))


def _sigmoid_method(self):
    return SIGMOID_IMPL[0](self)


def sigmoid_exp2(x):
    """sg2: sigmoid(x) = exp(min(x, 0)) / (1 + exp(-|x|)) with min(x, 0) = x - relu(x) and -|x| = x - 2 relu(x) (the
    softplus rewrite's form) and both EXP arguments clamped at -80: every EXP argument is <= 0, so nothing overflows
    fp16 (sg's exp(-x) does for x < -11) and nothing enters the NPU band [-181.6, -87.6]. fp32: the same value up to
    rounding (x >= 0: 1 / (1 + e^-x); x < 0: e^x / (1 + e^x)). Not used by the final files."""
    r = torch.relu(x)
    num = torch.exp(torch.clamp_min(x - r, EXP_FLOOR))
    den = 1.0 + torch.exp(torch.clamp_min(x - 2.0 * r, EXP_FLOOR))
    return num / den


SIGMOID_IMPL = [sigmoid_exp]


class _TorchShim(__import__("types").ModuleType):
    """modeling_qwen3_5's `torch` global with only `sigmoid` replaced (the attention output gate's torch.sigmoid(gate));
    torch._decomp's SiLU decomposition (x * torch.sigmoid(x)) keeps the real torch, so the 60 SiLU stay LOGISTIC x."""

    def __getattr__(self, name):
        return getattr(torch, name)


def rmsnorm_norm_newton(self, x):
    """Qwen3_5RMSNorm._norm with the Newton-refined rsqrt (nr) and / or the two-stage mean (rs)."""
    v = _mean(x.pow(2)) + self.eps
    return x * (rsqrt_newton(v) if NORM["newton"] else torch.rsqrt(v))


def l2norm_newton(x, dim=-1, eps=1e-6):
    """modeling_qwen3_5.l2norm with the Newton-refined rsqrt (nr) and / or the two-stage sum (rs)."""
    assert dim == -1
    v = _sum(x * x) + eps
    return x * (rsqrt_newton(v) if NORM["newton"] else torch.rsqrt(v))


# ---- apply / reset ----------------------------------------------------------------------------------------------------
KNOWN = ("loop", "R64", "sp", "ec", "dd", "gn", "nr", "rs", "sg", "sga", "sg2")   # + fc<n>k<k>, vs<k>, gs<j>
STATE = {"spec": None}


def parse(spec):
    """tokens: KNOWN, plus fc<nmax>k<kslice> (e.g. fc1024k256 = N-slices <= 1024, K-slices of 256 inputs)."""
    import re
    toks = [t for t in spec.split("+") if t]
    fc = [t for t in toks if re.fullmatch(r"fc\d+k\d+", t)]
    vs = [int(t[2:]) for t in toks if re.fullmatch(r"vs\d+", t)]
    bm = [int(t[3:]) for t in toks if re.fullmatch(r"bmc\d+", t)]
    gs = [int(t[2:]) for t in toks if re.fullmatch(r"gs\d+", t)]
    bad = [t for t in toks if t not in KNOWN and t not in fc and not re.fullmatch(r"(vs|gs|bmc)\d+", t)]
    assert len(bm) <= 1 and not (bm and fc), "one FC rewrite: fc<n>k<k> or bmc<chain>"
    assert not bad, f"unknown tokens {bad} (known {KNOWN}, fc<n>k<k>)"
    assert len(fc) <= 1
    inverse = "R64" if "R64" in toks else "loop"
    assert not ("R64" in toks and "loop" in toks)
    fcs = None
    if fc:
        mt = re.fullmatch(r"fc(\d+)k(\d+)", fc[0])
        fcs = (int(mt.group(1)), int(mt.group(2)))
    assert len(vs) <= 1 and len(gs) <= 1
    assert not ("gn" in toks and (vs or gs)), "gn = r11_fp16_safe's gated norm; vs / gs use this file's"
    return {"inverse": inverse, "softplus": "sp" in toks, "expclamp": "ec" in toks,
            "decay": "direct" if "dd" in toks else "cumsum", "gnorm": "gn" in toks, "fc_split": fcs,
            "vscale_log2": vs[0] if vs else 0, "gn_in_log2": gs[0] if gs else 0, "newton": "nr" in toks,
            "rsplit": "rs" in toks, "sigmoid_exp": "sg" in toks, "sigmoid_exp_all": "sga" in toks,
            "sigmoid_exp2": "sg2" in toks,
            "bmm_chain": bm[0] if bm else 0,
            "tokens": toks}


def _gdn_modules(model):
    return [m for m in model.modules() if isinstance(m, P.PatchedQwen3_5GatedDeltaNet)]


class _Bound:
    """A kernel bound to one GatedDeltaNet layer (passes _layer for the capture; same call signature otherwise)."""

    def __init__(self, kern, layer):
        self.kern, self.layer = kern, layer
        self.__name__ = kern.__name__
        self.r13 = kern.r13

    def __call__(self, *args, **kwargs):
        return self.kern(*args, _layer=self.layer, **kwargs)


def apply(model, spec):
    """Install the form: the kernel on every PatchedQwen3_5GatedDeltaNet; sp / gn through r11_fp16_safe (softplus = the
    guarded forward's F, gnorm = Qwen3_5RMSNormGated.forward with the 2^7 pre-scale); fc = SplitLinear swaps."""
    import r11_fp16_safe as S11
    assert STATE["spec"] is None, f"form {STATE['spec']} already applied: reset first"
    f = parse(spec)
    mods = _gdn_modules(model)
    assert mods and all(m.chunk_gated_delta_rule is P._rank4_chunk_gated_delta_rule for m in mods), \
        "a GatedDeltaNet does not hold the shipped kernel (reset other rewrites first)"
    kern = make_kernel(f["inverse"], f["expclamp"], f["decay"], vscale_log2=f["vscale_log2"])
    for m in mods:
        m.chunk_gated_delta_rule = _Bound(kern, m.layer_idx)
    if f["softplus"]:
        S11.apply("softplus")
    if f["gnorm"]:
        S11.apply("gnorm")
    from transformers.models.qwen3_5 import modeling_qwen3_5 as M
    if f["vscale_log2"] or f["gn_in_log2"] or f["newton"] or f["rsplit"]:
        assert not f["gnorm"], "nr / rs / vs / gs replace the gated norm: do not combine with r11_fp16_safe's gn"
        _ORIG_GN.setdefault("forward", M.Qwen3_5RMSNormGated.forward)
        GN["in_log2"] = f["gn_in_log2"]
        GN["eps_log4"] = f["vscale_log2"] + f["gn_in_log2"]
        GN["newton"] = f["newton"]
        M.Qwen3_5RMSNormGated.forward = gated_norm_scaled
    NORM["newton"], NORM["rsplit"] = f["newton"], f["rsplit"]
    if f["newton"] or f["rsplit"]:
        _ORIG_NR.setdefault("_norm", M.Qwen3_5RMSNorm._norm)
        _ORIG_NR.setdefault("l2norm", M.l2norm)
        M.Qwen3_5RMSNorm._norm = rmsnorm_norm_newton
        M.l2norm = l2norm_newton
    assert sum((f["sigmoid_exp"], f["sigmoid_exp_all"], f["sigmoid_exp2"])) <= 1
    SIGMOID_IMPL[0] = sigmoid_exp2 if f["sigmoid_exp2"] else sigmoid_exp
    if f["sigmoid_exp"] or f["sigmoid_exp2"]:      # sg: the 24 standalone sigmoids only — beta = b.sigmoid() (Tensor method, guarded GDN
        # forward) and torch.sigmoid(gate) (attention, through modeling_qwen3_5's torch global); the converter's SiLU
        # decomposition calls torch.sigmoid from torch._decomp, which keeps the real torch (60 SiLU stay LOGISTIC x)
        _ORIG_SG.setdefault("method", torch.Tensor.sigmoid)
        _ORIG_SG.setdefault("mod_torch", M.torch)
        shim = _TorchShim("torch_r13_sg")
        shim.sigmoid = SIGMOID_IMPL[0]
        torch.Tensor.sigmoid = _sigmoid_method
        M.torch = shim
    if f["sigmoid_exp_all"]:  # sga (first sg version): torch.sigmoid itself patched = all 84 LOGISTIC incl. the SiLU
        _ORIG_SG.setdefault("fn", torch.sigmoid)
        _ORIG_SG.setdefault("method", torch.Tensor.sigmoid)
        torch.sigmoid = sigmoid_exp
        torch.Tensor.sigmoid = _sigmoid_method
    nsplit = split_linears(model, *f["fc_split"]) if f["fc_split"] else 0
    if f["bmm_chain"]:
        nsplit = bmm_linears(model, f["bmm_chain"])
    STATE["spec"] = spec
    return {"spec": spec, **f, "kernel": kern.__name__, "gdn_modules": len(mods), "split_linears": nsplit,
            "r11_applied": list(S11.APPLIED)}


def reset(model):
    import r11_fp16_safe as S11
    for m in _gdn_modules(model):
        if isinstance(m.chunk_gated_delta_rule, _Bound):
            m.chunk_gated_delta_rule = P._rank4_chunk_gated_delta_rule
    unsplit_linears(model)
    S11.reset()
    from transformers.models.qwen3_5 import modeling_qwen3_5 as M
    if "forward" in _ORIG_GN:
        M.Qwen3_5RMSNormGated.forward = _ORIG_GN["forward"]
        GN["in_log2"] = GN["eps_log4"] = 0
        GN["newton"] = False
    NORM["newton"] = NORM["rsplit"] = False
    if "fn" in _ORIG_SG:
        torch.sigmoid = _ORIG_SG.pop("fn")
    if "method" in _ORIG_SG:
        torch.Tensor.sigmoid = _ORIG_SG.pop("method")
    if "mod_torch" in _ORIG_SG:
        M.torch = _ORIG_SG.pop("mod_torch")
    SIGMOID_IMPL[0] = sigmoid_exp
    if "_norm" in _ORIG_NR:
        M.Qwen3_5RMSNorm._norm = _ORIG_NR["_norm"]
        M.l2norm = _ORIG_NR["l2norm"]
    STATE["spec"] = None


def selftest():
    """Random kernel inputs: the window / suffix sums equal the cumsum differences (fp64), and every fp32 form vs the
    kernel computed in float64 (decays down to -6 / token, cumsum to -390; the float64 reference has no exp clamp)."""
    import json
    torch.manual_seed(0)
    c = 64
    g = -torch.rand(16, 2, c, dtype=torch.float64) * 6.0
    gc = g.cumsum(-1)
    W = window_matrix(c, torch.float64, "cpu")
    S = suffix_matrix(c, torch.float64, "cpu")
    diff_ref = (gc.unsqueeze(-1) - gc.unsqueeze(-2)) * torch.tril(torch.ones(c, c, dtype=torch.float64))
    diff_dd = (g @ W).reshape(16, 2, c, c)
    suf_ref = gc[..., -1:] - gc
    suf_dd = g @ S
    res = {"window_vs_cumsum_diff_fp64_max_abs": float((diff_dd - diff_ref).abs().max()),
           "suffix_vs_cumsum_fp64_max_abs": float((suf_dd - suf_ref).abs().max())}
    # fp32 kernels vs an fp64 kernel on the same inputs
    H, L, D = 4, 130, 16
    q = torch.randn(1, L, H, D, dtype=torch.float64)
    k = torch.randn(1, L, H, D, dtype=torch.float64)
    v = torch.randn(1, L, H, D, dtype=torch.float64)
    gg = -torch.rand(1, L, H, dtype=torch.float64) * 6.0
    b = torch.rand(1, L, H, dtype=torch.float64)
    ref, _ = make_kernel("loop", False, "direct", dtype=torch.float64)(q, k, v, gg, b, use_qk_l2norm_in_kernel=True)
    for inv in ("loop", "R64"):
        for dec in ("cumsum", "direct"):
            out, _ = make_kernel(inv, True, dec)(*(x.float() for x in (q, k, v, gg, b)), use_qk_l2norm_in_kernel=True)
            res[f"{inv}_ec_{dec}_fp32_vs_fp64_max_abs"] = float((out.double() - ref).abs().max())
    # fp16 emulation of the decay alone on real-looking magnitudes: abs error of the decay mask, cumsum vs direct
    gh = (-torch.rand(16, 2, c) * 4.0).half()
    gc_h = gh.float().cumsum(-1).half()                  # fp16 storage of the prefix sums
    diff_h = (gc_h.unsqueeze(-1).float() - gc_h.unsqueeze(-2).float()).half()
    dd_h = (gh.float() @ window_matrix(c, torch.float32, "cpu")).half().reshape(16, 2, c, c)
    true = (gh.double() @ window_matrix(c, torch.float64, "cpu")).reshape(16, 2, c, c)
    tri = torch.tril(torch.ones(c, c, dtype=torch.bool))
    res["fp16_storage_decay_mask_abs_err_cumsum_diff"] = float((diff_h.double().exp() - true.exp())[..., tri].abs().max())
    res["fp16_storage_decay_mask_abs_err_direct"] = float((dd_h.double().exp() - true.exp())[..., tri].abs().max())
    print(json.dumps(res, indent=1))
    return res


if __name__ == "__main__":
    selftest()