"""ÇABA v0-core: recurrent three-bank causal LM for Hugging Face Transformers.

The recurrent update here is an explicit delta-rule prototype, not a claim of
bit-for-bit equivalence with the official Gated DeltaNet-2 implementation.
Training on TPU uses PyTorch/XLA's scan operator; cached inference keeps a
constant-size recurrent state plus the unfinished fixed-size promotion window.
"""

from dataclasses import dataclass
import math
from typing import Any, Dict, List, Optional, Tuple

import torch
import transformers as _transformers
from torch import nn
from torch.nn import functional as F
from torch.utils.checkpoint import checkpoint as _torch_checkpoint
try:
    from torch_xla.utils.checkpoint import checkpoint as _xla_checkpoint
except ImportError:  # Keep non-XLA inference and training usable.
    _xla_checkpoint = None
from transformers import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_caba import CabaConfig

try:
    from torch_xla.experimental.scan import scan as _xla_scan
except Exception:  # Transformers inference does not require PyTorch/XLA.
    _xla_scan = None


@dataclass
class CabaCausalLMOutput(CausalLMOutputWithPast):
    """Causal LM output with small scalar diagnostics for training logs."""

    memory_metrics: Optional[Dict[str, torch.Tensor]] = None


def _inverse_softplus(x: float) -> float:
    return math.log(math.expm1(x))


def _scan(fn, carry: torch.Tensor, xs: torch.Tensor):
    """Run a scan on TPU, with an eager reference path for CPU smoke use."""
    if carry.device.type == "xla":
        if _xla_scan is None:
            raise RuntimeError(
                "TPU training needs torch_xla.experimental.scan. Install a matched "
                "PyTorch/PyTorch-XLA pair from the first notebook cell and restart runtime."
            )
        return _xla_scan(fn, carry, xs)

    outputs = []
    for index in range(xs.shape[0]):
        carry, output = fn(carry, xs[index])
        outputs.append(output)
    if not outputs:
        raise ValueError("scan input must contain at least one step")
    return carry, torch.stack(outputs, dim=0)


def _read(memory: torch.Tensor, query: torch.Tensor) -> torch.Tensor:
    # memory: [batch, heads, key_dim, value_dim]; query: [batch, heads, key_dim]
    return torch.einsum("bhkv,bhk->bhv", memory, query.float())


def _predict_error(memory: torch.Tensor, key: torch.Tensor, value: torch.Tensor):
    prediction = _read(memory, key)
    error = value.float() - prediction
    score = (
        torch.linalg.vector_norm(error, dim=-1)
        / (torch.linalg.vector_norm(value.float(), dim=-1) + 1e-6)
    )
    return error, score.detach()


def _outer(error: torch.Tensor, key: torch.Tensor) -> torch.Tensor:
    return torch.einsum("bhv,bhk->bhkv", error, key.float())


def _mark_data_sharding(tensor: torch.Tensor, mesh, spec):
    if mesh is None:
        return tensor
    import torch_xla.distributed.spmd as xs
    return xs.mark_sharding(tensor, mesh, spec)


class CabaRMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        scale = torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return (x.float() * scale * self.weight.float()).to(x.dtype)


