"""PyTorch implementation of the Weevil recurrent GDN2 decoder.

Prelude runs once. The four-layer core then repeats ``core_passes`` times,
adding ``input_scale * prelude`` on every pass. Coda runs once. Next-token
logits are tied through the factorized embedding.

Each mixer keeps a convolution window and a Gated DeltaNet-2 state. Prefill
runs the chunk schedule over the prompt and stores that state. Later tokens
advance the same state one step at a time. Depth 16 is the evaluated setting.
"""

from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn.functional as F
from torch import nn
from transformers import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_weevil_recurrent import WeevilRecurrentConfig


def _dtype(name: str) -> torch.dtype:
    return torch.bfloat16 if name == "bfloat16" else torch.float32


def apply_linear(x: torch.Tensor, layer: nn.Linear) -> torch.Tensor:
    weight = layer.weight.to(dtype=x.dtype)
    bias = None if layer.bias is None else layer.bias.to(dtype=x.dtype)
    return F.linear(x, weight, bias)


def rmsnorm(x: torch.Tensor, eps: float) -> torch.Tensor:
    x32 = x.float()
    inv = torch.rsqrt(x32.square().mean(dim=-1, keepdim=True) + eps)
    return (x32 * inv).to(dtype=x.dtype)


class WeevilRecurrentState:
    """Convolution windows and GDN2 states for one batch of sequences.

    ``recurrent_s`` has a separate state for every core pass. Passes share
    weights, but each pass sees a different residual and keeps its own time
    state. ``seq_length`` is how many tokens have already been folded in.
    """

    def __init__(
        self,
        prelude_s: torch.Tensor,
        prelude_hist: torch.Tensor,
        recurrent_s: torch.Tensor,
        recurrent_hist: torch.Tensor,
        coda_s: torch.Tensor,
        coda_hist: torch.Tensor,
        seq_length: int,
        core_passes: int,
    ) -> None:
        self.prelude_s = prelude_s
        self.prelude_hist = prelude_hist
        self.recurrent_s = recurrent_s
        self.recurrent_hist = recurrent_hist
        self.coda_s = coda_s
        self.coda_hist = coda_hist
        self.seq_length = seq_length
        self.core_passes = core_passes

    def get_seq_length(self, layer_idx: int = 0) -> int:
        return self.seq_length

    def get_max_length(self) -> int:
        return -1

    @property
    def batch_size(self) -> int:
        return self.prelude_s.shape[1]

    def clone(self) -> "WeevilRecurrentState":
        """Independent copy. Token steps write the state in place, so a fork has to own its tensors."""

        return WeevilRecurrentState(
            self.prelude_s.clone(),
            self.prelude_hist.clone(),
            self.recurrent_s.clone(),
            self.recurrent_hist.clone(),
            self.coda_s.clone(),
            self.coda_hist.clone(),
            self.seq_length,
            self.core_passes,
        )

    def repeat(self, count: int) -> "WeevilRecurrentState":
        """``count`` independent copies of each sequence, in batch order."""

        if count < 1:
            raise ValueError("count must be positive")

        def copies(tensor: torch.Tensor, dim: int) -> torch.Tensor:
            return tensor.repeat_interleave(count, dim=dim)

        return WeevilRecurrentState(
            copies(self.prelude_s, 1),
            copies(self.prelude_hist, 2),
            copies(self.recurrent_s, 2),
            copies(self.recurrent_hist, 3),
            copies(self.coda_s, 1),
            copies(self.coda_hist, 2),
            self.seq_length,
            self.core_passes,
        )


@dataclass
class WeevilCausalLMOutput(CausalLMOutputWithPast):
    """Causal LM output. ``recurrent_state`` is the carried GDN2 state."""

    recurrent_state: WeevilRecurrentState | None = None