class CabaBlock(nn.Module):
    def __init__(self, config: CabaConfig):
        super().__init__()
        self.hidden = config.hidden_size
        self.heads = config.num_heads
        self.head_dim = config.head_dim
        self.medium_clock = config.medium_clock
        self.slow_clock = config.slow_clock
        self.slow_span = self.medium_clock * self.slow_clock
        self.conv_kernel = config.conv_kernel
        self.eps = config.rms_norm_eps
        self.data_mesh = None

        self.input_norm = CabaRMSNorm(self.hidden, self.eps)
        self.local_conv = nn.Conv1d(
            self.hidden,
            self.hidden,
            kernel_size=self.conv_kernel,
            groups=self.hidden,
            bias=False,
        )
        self.qkv = nn.Linear(self.hidden, 3 * self.hidden, bias=False)
        self.fast_write = nn.Linear(self.hidden, self.heads, bias=False)
        self.merge_gate = nn.Linear(self.hidden, 3 * self.heads, bias=False)
        self.out_proj = nn.Linear(self.hidden, self.hidden, bias=False)

        self.fast_score_alpha = nn.Parameter(
            torch.full((self.heads,), _inverse_softplus(1.0))
        )
        self.fast_score_tau = nn.Parameter(torch.ones(self.heads))
        self.medium_score_alpha = nn.Parameter(
            torch.full((self.heads,), _inverse_softplus(1.0))
        )
        self.medium_score_tau = nn.Parameter(torch.ones(self.heads))
        self.medium_admit_threshold = nn.Parameter(torch.full((self.heads,), 2.0))
        self.slow_admit_threshold = nn.Parameter(torch.full((self.heads,), 3.0))

        # Positive increments structurally preserve T_slow > T_medium > T_fast.
        self.raw_fast_timescale = nn.Parameter(
            torch.full((self.heads,), _inverse_softplus(7.0))
        )
        self.raw_medium_gap = nn.Parameter(
            torch.full((self.heads,), _inverse_softplus(56.0))
        )
        self.raw_slow_gap = nn.Parameter(
            torch.full((self.heads,), _inverse_softplus(448.0))
        )

        self.ffn_norm = CabaRMSNorm(self.hidden, self.eps)
        self.ffn_gate_up = nn.Linear(self.hidden, 2 * config.intermediate_size, bias=False)
        self.ffn_down = nn.Linear(config.intermediate_size, self.hidden, bias=False)

    def _decays(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        t_fast = F.softplus(self.raw_fast_timescale.float()) + 1.0
        t_medium = t_fast + F.softplus(self.raw_medium_gap.float())
        t_slow = t_medium + F.softplus(self.raw_slow_gap.float())
        return torch.exp(-1.0 / t_fast), torch.exp(-1.0 / t_medium), torch.exp(-1.0 / t_slow)

    def _project(self, x: torch.Tensor):
        batch, length, _ = x.shape
        normalized = self.input_norm(x)
        conv_in = F.pad(normalized.transpose(1, 2), (self.conv_kernel - 1, 0))
        h = F.silu(self.local_conv(conv_in).transpose(1, 2))
        q, k, v = self.qkv(h).view(
            batch, length, 3, self.heads, self.head_dim
        ).unbind(dim=2)
        q = F.normalize(q.float(), dim=-1, eps=1e-6)
        k = F.normalize(k.float(), dim=-1, eps=1e-6)
        v = v.float()
        write = torch.sigmoid(self.fast_write(h).float()).unsqueeze(-1)
        return h, q, k, v, write

    def _fast_sequence(self, q, k, v, write):
        batch, length, heads, dim = q.shape
        decay_fast, _, _ = self._decays()
        decay_fast = decay_fast.view(1, heads, 1, 1)
        xs = torch.cat((q, k, v, write), dim=-1).permute(1, 0, 2, 3).contiguous()
        decay_steps = decay_fast.reshape(1, 1, heads, 1).expand(length, batch, heads, 1)
        xs = torch.cat((xs, decay_steps), dim=-1).contiguous()
        xs = _mark_data_sharding(xs, self.data_mesh, (None, "data", None, None))
        initial = torch.zeros(
            batch, heads, dim, dim, device=q.device, dtype=torch.float32
        )
        initial = _mark_data_sharding(initial, self.data_mesh, ("data", None, None, None))

        def step(memory, item):
            query, key, value, gate, decay = torch.split(
                item, (dim, dim, dim, 1, 1), dim=-1
            )
            key = F.normalize(key.float(), dim=-1, eps=1e-6)
            error, score = _predict_error(memory, key.squeeze(-1) if key.shape[-1] == 1 else key, value)
            next_memory = decay.unsqueeze(-1) * memory + gate.float().unsqueeze(-1) * _outer(error, key)
            read = _read(next_memory, query)
            packed = torch.cat((read, error, score.unsqueeze(-1)), dim=-1)
            return next_memory, packed

        final, packed = _scan(step, initial, xs)
        final = _mark_data_sharding(final, self.data_mesh, ("data", None, None, None))
        packed = _mark_data_sharding(packed, self.data_mesh, (None, "data", None, None))
        packed = packed.permute(1, 0, 2, 3).contiguous()
        read = _mark_data_sharding(packed[..., :dim], self.data_mesh, ("data", None, None, None))
        error = _mark_data_sharding(packed[..., dim : 2 * dim], self.data_mesh, ("data", None, None, None))
        score = _mark_data_sharding(packed[..., -1], self.data_mesh, ("data", None, None))
        return final, read, error, score

    def _medium_sequence(self, q, k, v, scores):
        batch, length, heads, dim = q.shape
        clock = self.medium_clock
        if length % clock:
            raise ValueError(
                f"Training sequence length {length} must be divisible by medium_clock={clock}."
            )
        windows = length // clock
        q_windows = q.reshape(batch, windows, clock, heads, dim).permute(0, 1, 3, 2, 4)
        k_windows = k.reshape(batch, windows, clock, heads, dim).permute(0, 1, 3, 2, 4)
        v_windows = v.reshape(batch, windows, clock, heads, dim).permute(0, 1, 3, 2, 4)
        s_windows = scores.reshape(batch, windows, clock, heads).permute(0, 1, 3, 2)

        selection = torch.softmax(s_windows, dim=-1)
        summary_k = (selection.unsqueeze(-1) * k_windows).sum(dim=-2)
        summary_v = (selection.unsqueeze(-1) * v_windows).sum(dim=-2)
        log_mean_exp = torch.logsumexp(s_windows, dim=-1) - math.log(clock)
        admission = torch.sigmoid(
            log_mean_exp - self.medium_admit_threshold.float().view(1, 1, heads)
        ).unsqueeze(-1)

        q_flat = q_windows.reshape(batch, windows, heads, clock * dim).permute(1, 0, 2, 3)
        xs = torch.cat(
            (
                q_flat,
                summary_k.permute(1, 0, 2, 3),
                summary_v.permute(1, 0, 2, 3),
                admission.permute(1, 0, 2, 3),
            ),
            dim=-1,
        ).contiguous()
        initial = torch.zeros(batch, heads, dim, dim, device=q.device, dtype=torch.float32)
        _, decay_medium, _ = self._decays()
        decay_medium = decay_medium.pow(clock).view(1, heads, 1, 1)
        decay_steps = decay_medium.reshape(1, 1, heads, 1).expand(windows, batch, heads, 1)
        xs = torch.cat((xs, decay_steps), dim=-1).contiguous()
        xs = _mark_data_sharding(xs, self.data_mesh, (None, "data", None, None))
        initial = _mark_data_sharding(initial, self.data_mesh, ("data", None, None, None))

        def step(memory, item):
            offset = clock * dim
            query = item[..., :offset].reshape(batch, heads, clock, dim).permute(0, 2, 1, 3)
            key = item[..., offset : offset + dim]
            value = item[..., offset + dim : offset + 2 * dim]
            gate = item[..., -2:-1]
            decay = item[..., -1:]
            read = torch.einsum("bhkv,brhk->brhv", memory, query.float())
            error, score = _predict_error(memory, key, value)
            next_memory = (
                decay.unsqueeze(-1) * memory + gate.float().unsqueeze(-1) * _outer(error, key)
            )
            repeated_error = error.unsqueeze(2).expand(-1, -1, clock, -1)
            repeated_score = score.unsqueeze(-1).unsqueeze(2).expand(-1, -1, clock, -1)
            packed = torch.cat(
                (read.permute(0, 2, 1, 3), repeated_error, repeated_score), dim=-1
            ).reshape(batch, heads, clock * (2 * dim + 1))
            return next_memory, packed

        final, packed = _scan(step, initial, xs)
        final = _mark_data_sharding(final, self.data_mesh, ("data", None, None, None))
        packed = _mark_data_sharding(packed, self.data_mesh, (None, "data", None, None))
        packed = packed.reshape(windows, batch, heads, clock, 2 * dim + 1)
        packed = packed.permute(1, 0, 3, 2, 4).contiguous()
        read = _mark_data_sharding(packed[..., :dim].reshape(batch, length, heads, dim), self.data_mesh, ("data", None, None, None))
        # Every chunk repeats its single medium residual/score for fixed-shape output.
        error = _mark_data_sharding(packed[:, :, 0, :, dim : 2 * dim], self.data_mesh, ("data", None, None, None))
        score = _mark_data_sharding(packed[:, :, 0, :, -1], self.data_mesh, ("data", None, None))
        summary_k = _mark_data_sharding(summary_k, self.data_mesh, ("data", None, None, None))
        summary_v = _mark_data_sharding(summary_v, self.data_mesh, ("data", None, None, None))
        admission = _mark_data_sharding(admission.squeeze(-1), self.data_mesh, ("data", None, None))
        return final, read, summary_k, summary_v, error, score, admission

    def _slow_sequence(self, q, medium_k, medium_v, medium_scores):
        batch, length, heads, dim = q.shape
        clock = self.slow_span
        if length % clock:
            raise ValueError(
                f"Training sequence length {length} must be divisible by slow span={clock}."
            )
        windows = length // clock
        medium_windows_per_slow = self.slow_clock
        score_windows = medium_scores.reshape(
            batch, windows, medium_windows_per_slow, heads
        ).permute(0, 1, 3, 2)
        mk = medium_k.reshape(batch, windows, medium_windows_per_slow, heads, dim).permute(
            0, 1, 3, 2, 4
        )
        mv = medium_v.reshape(batch, windows, medium_windows_per_slow, heads, dim).permute(
            0, 1, 3, 2, 4
        )
        selection = torch.softmax(score_windows, dim=-1)
        summary_k = (selection.unsqueeze(-1) * mk).sum(dim=-2)
        summary_v = (selection.unsqueeze(-1) * mv).sum(dim=-2)
        log_mean_exp = torch.logsumexp(score_windows, dim=-1) - math.log(
            medium_windows_per_slow
        )
        admission = torch.sigmoid(
            log_mean_exp - self.slow_admit_threshold.float().view(1, 1, heads)
        ).unsqueeze(-1)

        q_windows = q.reshape(batch, windows, clock, heads, dim).permute(0, 1, 3, 2, 4)
        q_flat = q_windows.reshape(batch, windows, heads, clock * dim).permute(1, 0, 2, 3)
        xs = torch.cat(
            (
                q_flat,
                summary_k.permute(1, 0, 2, 3),
                summary_v.permute(1, 0, 2, 3),
                admission.permute(1, 0, 2, 3),
            ),
            dim=-1,
        ).contiguous()
        initial = torch.zeros(batch, heads, dim, dim, device=q.device, dtype=torch.float32)
        _, _, decay_slow = self._decays()
        decay_slow = decay_slow.pow(clock).view(1, heads, 1, 1)
        decay_steps = decay_slow.reshape(1, 1, heads, 1).expand(windows, batch, heads, 1)
        xs = torch.cat((xs, decay_steps), dim=-1).contiguous()
        xs = _mark_data_sharding(xs, self.data_mesh, (None, "data", None, None))
        initial = _mark_data_sharding(initial, self.data_mesh, ("data", None, None, None))

        def step(memory, item):
            offset = clock * dim
            query = item[..., :offset].reshape(batch, heads, clock, dim).permute(0, 2, 1, 3)
            key = item[..., offset : offset + dim]
            value = item[..., offset + dim : offset + 2 * dim]
            gate = item[..., -2:-1]
            decay = item[..., -1:]
            read = torch.einsum("bhkv,brhk->brhv", memory, query.float())
            error, _ = _predict_error(memory, key, value)
            next_memory = (
                decay.unsqueeze(-1) * memory + gate.float().unsqueeze(-1) * _outer(error, key)
            )
            return next_memory, read.permute(0, 2, 1, 3).reshape(
                batch, heads, clock * dim
            )

        final, read_flat = _scan(step, initial, xs)
        final = _mark_data_sharding(final, self.data_mesh, ("data", None, None, None))
        read_flat = _mark_data_sharding(read_flat, self.data_mesh, (None, "data", None, None))
        read = read_flat.reshape(windows, batch, heads, clock, dim)
        read = _mark_data_sharding(read.permute(1, 0, 3, 2, 4).reshape(batch, length, heads, dim), self.data_mesh, ("data", None, None, None))
        admission = _mark_data_sharding(admission.squeeze(-1), self.data_mesh, ("data", None, None))
        return final, read, admission

    def forward_sequence(self, x: torch.Tensor):
        batch, length, _ = x.shape
        if length % (self.medium_clock * self.slow_clock):
            raise ValueError(
                "use_cache=False training/evaluation sequence length must be divisible "
                f"by medium_clock*slow_clock={self.medium_clock * self.slow_clock}; "
                "the notebook default is 512."
            )
        h, q, k, v, write = self._project(x)
        h = _mark_data_sharding(h, self.data_mesh, ("data", None, None))
        q = _mark_data_sharding(q, self.data_mesh, ("data", None, None, None))
        k = _mark_data_sharding(k, self.data_mesh, ("data", None, None, None))
        v = _mark_data_sharding(v, self.data_mesh, ("data", None, None, None))
        write = _mark_data_sharding(write, self.data_mesh, ("data", None, None, None))
        _, fast_read, fast_error, fast_score = self._fast_sequence(q, k, v, write)
        unresolvedness = fast_score.detach().mean()
        fast_score = F.softplus(self.fast_score_alpha.float()).view(1, 1, -1) * (
            fast_score - self.fast_score_tau.float().view(1, 1, -1)
        )
        _, medium_read, medium_k, medium_v, medium_error, medium_score, medium_admission = (
            self._medium_sequence(q, k, v, fast_score)
        )
        medium_score = F.softplus(self.medium_score_alpha.float()).view(1, 1, -1) * (
            medium_score - self.medium_score_tau.float().view(1, 1, -1)
        )
        _, slow_read, slow_admission = self._slow_sequence(
            q, medium_k, medium_v, medium_score
        )

        gates = torch.softmax(
            self.merge_gate(h).view(batch, length, self.heads, 3).float(), dim=-1
        )
        reads = torch.stack((fast_read, medium_read, slow_read), dim=-2)
        merged = (reads * gates.unsqueeze(-1)).sum(dim=-2).reshape(batch, length, self.hidden)
        x = x + self.out_proj(merged.to(x.dtype))
        gate_up = self.ffn_gate_up(self.ffn_norm(x))
        gate, up = gate_up.chunk(2, dim=-1)
        x = x + self.ffn_down(F.silu(gate) * up)
        metrics = {
            "unresolvedness": unresolvedness,
            "medium_admission": medium_admission.detach().mean(),
            "slow_admission": slow_admission.detach().mean(),
        }
        return x, metrics

    def _new_cache(self, batch: int, device: torch.device) -> Dict[str, Any]:
        h, d = self.heads, self.head_dim
        return {
            "fast": torch.zeros(batch, h, d, d, device=device, dtype=torch.float32),
            "medium": torch.zeros(batch, h, d, d, device=device, dtype=torch.float32),
            "slow": torch.zeros(batch, h, d, d, device=device, dtype=torch.float32),
            "conv": torch.zeros(
                batch, max(self.conv_kernel - 1, 0), self.hidden, device=device
            ),
            "medium_pending_k": [],
            "medium_pending_v": [],
            "medium_pending_score": [],
            "slow_pending_k": [],
            "slow_pending_v": [],
            "slow_pending_score": [],
            "position": 0,
        }

    def step(self, x: torch.Tensor, cache: Dict[str, Any], position: int):
        batch = x.shape[0]
        normalized = self.input_norm(x)
        context = torch.cat((cache["conv"], normalized), dim=1)
        conv = self.local_conv(context.transpose(1, 2)).transpose(1, 2)
        h = F.silu(conv)
        conv_state = (
            context[:, -(self.conv_kernel - 1) :, :]
            if self.conv_kernel > 1
            else context[:, :0, :]
        )

        q, k, v = self.qkv(h).view(batch, 1, 3, self.heads, self.head_dim).unbind(dim=2)
        q = F.normalize(q[:, 0].float(), dim=-1, eps=1e-6)
        k = F.normalize(k[:, 0].float(), dim=-1, eps=1e-6)
        v = v[:, 0].float()
        write = torch.sigmoid(self.fast_write(h).float()[:, 0]).unsqueeze(-1)

        decay_fast, decay_medium, decay_slow = self._decays()
        fast_error, fast_unresolvedness = _predict_error(cache["fast"], k, v)
        fast_next = (
            decay_fast.view(1, self.heads, 1, 1) * cache["fast"]
            + write.float().unsqueeze(-1) * _outer(fast_error, k)
        )
        fast_read = _read(fast_next, q)
        medium_read = _read(cache["medium"], q)
        slow_read = _read(cache["slow"], q)

        score_fast = F.softplus(self.fast_score_alpha.float()).view(1, -1) * (
            fast_unresolvedness - self.fast_score_tau.float().view(1, -1)
        )
        pending_k = cache["medium_pending_k"] + [k]
        pending_v = cache["medium_pending_v"] + [v]
        pending_score = cache["medium_pending_score"] + [score_fast]
        medium_next = cache["medium"]
        slow_next = cache["slow"]
        slow_pending_k = cache["slow_pending_k"]
        slow_pending_v = cache["slow_pending_v"]
        slow_pending_score = cache["slow_pending_score"]
        medium_admission_value = torch.zeros((), device=x.device)
        slow_admission_value = torch.zeros((), device=x.device)

        if (position + 1) % self.medium_clock == 0:
            k_stack = torch.stack(pending_k, dim=2)
            v_stack = torch.stack(pending_v, dim=2)
            s_stack = torch.stack(pending_score, dim=2)
            selected = torch.softmax(s_stack, dim=2)
            summary_k = (selected.unsqueeze(-1) * k_stack).sum(dim=2)
            summary_v = (selected.unsqueeze(-1) * v_stack).sum(dim=2)
            log_mean_exp = torch.logsumexp(s_stack, dim=2) - math.log(self.medium_clock)
            medium_admission = torch.sigmoid(
                log_mean_exp
                - self.medium_admit_threshold.float().view(1, self.heads)
            ).unsqueeze(-1)
            medium_error, medium_unresolvedness = _predict_error(
                cache["medium"], summary_k, summary_v
            )
            medium_next = (
                decay_medium.pow(self.medium_clock).view(1, self.heads, 1, 1)
                * cache["medium"]
                + medium_admission.unsqueeze(-1) * _outer(medium_error, summary_k)
            )
            score_medium = F.softplus(self.medium_score_alpha.float()).view(1, -1) * (
                medium_unresolvedness - self.medium_score_tau.float().view(1, -1)
            )
            slow_pending_k = slow_pending_k + [summary_k]
            slow_pending_v = slow_pending_v + [summary_v]
            slow_pending_score = slow_pending_score + [score_medium]
            pending_k, pending_v, pending_score = [], [], []
            medium_admission_value = medium_admission.mean()

            if len(slow_pending_k) == self.slow_clock:
                sk = torch.stack(slow_pending_k, dim=2)
                sv = torch.stack(slow_pending_v, dim=2)
                ss = torch.stack(slow_pending_score, dim=2)
                slow_selected = torch.softmax(ss, dim=2)
                slow_key = (slow_selected.unsqueeze(-1) * sk).sum(dim=2)
                slow_value = (slow_selected.unsqueeze(-1) * sv).sum(dim=2)
                slow_log_mean_exp = torch.logsumexp(ss, dim=2) - math.log(self.slow_clock)
                slow_admission = torch.sigmoid(
                    slow_log_mean_exp
                    - self.slow_admit_threshold.float().view(1, self.heads)
                ).unsqueeze(-1)
                slow_error, _ = _predict_error(cache["slow"], slow_key, slow_value)
                slow_next = (
                    decay_slow.pow(self.slow_span).view(1, self.heads, 1, 1)
                    * cache["slow"]
                    + slow_admission.unsqueeze(-1) * _outer(slow_error, slow_key)
                )
                slow_pending_k, slow_pending_v, slow_pending_score = [], [], []
                slow_admission_value = slow_admission.mean()

        gates = torch.softmax(
            self.merge_gate(h).view(batch, 1, self.heads, 3).float(), dim=-1
        )[:, 0]
        reads = torch.stack((fast_read, medium_read, slow_read), dim=-2)
        merged = (reads * gates.unsqueeze(-1)).sum(dim=-2).reshape(batch, 1, self.hidden)
        x = x + self.out_proj(merged.to(x.dtype))
        gate_up = self.ffn_gate_up(self.ffn_norm(x))
        gate, up = gate_up.chunk(2, dim=-1)
        x = x + self.ffn_down(F.silu(gate) * up)

        cache = {
            "fast": fast_next,
            "medium": medium_next,
            "slow": slow_next,
            "conv": conv_state,
            "medium_pending_k": pending_k,
            "medium_pending_v": pending_v,
            "medium_pending_score": pending_score,
            "slow_pending_k": slow_pending_k,
            "slow_pending_v": slow_pending_v,
            "slow_pending_score": slow_pending_score,
            "position": position + 1,
        }
        metrics = {
            "unresolvedness": fast_unresolvedness.detach().mean(),
            "medium_admission": medium_admission_value.detach(),
            "slow_admission": slow_admission_value.detach(),
        }
        return x, cache, metrics


class CabaPreTrainedModel(PreTrainedModel):
    config_class = CabaConfig
    base_model_prefix = "caba"
    supports_gradient_checkpointing = False
    _supports_cache_class = False

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, (nn.Linear, nn.Embedding, nn.Conv1d)):
            nn.init.normal_(module.weight, mean=0.0, std=std)
            if getattr(module, "bias", None) is not None:
                nn.init.zeros_(module.bias)
            if isinstance(module, nn.Embedding) and module.padding_idx is not None:
                with torch.no_grad():
                    module.weight[module.padding_idx].zero_()
        elif isinstance(module, CabaRMSNorm):
            nn.init.ones_(module.weight)


class CabaModel(CabaPreTrainedModel):
    def __init__(self, config: CabaConfig):
        super().__init__(config)
        self.data_mesh = None
        self.embed_tokens = nn.Embedding(
            config.vocab_size,
            config.hidden_size,
            padding_idx=config.pad_token_id,
        )
        self.layers = nn.ModuleList(
            [CabaBlock(config) for _ in range(config.num_hidden_layers)]
        )
        self.norm = CabaRMSNorm(config.hidden_size, config.rms_norm_eps)

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    def set_data_mesh(self, mesh):
        self.data_mesh = mesh
        for layer in self.layers:
            layer.data_mesh = mesh

    def forward(
        self,
        input_ids: torch.LongTensor,
        past_key_values=None,
        use_cache: bool = False,
        **kwargs,
    ):
        if input_ids.ndim != 2:
            raise ValueError("input_ids must have shape [batch, sequence]")
        x = self.embed_tokens(input_ids)
        x = _mark_data_sharding(x, self.data_mesh, ("data", None, None))
        metrics_by_layer = []

        if use_cache or past_key_values is not None:
            caches = (
                list(past_key_values)
                if past_key_values is not None
                else [
                    layer._new_cache(input_ids.shape[0], input_ids.device)
                    for layer in self.layers
                ]
            )
            position = caches[0]["position"] if caches else 0
            pieces = []
            step_metrics = []
            for offset in range(input_ids.shape[1]):
                hidden = x[:, offset : offset + 1]
                new_caches = []
                current_metrics = []
                for layer, cache in zip(self.layers, caches):
                    hidden, cache, metrics = layer.step(hidden, cache, position + offset)
                    new_caches.append(cache)
                    current_metrics.append(metrics)
                caches = new_caches
                pieces.append(hidden)
                step_metrics.append(current_metrics)
            x = torch.cat(pieces, dim=1)
            for layer_index in range(len(self.layers)):
                metrics_by_layer.append(
                    {
                        key: torch.stack(
                            [m[layer_index][key] for m in step_metrics]
                        ).mean()
                        for key in ("unresolvedness", "medium_admission", "slow_admission")
                    }
                )
            cache_output = tuple(caches)
        else:
            for layer in self.layers:
                if self.config.gradient_checkpointing and self.training:
                    def _run_checkpointed(hidden, _layer=layer):
                        hidden_out, layer_metrics = _layer.forward_sequence(hidden)
                        values = torch.stack([
                            layer_metrics["unresolvedness"].detach().float(),
                            layer_metrics["medium_admission"].detach().float(),
                            layer_metrics["slow_admission"].detach().float(),
                        ])
                        return hidden_out, values
                    checkpoint_fn = _xla_checkpoint if x.device.type == "xla" else _torch_checkpoint
                    if checkpoint_fn is None:
                        raise RuntimeError("TPU gradient checkpointing requires torch_xla.utils.checkpoint")
                    x, values = checkpoint_fn(_run_checkpointed, x, use_reentrant=True)
                    metrics_by_layer.append({
                        "unresolvedness": values[0],
                        "medium_admission": values[1],
                        "slow_admission": values[2],
                    })
                else:
                    x, metrics = layer.forward_sequence(x)
                    metrics_by_layer.append(metrics)
                x = _mark_data_sharding(x, self.data_mesh, ("data", None, None))
            cache_output = None

        x = self.norm(x)
        x = _mark_data_sharding(x, self.data_mesh, ("data", None, None))
        keys = ("unresolvedness", "medium_admission", "slow_admission")
        metrics = {
            key: torch.stack([entry[key] for entry in metrics_by_layer]).mean()
            for key in keys
        }
        return x, cache_output, metrics