class CausalConv(nn.Module):
    def __init__(self, channels: int, kernel: int) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.empty(channels, kernel))

    def forward(
        self, x: torch.Tensor, hist: torch.Tensor | None
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Causal depthwise conv. ``weight[:, -1]`` multiplies the current token.

        ``hist`` is the previous ``kernel - 1`` inputs, or ``None`` for zeros.
        Returns the activated outputs and the window to carry forward.
        """

        channels, kernel = self.weight.shape
        lookback = kernel - 1
        if hist is None:
            hist = x.new_zeros(x.shape[0], lookback, channels)
        full = x if lookback == 0 else torch.cat((hist, x), dim=1)
        rhs = self.weight.to(dtype=x.dtype).view(channels, 1, kernel)
        y = F.conv1d(full.transpose(1, 2), rhs, groups=channels)
        new_hist = full[:, -lookback:] if lookback else full[:, :0]
        return F.silu(y.transpose(1, 2)), new_hist


def _gdn_prepare(q, k, v, g, b, w, valid, chunk_size):
    """Mask, normalize, and clamp. Shared by the chunk prefill and the token step."""

    if valid is not None:
        mask = valid.to(dtype=q.dtype)[..., None, None]
        q, k, v, g, b, w = q * mask, k * mask, v * mask, g * mask, b * mask, w * mask
    q = q.float()
    k = k.float()
    q = q * torch.rsqrt(q.square().sum(dim=-1, keepdim=True) + 1e-6)
    k = k * torch.rsqrt(k.square().sum(dim=-1, keepdim=True) + 1e-6)
    q = q * (q.shape[-1] ** -0.5)
    v = v.float()
    g = torch.clamp(g.float(), min=-60.0 / chunk_size)
    return q, k, v, g, b.float(), w.float()


def _forward_substitute(lower: torch.Tensor) -> torch.Tensor:
    size = lower.shape[-1]
    solved = lower
    index = torch.arange(size, device=lower.device)
    for row_index in range(1, size):
        mask = (index < row_index).to(dtype=lower.dtype)
        row = solved[..., row_index, :]
        masked_row = row * mask
        masked_block = solved * mask[..., :, None] * mask[..., None, :]
        updated = row + torch.matmul(masked_row.unsqueeze(-2), masked_block).squeeze(-2)
        updated = torch.where(mask.bool(), updated, row)
        solved = solved.clone()
        solved[..., row_index, :] = updated
    eye = torch.eye(size, dtype=solved.dtype, device=solved.device)
    return solved + eye


def gdn2_chunk(q, k, v, g, b, w, chunk_size: int) -> tuple[torch.Tensor, torch.Tensor]:
    """Chunk prefill from a zero state. Inputs are already prepared, ``[B, T, H, D]``."""

    def heads_first(tensor: torch.Tensor) -> torch.Tensor:
        return tensor.permute(0, 2, 1, 3)

    q, k, v, g, b, w = (heads_first(tensor) for tensor in (q, k, v, g, b, w))
    batch, heads, length, k_dim = k.shape
    v_dim = v.shape[-1]
    pad = (chunk_size - length % chunk_size) % chunk_size
    if pad:
        padding = (0, 0, 0, pad)
        q, k, v, g, b, w = (F.pad(tensor, padding) for tensor in (q, k, v, g, b, w))
    total = length + pad
    n_chunks = total // chunk_size

    def as_chunks(tensor: torch.Tensor) -> torch.Tensor:
        return tensor.reshape(batch, heads, n_chunks, chunk_size, tensor.shape[-1])

    q, k, v, g, b, w = (as_chunks(tensor) for tensor in (q, k, v, g, b, w))
    g = torch.cumsum(g, dim=-2)
    decay = torch.exp(g)
    center = g[..., -1:, :] * 0.5
    decay_centered = torch.exp(g - center)
    inv_decay_centered = torch.exp(center - g)
    erase = decay * b * k
    erase_centered = decay_centered * b * k
    key_inv = k * inv_decay_centered
    write = w * v
    interaction = -torch.matmul(erase_centered, key_inv.transpose(-1, -2))
    strict = torch.tril(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=-1)
    lower = torch.where(strict, interaction, torch.zeros_like(interaction))
    transition = _forward_substitute(lower)
    write = torch.matmul(transition, write)
    k_cumdecay = torch.matmul(transition, erase)
    query_decay = q * decay
    query_decay_centered = q * decay_centered
    g_last = g[..., -1:, :]

    state = q.new_zeros(batch, heads, k_dim, v_dim)
    outputs = []
    causal = torch.tril(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device))
    for index in range(n_chunks):
        q_decay = query_decay[:, :, index]
        q_centered = query_decay_centered[:, :, index]
        key_inv_c = key_inv[:, :, index]
        write_c = write[:, :, index]
        k_cum_c = k_cumdecay[:, :, index]
        key_c = k[:, :, index]
        g_c = g[:, :, index]
        g_end = g_last[:, :, index]
        attn = torch.matmul(q_centered, key_inv_c.transpose(-1, -2))
        attn = torch.where(causal, attn, torch.zeros_like(attn))
        residual = write_c - torch.matmul(k_cum_c, state)
        out = torch.matmul(q_decay, state) + torch.matmul(attn, residual)
        key_bar = key_c * torch.exp(g_end - g_c)
        decay_state = torch.exp(g_end.squeeze(-2)).unsqueeze(-1)
        state = state * decay_state + torch.matmul(key_bar.transpose(-1, -2), residual)
        outputs.append(out)
    stacked = torch.stack(outputs, dim=2).reshape(batch, heads, total, v_dim)[:, :, :length]
    return stacked.permute(0, 2, 1, 3), state


def _gdn2_token_pytorch(q, k, v, g, b, w, state: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Reference one-step update. ``q``..``w`` are ``[B, H, D]``, state is ``[B, H, D, D]``."""

    alpha = g.exp()
    decayed = state * alpha.unsqueeze(-1)
    erase = torch.einsum("bhd,bhde->bhe", b * k, decayed)
    updated = decayed - torch.einsum("bhd,bhe->bhde", k, erase)
    updated = updated + torch.einsum("bhd,bhe->bhde", k, w * v)
    out = torch.einsum("bhde,bhd->bhe", updated, q)
    return out, updated


try:
    import triton
    import triton.language as tl

    @triton.jit
    def _gdn2_token_kernel(
        q_ptr, k_ptr, v_ptr, g_ptr, b_ptr, w_ptr, s_ptr, o_ptr,
        stride_vb, stride_vh, stride_sb, stride_sh, stride_sd,
        D: tl.constexpr,
    ):
        """One head. Load the DxD state, apply the delta rule, store it back."""

        pid_b = tl.program_id(0)
        pid_h = tl.program_id(1)
        offs = tl.arange(0, D)
        base = pid_b * stride_vb + pid_h * stride_vh
        sbase = pid_b * stride_sb + pid_h * stride_sh
        q = tl.load(q_ptr + base + offs)
        k = tl.load(k_ptr + base + offs)
        v = tl.load(v_ptr + base + offs)
        g = tl.load(g_ptr + base + offs)
        b = tl.load(b_ptr + base + offs)
        w = tl.load(w_ptr + base + offs)
        state = tl.load(s_ptr + sbase + offs[:, None] * stride_sd + offs[None, :])
        alpha = tl.exp(g)
        decayed = state * alpha[:, None]
        erase = tl.sum(decayed * (b * k)[:, None], axis=0)
        updated = decayed + k[:, None] * ((w * v)[None, :] - erase[None, :])
        tl.store(s_ptr + sbase + offs[:, None] * stride_sd + offs[None, :], updated)
        tl.store(o_ptr + base + offs, tl.sum(updated * q[:, None], axis=0))

    @triton.jit
    def _conv_token_kernel(
        proj_ptr, hist_ptr, weight_ptr, out_ptr,
        proj_stride, hist_stride_stream, hist_stride_batch, hist_stride_tap, weight_stride,
        channels,
        BLOCK: tl.constexpr,
    ):
        """Four-tap depthwise conv for one token, then slide the window down."""

        batch = tl.program_id(0)
        stream = tl.program_id(1)
        block = tl.program_id(2)
        offs = block * BLOCK + tl.arange(0, BLOCK)
        mask = offs < channels
        hist_base = stream * hist_stride_stream + batch * hist_stride_batch
        weight_base = (stream * channels + offs) * weight_stride
        h0 = tl.load(hist_ptr + hist_base + 0 * hist_stride_tap + offs, mask=mask, other=0.0).to(tl.float32)
        h1 = tl.load(hist_ptr + hist_base + 1 * hist_stride_tap + offs, mask=mask, other=0.0).to(tl.float32)
        h2 = tl.load(hist_ptr + hist_base + 2 * hist_stride_tap + offs, mask=mask, other=0.0).to(tl.float32)
        current = tl.load(proj_ptr + batch * proj_stride + stream * channels + offs, mask=mask, other=0.0)
        w0 = tl.load(weight_ptr + weight_base + 0, mask=mask, other=0.0).to(tl.float32)
        w1 = tl.load(weight_ptr + weight_base + 1, mask=mask, other=0.0).to(tl.float32)
        w2 = tl.load(weight_ptr + weight_base + 2, mask=mask, other=0.0).to(tl.float32)
        w3 = tl.load(weight_ptr + weight_base + 3, mask=mask, other=0.0).to(tl.float32)
        acc = h0 * w0 + h1 * w1 + h2 * w2 + current.to(tl.float32) * w3
        out = (acc * tl.sigmoid(acc)).to(current.dtype)
        tl.store(out_ptr + (batch * 3 + stream) * channels + offs, out, mask=mask)
        tl.store(hist_ptr + hist_base + 0 * hist_stride_tap + offs, h1.to(current.dtype), mask=mask)
        tl.store(hist_ptr + hist_base + 1 * hist_stride_tap + offs, h2.to(current.dtype), mask=mask)
        tl.store(hist_ptr + hist_base + 2 * hist_stride_tap + offs, current, mask=mask)

except ImportError:  # pragma: no cover - CPU wheels have no Triton
    _gdn2_token_kernel = None
    _conv_token_kernel = None


def _conv_token_cuda(projected: torch.Tensor, hist: torch.Tensor, weight: torch.Tensor):
    """Slide the q/k/v windows and return the activated conv outputs, each ``[B, 1, C]``."""

    batch = projected.shape[0]
    channels = hist.shape[-1]
    out = torch.empty(batch, 3, channels, device=projected.device, dtype=projected.dtype)
    block = 128
    _conv_token_kernel[(batch, 3, (channels + block - 1) // block)](
        projected, hist, weight, out,
        projected.stride(0), hist.stride(0), hist.stride(1), hist.stride(2), weight.stride(0),
        channels, BLOCK=block,
    )
    return out[:, 0:1], out[:, 1:2], out[:, 2:3]


def _gdn2_token_cuda(q, k, v, g, b, w, state: torch.Tensor) -> torch.Tensor:
    batch, heads, dim = q.shape
    out = torch.empty(batch, heads, dim, device=q.device, dtype=torch.float32)
    _gdn2_token_kernel[batch, heads](
        q, k, v, g, b, w, state, out,
        q.stride(0), q.stride(1), state.stride(0), state.stride(1), state.stride(2),
        D=dim, num_warps=4,
    )
    return out


def gdn2_token(q, k, v, g, b, w, state: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """One time step. Inputs are prepared ``[B, 1, H, D]``. ``state`` is ``[B, H, D, D]``.

    On CUDA the state is updated in place. The returned state is that same tensor.
    """

    q, k, v, g, b, w = (tensor[:, 0] for tensor in (q, k, v, g, b, w))
    vectors = (q, k, v, g, b, w)
    inplace = (
        _gdn2_token_kernel is not None
        and state.is_cuda
        and state.is_contiguous()
        and state.dtype == torch.float32
        and all(tensor.is_contiguous() and tensor.dtype == torch.float32 for tensor in vectors)
    )
    if inplace:
        out = _gdn2_token_cuda(*vectors, state)
        return out[:, None], state
    out, updated = _gdn2_token_pytorch(*vectors, state)
    return out[:, None], updated


class Mixer(nn.Module):
    def __init__(self, config: WeevilRecurrentConfig) -> None:
        super().__init__()
        hidden = config.hidden_size
        width = config.num_heads * config.head_dim
        head = config.head_dim
        self.config = config
        self.q_proj = nn.Linear(hidden, width, bias=False)
        self.k_proj = nn.Linear(hidden, width, bias=False)
        self.v_proj = nn.Linear(hidden, width, bias=False)
        self.q_conv = CausalConv(width, config.conv_kernel)
        self.k_conv = CausalConv(width, config.conv_kernel)
        self.v_conv = CausalConv(width, config.conv_kernel)
        self.f_in = nn.Linear(hidden, head, bias=False)
        self.f_out = nn.Linear(head, width, bias=False)
        self.b_proj = nn.Linear(hidden, width, bias=False)
        self.w_proj = nn.Linear(hidden, width, bias=False)
        self.g_in = nn.Linear(hidden, head, bias=False)
        self.g_out = nn.Linear(head, width, bias=True)
        self.o_proj = nn.Linear(width, hidden, bias=False)
        self.o_norm = nn.Parameter(torch.ones(head))
        self.A_log = nn.Parameter(torch.empty(config.num_heads))
        self.dt_bias = nn.Parameter(torch.empty(width))
        self._packed_input = None
        self._packed_conv = None

    def _pack(self, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
        """One input GEMM and one depthwise conv cover q, k, v, erase, write, and the two small ins."""

        if (
            self._packed_input is not None
            and self._packed_input.dtype == dtype
            and self._packed_input.device == self.q_proj.weight.device
        ):
            return self._packed_input, self._packed_conv
        parts = (
            self.q_proj.weight,
            self.k_proj.weight,
            self.v_proj.weight,
            self.b_proj.weight,
            self.w_proj.weight,
            self.f_in.weight,
            self.g_in.weight,
        )
        self._packed_input = torch.cat([part.detach().to(dtype=dtype) for part in parts], dim=0)
        convs = (self.q_conv.weight, self.k_conv.weight, self.v_conv.weight)
        self._packed_conv = torch.cat([part.detach().to(dtype=dtype) for part in convs], dim=0)
        return self._packed_input, self._packed_conv

    def forward(
        self,
        x: torch.Tensor,
        valid: torch.Tensor | None,
        s: torch.Tensor | None,
        hist: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        config = self.config
        compute = _dtype(config.compute_dtype)
        batch, length, _ = x.shape
        heads, head = config.num_heads, config.head_dim
        width = heads * head
        x_c = x.to(dtype=compute)
        packed_w, packed_conv = self._pack(compute)
        projected = F.linear(x_c, packed_w)
        q, k, v, erase_in, write_in, decay_in, gate_in = projected.split(
            (width, width, width, width, width, head, head), dim=-1
        )
        projected_q, projected_k, projected_v = q, k, v
        channels = packed_conv.shape[0]
        lookback = packed_conv.shape[1] - 1
        fast_conv = (
            _conv_token_kernel is not None
            and length == 1
            and hist is not None
            and lookback == 3
            and projected.is_cuda
        )
        if fast_conv:
            q, k, v = _conv_token_cuda(projected, hist, packed_conv)
            new_hist = hist
        else:
            qkv = torch.cat((projected_q, projected_k, projected_v), dim=-1)
            if hist is None:
                conv_hist = qkv.new_zeros(batch, lookback, channels)
            else:
                conv_hist = torch.cat((hist[0], hist[1], hist[2]), dim=-1)
            full = torch.cat((conv_hist, qkv), dim=1)
            if length == 1:
                activated = F.silu((full * packed_conv.transpose(0, 1)).sum(dim=1, keepdim=True))
            else:
                y = F.conv1d(
                    full.transpose(1, 2),
                    packed_conv.view(channels, 1, lookback + 1),
                    groups=channels,
                )
                activated = F.silu(y.transpose(1, 2))
            q, k, v = activated.split((width, width, width), dim=-1)
            if length == 1 and hist is not None and lookback:
                for shift in range(lookback - 1):
                    hist[:, :, shift].copy_(hist[:, :, shift + 1])
                hist[0, :, -1].copy_(projected_q[:, 0])
                hist[1, :, -1].copy_(projected_k[:, 0])
                hist[2, :, -1].copy_(projected_v[:, 0])
                new_hist = hist
            else:
                new_window = full[:, -lookback:]
                q_hist, k_hist, v_hist = new_window.split((width, width, width), dim=-1)
                new_hist = torch.stack((q_hist, k_hist, v_hist), dim=0)
        decay = F.softplus(
            apply_linear(decay_in, self.f_out).float() + self.dt_bias.float()
        )
        decay = decay.reshape(batch, length, heads, head)
        rate = torch.exp(self.A_log.float())
        log_decay = -rate[None, None, :, None] * decay
        erase = torch.sigmoid(erase_in).reshape(batch, length, heads, head)
        write = torch.sigmoid(write_in).reshape(batch, length, heads, head)
        q, k, v, g, b, w = _gdn_prepare(
            q.reshape(batch, length, heads, head),
            k.reshape(batch, length, heads, head),
            v.reshape(batch, length, heads, head),
            log_decay,
            erase,
            write,
            valid,
            config.chunk_size,
        )
        if length == 1:
            if s is None:
                s = q.new_zeros(batch, heads, head, head)
            out, s_new = gdn2_token(q, k, v, g, b, w, s)
        else:
            out, s_new = gdn2_chunk(q, k, v, g, b, w, config.chunk_size)
        gate = apply_linear(gate_in, self.g_out).reshape(batch, length, heads, head)
        out32 = out.float()
        inv = torch.rsqrt(out32.square().mean(dim=-1, keepdim=True) + config.norm_eps)
        out = out32 * inv * self.o_norm * torch.sigmoid(gate.float())
        flat = out.reshape(batch, length, config.num_heads * head).to(dtype=compute)
        return apply_linear(flat, self.o_proj), s_new, new_hist


class SwiGLU(nn.Module):
    def __init__(self, config: WeevilRecurrentConfig) -> None:
        super().__init__()
        self.config = config
        self.in_proj = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)
        self.out_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        compute = _dtype(self.config.compute_dtype)
        gate, up = apply_linear(x.to(dtype=compute), self.in_proj).chunk(2, dim=-1)
        return apply_linear(F.silu(gate) * up, self.out_proj)


class Gate(nn.Module):
    def __init__(self, hidden: int) -> None:
        super().__init__()
        self.pre = nn.Parameter(torch.empty(hidden))
        self.post = nn.Parameter(torch.empty(hidden))
        self.alpha = nn.Parameter(torch.empty(2))
        self.bias = nn.Parameter(torch.empty(2))

    def forward(self, normed: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        flat = normed.float()
        pre = (flat * self.pre.float()).sum(dim=-1)
        post = (flat * self.post.float()).sum(dim=-1)
        pre = self.alpha[0].float() * pre + self.bias[0].float()
        post = self.alpha[1].float() * post + self.bias[1].float()
        return torch.sigmoid(pre), 2 * torch.sigmoid(post)


class Block(nn.Module):
    def __init__(self, config: WeevilRecurrentConfig) -> None:
        super().__init__()
        self.config = config
        self.mixer = Mixer(config)
        self.mlp = SwiGLU(config)
        self.mixer_gate = Gate(config.hidden_size)
        self.mlp_gate = Gate(config.hidden_size)

    def _read(self, residual: torch.Tensor, gate: Gate) -> tuple[torch.Tensor, torch.Tensor]:
        normed = rmsnorm(residual, self.config.norm_eps)
        pre, post = gate(normed)
        hidden = (pre.unsqueeze(-1) * normed).to(dtype=_dtype(self.config.compute_dtype))
        return hidden, post

    def forward(
        self,
        residual: torch.Tensor,
        valid: torch.Tensor | None,
        s: torch.Tensor | None,
        hist: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        hidden, post = self._read(residual, self.mixer_gate)
        mixed, s_new, hist_new = self.mixer(hidden, valid, s, hist)
        residual = residual + post.unsqueeze(-1) * mixed.float()
        hidden, post = self._read(residual, self.mlp_gate)
        residual = residual + post.unsqueeze(-1) * self.mlp(hidden).float()
        if valid is not None:
            residual = residual * valid.to(dtype=residual.dtype).unsqueeze(-1)
        return residual, s_new, hist_new


class WeevilRecurrentPreTrainedModel(PreTrainedModel):
    config_class = WeevilRecurrentConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False
    _no_split_modules = ["Block"]


class Stage(nn.Module):
    def __init__(self, config: WeevilRecurrentConfig, count: int) -> None:
        super().__init__()
        self.layers = nn.ModuleList(Block(config) for _ in range(count))

    def forward(
        self,
        hidden: torch.Tensor,
        valid: torch.Tensor | None,
        s: torch.Tensor | None,
        hist: torch.Tensor | None,
    ) -> torch.Tensor:
        for index, layer in enumerate(self.layers):
            layer_s = None if s is None else s[index]
            layer_hist = None if hist is None else hist[index]
            hidden, s_new, hist_new = layer(hidden, valid, layer_s, layer_hist)
            if s is not None:
                if s_new.data_ptr() != s[index].data_ptr():
                    s[index].copy_(s_new)
                if hist_new.data_ptr() != hist[index].data_ptr():
                    hist[index].copy_(hist_new)
        return hidden


class WeevilRecurrentModel(WeevilRecurrentPreTrainedModel):
    def __init__(self, config: WeevilRecurrentConfig) -> None:
        super().__init__(config)
        self.embed = nn.Embedding(config.vocab_size, config.embedding_dim)
        self.proj = nn.Linear(config.embedding_dim, config.hidden_size, bias=False)
        self.input_scale = nn.Parameter(torch.zeros(config.hidden_size))
        self.prelude = Stage(config, config.prelude_layers)
        self.recurrent = Stage(config, config.recurrent_layers)
        self.coda = Stage(config, config.coda_layers)
        self.post_init()

    def empty_state(self, batch: int, passes: int, device: torch.device) -> WeevilRecurrentState:
        config = self.config
        heads, dim = config.num_heads, config.head_dim
        width = heads * dim
        lookback = config.conv_kernel - 1
        compute = _dtype(config.compute_dtype)

        def states(count: int) -> torch.Tensor:
            return torch.zeros(count, batch, heads, dim, dim, device=device)

        def windows(count: int) -> torch.Tensor:
            return torch.zeros(count, 3, batch, lookback, width, device=device, dtype=compute)

        return WeevilRecurrentState(
            prelude_s=states(config.prelude_layers),
            prelude_hist=windows(config.prelude_layers),
            recurrent_s=states(passes * config.recurrent_layers).view(
                passes, config.recurrent_layers, batch, heads, dim, dim
            ),
            recurrent_hist=windows(passes * config.recurrent_layers).view(
                passes, config.recurrent_layers, 3, batch, lookback, width
            ),
            coda_s=states(config.coda_layers),
            coda_hist=windows(config.coda_layers),
            seq_length=0,
            core_passes=passes,
        )

    def embed_tokens(self, input_ids: torch.Tensor, valid: torch.Tensor | None) -> torch.Tensor:
        compute = _dtype(self.config.compute_dtype)
        tokens = self.embed(input_ids).to(dtype=compute)
        hidden = apply_linear(tokens, self.proj).float()
        if valid is not None:
            hidden = hidden * valid.to(dtype=hidden.dtype).unsqueeze(-1)
        return hidden

    def _run(
        self,
        embedded: torch.Tensor,
        valid: torch.Tensor | None,
        state: WeevilRecurrentState,
        passes: int,
    ) -> torch.Tensor:
        if embedded.shape[1] > 1 and state.seq_length != 0:
            hidden = embedded
            for step in range(embedded.shape[1]):
                step_valid = None if valid is None else valid[:, step : step + 1]
                hidden_step = self._run(hidden[:, step : step + 1], step_valid, state, passes)
                if step == 0:
                    outputs = hidden_step
                else:
                    outputs = torch.cat((outputs, hidden_step), dim=1)
            return outputs
        prelude = self.prelude(embedded, valid, state.prelude_s, state.prelude_hist)
        injected = self.input_scale.float() * prelude
        hidden = prelude
        for index in range(passes):
            hidden = self.recurrent(
                hidden + injected, valid, state.recurrent_s[index], state.recurrent_hist[index]
            )
        hidden = self.coda(hidden, valid, state.coda_s, state.coda_hist)
        hidden = rmsnorm(hidden, self.config.norm_eps).to(dtype=_dtype(self.config.compute_dtype))
        if valid is not None:
            hidden = hidden * valid.to(dtype=hidden.dtype).unsqueeze(-1)
        state.seq_length += embedded.shape[1]
        return hidden

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        core_passes: int | None = None,
        past_key_values: WeevilRecurrentState | None = None,
        use_cache: bool | None = None,
        **kwargs,
    ) -> tuple[torch.Tensor, WeevilRecurrentState | None]:
        passes = self.config.core_passes if core_passes is None else core_passes
        if passes < 1:
            raise ValueError("core_passes must be positive")
        if attention_mask is not None and attention_mask.shape[-1] != input_ids.shape[-1]:
            attention_mask = attention_mask[:, -input_ids.shape[-1] :]
        valid = None if attention_mask is None else attention_mask.to(dtype=torch.bool)
        state = past_key_values
        if state is None and use_cache:
            state = self.empty_state(input_ids.shape[0], passes, input_ids.device)
        if state is not None:
            if state.core_passes != passes:
                raise ValueError(
                    f"state was built for {state.core_passes} core passes, not {passes}"
                )
            if state.batch_size != input_ids.shape[0]:
                raise ValueError("state batch does not match input_ids")
        embedded = self.embed_tokens(input_ids, valid)
        if state is None:
            state = self.empty_state(input_ids.shape[0], passes, input_ids.device)
            hidden = self._run(embedded, valid, state, passes)
            state = None
        else:
            hidden = self._run(embedded, valid, state, passes)
        return hidden, state


class WeevilRecurrentForCausalLM(WeevilRecurrentPreTrainedModel, GenerationMixin):
    def __init__(self, config: WeevilRecurrentConfig) -> None:
        super().__init__(config)
        self.model = WeevilRecurrentModel(config)
        self.post_init()
        self.generation_config.use_cache = True

    @classmethod
    def _supports_default_dynamic_cache(cls) -> bool:
        return False

    def forward(
        self,
        input_ids: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        core_passes: int | None = None,
        past_key_values: WeevilRecurrentState | None = None,
        use_cache: bool | None = None,
        labels: torch.Tensor | None = None,
        **kwargs,
    ) -> WeevilCausalLMOutput:
        if input_ids is None:
            raise ValueError("input_ids is required")
        if use_cache is None:
            use_cache = past_key_values is not None or labels is None
        hidden, state = self.model(
            input_ids,
            attention_mask=attention_mask,
            core_passes=core_passes,
            past_key_values=past_key_values,
            use_cache=use_cache,
        )
        compute = _dtype(self.config.compute_dtype)
        low = hidden.to(dtype=compute) @ self.model.proj.weight.to(dtype=compute)
        logits = (low @ self.model.embed.weight.to(dtype=compute).T).float()
        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(shift_logits.view(-1, shift_logits.shape[-1]), shift_labels.view(-1))
        return WeevilCausalLMOutput(
            loss=loss,
            logits=logits,
            past_key_values=state,
            recurrent_state=state,
        )

    def score_continuations(
        self,
        input_ids: torch.Tensor,
        continuations: torch.Tensor | list,
        continuation_mask: torch.Tensor | None = None,
        core_passes: int | None = None,
    ) -> torch.Tensor:
        """Mean log-probability of each continuation after one prefill of ``input_ids``.

        ``continuations`` is a list of token-id sequences or a right-padded ``[choices, length]``
        tensor. A padded tensor uses ``pad_token_id`` as padding unless ``continuation_mask``
        marks the real tokens. The prompt state is copied once per choice, so the choices share
        the prefill and do not interfere.
        """

        if input_ids.ndim == 1:
            input_ids = input_ids.unsqueeze(0)
        if input_ids.shape[0] != 1:
            raise ValueError("score_continuations prefills one prompt")
        encoded, mask = _pack_continuations(
            continuations, continuation_mask, self.config.pad_token_id, input_ids.device
        )
        prefill = self(input_ids, core_passes=core_passes, use_cache=True)
        first = F.log_softmax(prefill.logits[0, -1].float(), dim=-1)
        token_scores = first.new_zeros(encoded.shape)
        token_scores[:, 0] = first[encoded[:, 0]]
        if encoded.shape[1] > 1 and bool(mask[:, 1:].any()):
            branch = prefill.recurrent_state.repeat(encoded.shape[0])
            scored = self(
                encoded[:, :-1],
                attention_mask=mask[:, :-1].to(dtype=torch.long),
                past_key_values=branch,
                core_passes=core_passes,
                use_cache=True,
            )
            taken = F.log_softmax(scored.logits.float(), dim=-1)
            token_scores[:, 1:] = taken.gather(-1, encoded[:, 1:].unsqueeze(-1)).squeeze(-1)
        return (token_scores * mask).sum(dim=-1) / mask.sum(dim=-1)


def _pack_continuations(continuations, mask, pad_id: int, device: torch.device):
    if isinstance(continuations, torch.Tensor):
        if continuations.ndim == 1:
            continuations = continuations.unsqueeze(0)
        if continuations.ndim != 2:
            raise ValueError("continuations tensor must have shape [choices, length]")
        encoded = continuations.to(device=device, dtype=torch.long)
        if encoded.shape[0] == 0 or encoded.shape[1] == 0:
            raise ValueError("each continuation must contain at least one token")
        if mask is None:
            mask = encoded.ne(pad_id)
        else:
            mask = mask.to(device=device, dtype=torch.bool)
            if mask.shape != encoded.shape:
                raise ValueError("continuation_mask must match continuations")
        if not bool(mask[:, 0].all()):
            raise ValueError("continuations must be left-aligned")
        return encoded, mask
    rows = [list(row) for row in continuations]
    if not rows or any(len(row) == 0 for row in rows):
        raise ValueError("each continuation must contain at least one token")
    width = max(len(row) for row in rows)
    encoded = torch.full((len(rows), width), pad_id, dtype=torch.long, device=device)
    valid = torch.zeros(len(rows), width, dtype=torch.bool, device=device)
    for index, row in enumerate(rows):
        encoded[index, : len(row)] = torch.tensor(row, dtype=torch.long, device=device)
        valid[index, : len(row)] = True
    return encoded, valid


WeevilRecurrentConfig.register_for_auto_class()
WeevilRecurrentModel.register_for_auto_class("AutoModel")
WeevilRecurrentForCausalLM.register_for_auto_class("AutoModelForCausalLM")