class CabaForCausalLM(CabaPreTrainedModel, GenerationMixin):
    """Hugging Face causal-LM wrapper with recurrent cached generation."""

    # Transformers 4 expects a list; Transformers 5 expects tied-key -> source-key mapping.
    _tied_weights_keys = (
        {"lm_head.weight": "model.embed_tokens.weight"}
        if int(_transformers.__version__.split(".", 1)[0]) >= 5
        else ["lm_head.weight"]
    )
    _supports_cache_class = False
    main_input_name = "input_ids"

    def __init__(self, config: CabaConfig):
        super().__init__(config)
        self.model = CabaModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

        if config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.set_input_embeddings(value)

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, value):
        self.lm_head = value

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        past_key_values=None,
        use_cache: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ):
        if use_cache is None:
            use_cache = bool(self.config.use_cache and not self.training)
        hidden, cache, metrics = self.model(
            input_ids,
            past_key_values=past_key_values,
            use_cache=bool(use_cache),
        )
        logits = self.lm_head(hidden).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, self.config.vocab_size),
                shift_labels.view(-1),
                ignore_index=-100,
            )
        output = CabaCausalLMOutput(
            loss=loss,
            logits=logits,
            past_key_values=cache,
            memory_metrics=metrics,
        )
        if return_dict is False:
            items = (logits, cache, metrics)
            return ((loss,) + items) if loss is not None else items
        return output

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.LongTensor,
        max_new_tokens: int = 64,
        do_sample: bool = False,
        temperature: float = 1.0,
        top_k: int = 0,
        top_p: float = 1.0,
        eos_token_id: Optional[int] = None,
        pad_token_id: Optional[int] = None,
        attention_mask: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> torch.LongTensor:
        """Generate with the recurrent state; supports greedy and top-k/top-p sampling."""
        if input_ids.ndim != 2 or input_ids.shape[1] == 0:
            raise ValueError("input_ids must be a non-empty [batch, sequence] tensor")
        if temperature <= 0:
            raise ValueError("temperature must be > 0")
        if kwargs:
            unsupported = ", ".join(sorted(kwargs))
            raise TypeError(f"Unsupported generation options in ÇABA v0-core: {unsupported}")
        self.eval()
        eos = self.config.eos_token_id if eos_token_id is None else eos_token_id
        pad = self.config.pad_token_id if pad_token_id is None else pad_token_id
        sequences = input_ids
        output = self(input_ids, use_cache=True)
        cache = output.past_key_values
        logits = output.logits[:, -1, :]
        finished = torch.zeros(input_ids.shape[0], dtype=torch.bool, device=input_ids.device)

        for _ in range(max_new_tokens):
            scores = logits.float()
            if do_sample:
                scores = scores / temperature
                if top_k > 0:
                    k = min(int(top_k), scores.shape[-1])
                    threshold = torch.topk(scores, k, dim=-1).values[:, -1, None]
                    scores = scores.masked_fill(scores < threshold, -float("inf"))
                if 0.0 < top_p < 1.0:
                    sorted_scores, sorted_ids = torch.sort(scores, descending=True, dim=-1)
                    cumulative = torch.softmax(sorted_scores, dim=-1).cumsum(dim=-1)
                    remove = cumulative > top_p
                    remove[:, 1:] = remove[:, :-1].clone()
                    remove[:, 0] = False
                    sorted_scores = sorted_scores.masked_fill(remove, -float("inf"))
                    scores = torch.full_like(scores, -float("inf")).scatter(
                        1, sorted_ids, sorted_scores
                    )
                next_token = torch.multinomial(torch.softmax(scores, dim=-1), num_samples=1)
            else:
                next_token = scores.argmax(dim=-1, keepdim=True)
            if eos is not None:
                next_token = torch.where(
                    finished[:, None],
                    torch.full_like(next_token, int(pad)),
                    next_token,
                )
            sequences = torch.cat((sequences, next_token), dim=1)
            if eos is not None:
                finished |= next_token[:, 0].eq(int(eos))
                if bool(finished.all()):
                    break
            output = self(next_token, past_key_values=cache, use_cache=True)
            cache = output.past_key_values
            logits = output.logits[:, -1, :]
        return sequences