"""ATOM-native text model built around ``AtomPacket`` observations.

The canonical toroidal modules are intentionally not modified here.  This
adapter replaces the GPT-2 ID/embedding input path with:

    AtomPacket.features -> eight toroidal atom properties

and replaces the 50k-class surface target with a bounded byte packet head.
The recurrent field, RK4 dynamics, interactions, aggregation, abstraction and
consolidation are the existing ATOM components.

This is the primary text path. Legacy GPT-2 tokenizer wrappers live under
``_quarantine_transformer_slop/``. The toroidal core remains in ``src.toroidal``.
"""

from __future__ import annotations

import math
from pathlib import Path
from typing import Iterable

import torch
import torch.nn as nn
import torch.nn.functional as F

from .io.atomizer import AtomPacket, Atomizer
from .speech_lock import format_dialogue_prompt, speech_ok
from .toroidal.atom import ToroidalAtom
from .toroidal.model import ToroidalFractalIntelligence

# Sane band for shared energy_decay (always applied).  Values near 0 turn the
# decay term into near-total wipe; values >1 flip it into growth.
DEFAULT_ENERGY_DECAY_BOUNDS: tuple[float, float] = (0.3, 0.95)

# Fixed distinct prompts for L_ign α_bar when --field-ignorance-prompt-bank is on.
# Probe-style set (not adjacent train ticks). Keep short for refresh cost.
DEFAULT_IGNORANCE_PROMPTS: tuple[str, ...] = (
    "Bonjour",
    "Qui es-tu ?",
    "Il était une fois",
    "Utilisateur: Bonjour\nAssistant:",
    "The future of AI is",
    "Valkyria Chronicles III is",
    "Hello, how are you?",
    "Comment ça va aujourd'hui ?",
)


class AtomCompiler(nn.Module):
    """Compile a packet observation into the eight atom properties.

    There is no vocabulary lookup here.  The only learned input map receives
    the fixed-size physical observation produced by ``Atomizer``.
    """

    def __init__(self, feature_dim: int, d_model: int, hidden_dim: int | None = None) -> None:
        super().__init__()
        hidden_dim = hidden_dim or max(32, 2 * d_model)
        self.feature_dim = feature_dim
        self.d_model = d_model
        self.mapping = nn.Sequential(
            nn.Linear(feature_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, 8 * d_model),
        )
        self.operation_gate = nn.Linear(feature_dim, 7)

    def forward(self, features: torch.Tensor, atom_count: int = 0) -> tuple[ToroidalAtom, torch.Tensor, torch.Tensor]:
        if features.ndim == 1:
            features = features.unsqueeze(0)
        if features.shape[-1] != self.feature_dim:
            raise ValueError(
                f"expected packet features with {self.feature_dim} values, got {features.shape[-1]}"
            )
        raw = self.mapping(features).view(features.shape[0], 8, self.d_model)
        r = F.normalize(raw[:, 0], dim=-1)
        phi = raw[:, 1]
        omega = F.softplus(raw[:, 2]) + 1e-3
        energy = torch.sigmoid(raw[:, 3])
        kappa = F.softplus(raw[:, 4])
        memory = F.normalize(raw[:, 5], dim=-1)
        tau = F.softplus(raw[:, 6]) + 1e-6
        rho = torch.full(
            (features.shape[0], 1),
            min(atom_count // 64, 4),
            dtype=features.dtype,
            device=features.device,
        )
        atom = ToroidalAtom(
            r=r,
            phi=phi,
            omega=omega,
            E=energy,
            kappa=kappa,
            M=memory,
            tau=tau,
            rho=rho,
        )
        operation_logits = self.operation_gate(features)
        operation = operation_logits.argmax(dim=-1)
        confidence = operation_logits.softmax(dim=-1).max(dim=-1).values
        return atom, operation, confidence


class AlphaOnlyMLP(nn.Module):
    """α-conditioned 2-layer map: Linear → GELU → Linear (no byte_decoder bypass).

    Used under ``field_obligatory_hard`` as the *only* trainable α→logit adapter.
    Input is the JL fingerprint (or any α-only feature); hidden defaults to 4*d.
    """

    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        hidden_dim: int,
        *,
        bias_out: bool = False,
    ) -> None:
        super().__init__()
        if hidden_dim < 1:
            raise ValueError("hidden_dim must be positive")
        self.in_dim = int(in_dim)
        self.hidden_dim = int(hidden_dim)
        self.out_dim = int(out_dim)
        self.fc1 = nn.Linear(self.in_dim, self.hidden_dim)
        self.fc2 = nn.Linear(self.hidden_dim, self.out_dim, bias=bias_out)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc2(F.gelu(self.fc1(x)))


class LastAtomReadout(nn.Module):
    """Last-atom readout + dentate 2-gram (ARCHITECTURE.md, bridges #20/#1).

    Next-byte logits from the last committed byte + previous byte + last atom
    phase. NOT payload-copy (no hist/winner broadcast).

    1. embed last byte, embed previous byte (2-gram)
    2. dentate: expand 4x -> ReLU -> keep top 25% -> project back
    3. add phase(cos phi, sin phi) projection, GELU
    4. linear to 256 logits
    """

    def __init__(self, d_model: int, embed_dim: int | None = None, expand: int = 4) -> None:
        super().__init__()
        embed_dim = int(embed_dim or d_model)
        self.embed_dim = embed_dim
        self.d_model = int(d_model)
        self.byte_embed = nn.Embedding(256, embed_dim)
        hidden = expand * 2 * embed_dim
        self.expand = nn.Linear(2 * embed_dim, hidden)
        self.contract = nn.Linear(hidden, embed_dim)
        self.phase_proj = nn.Linear(2 * d_model, embed_dim)
        self.head = nn.Linear(embed_dim, 256)

    def dentate_keep(self, hidden_dim: int) -> int:
        return max(1, hidden_dim // 4)

    def forward(self, last_byte: int, prev_byte: int, phi: torch.Tensor) -> torch.Tensor:
        """Return 256 next-byte logits. ``phi`` is (d_model,) live phase."""
        device = self.head.weight.device
        last = torch.tensor([int(last_byte) % 256], dtype=torch.long, device=device)
        prev = torch.tensor([int(prev_byte) % 256], dtype=torch.long, device=device)
        pair = torch.cat([self.byte_embed(last), self.byte_embed(prev)], dim=-1)
        h = F.relu(self.expand(pair))
        keep = self.dentate_keep(h.shape[-1])
        top_values, top_index = torch.topk(h, keep, dim=-1)
        sparse = torch.zeros_like(h).scatter(-1, top_index, top_values)
        z = self.contract(sparse)
        flat = phi.reshape(-1).to(device=device, dtype=z.dtype)
        if flat.numel() != self.d_model:
            raise ValueError(f"last-atom phi needs {self.d_model} elems, got {flat.numel()}")
        phase = torch.cat([torch.cos(flat), torch.sin(flat)], dim=-1)
        out = F.gelu(z + self.phase_proj(phase).unsqueeze(0))
        return self.head(out).reshape(256)


def _byte_cycle(buf: bytearray | bytes) -> bool:
    """True if the tail of ``buf`` closes a 2-8 byte repeat (anti-repeat).

    ARCHITECTURE.md says 2-6; extended to 8 because byte-tick outputs loop
    on 7-byte attractors (e.g. "Tu ont: "). Decode only, weights untouched.
    """
    n_buf = len(buf)
    for n in (2, 3, 4, 5, 6, 7, 8):
        if n_buf >= 2 * n and bytes(buf[-n:]) == bytes(buf[-2 * n : -n]):
            return True
    return False


def _tail_bytes(atom_payloads: list[bytes] | None) -> tuple[int, int]:
    """(last byte, previous byte) of living payloads; prev falls back to last."""
    tails = [bytes(p or b"") for p in (atom_payloads or [])]
    tails = [p for p in tails if p]
    last = tails[-1][-1] if tails else 0
    prev = tails[-2][-1] if len(tails) > 1 else last
    return last, prev


class AtomSurfaceHead(nn.Module):

    """Decode one variable-length surface packet from the living toroidal field.

    The output alphabet is raw bytes (256 classes), not GPT-2's 50,257 token
    vocabulary.  A bounded packet contains 1..``max_payload_bytes`` bytes and
    has a separate length distribution.

    Readout contract (field-first):
      surface_input = f(spectral field α, optional consolidation, atom.r)
    A LayerNorm + legacy linear path is retained for checkpoint compatibility,
    but every decode also mixes an explicit spectral summary of α (mean‖std)
    and a bias-free field→logit skip so prompt-conditioned α is not drowned by
    a shared decoder bias (the failure mode where inter-prompt logit cosine
    stayed ≈ 0.998 while α cosine was ≈ 0.74).
    """

    # Printable UTF-8 / Latin-1 friendly prior used when ``prefer_printable``.
    _PRINTABLE = set(range(0x20, 0x7F)) | {0x09, 0x0A, 0x0D} | set(range(0xC2, 0xF5)) | set(range(0x80, 0xC0))

    def __init__(
        self,
        d_model: int,
        max_payload_bytes: int = 32,
        n_modes: int | None = None,
        last_atom_readout: bool = False,
    ) -> None:
        super().__init__()
        if max_payload_bytes < 1:
            raise ValueError("max_payload_bytes must be positive")
        self.d_model = d_model
        self.n_modes = int(n_modes) if n_modes is not None else d_model
        self.max_payload_bytes = max_payload_bytes
        # Spectral living-field features: [mean(α) ‖ std(α) ‖ persist ‖ atom.r]
        self.field_feat_dim = 4 * d_model
        self.field_feat_norm = nn.LayerNorm(self.field_feat_dim)
        self.field_to_state = nn.Linear(self.field_feat_dim, d_model)
        # sigmoid(field_gate)≈0.88 at init → field-conditioned state dominates.
        self.field_gate = nn.Parameter(torch.tensor(2.0))
        # Bias-free skip so α differences reach logits even when |b| is large.
        self.field_byte_skip = nn.Linear(self.field_feat_dim, max_payload_bytes * 256, bias=False)
        self.field_length_skip = nn.Linear(self.field_feat_dim, max_payload_bytes, bias=False)
        # softplus(4)≈4.0 — enough for inter-prompt logit cosine << 0.998 at load.
        self.skip_gate = nn.Parameter(torch.tensor(1.0))
        # Hard-path couple: dedicated (NOT in _freeze_non_alpha_bypass). Soft hard
        # zeros field_byte_skip; this head re-opens persist‖atom.r → logits when
        # field_hard_couple_scale > 0.
        self.field_hard_couple_byte = nn.Linear(
            self.field_feat_dim, max_payload_bytes * 256, bias=False
        )
        self.field_hard_couple_length = nn.Linear(
            self.field_feat_dim, max_payload_bytes, bias=False
        )
        # Obligatory α readout (main CE path): full-α JL fingerprint → logits.
        # Frozen random map cannot be CE-collapsed; trainable residual can adapt.
        # zeros_like(α) ⇒ exact zero branch. Persist/atom_r cannot bypass.
        self.alpha_only_dim = 2 * d_model
        self._alpha_flat_dim = int(self.n_modes) * int(d_model)
        # Fixed JL projection of flattened α (buffer — not trained).
        jl = torch.empty(self.alpha_only_dim, self._alpha_flat_dim)
        torch.manual_seed(20260918)
        nn.init.normal_(jl, mean=0.0, std=(1.0 / self._alpha_flat_dim) ** 0.5)
        self.register_buffer("alpha_jl", jl, persistent=True)
        # Frozen random α→logit map (buffer) + trainable adapter.
        fr_b = torch.empty(max_payload_bytes * 256, self.alpha_only_dim)
        fr_l = torch.empty(max_payload_bytes, self.alpha_only_dim)
        nn.init.normal_(fr_b, mean=0.0, std=(1.0 / self.alpha_only_dim) ** 0.5)
        nn.init.normal_(fr_l, mean=0.0, std=(1.0 / self.alpha_only_dim) ** 0.5)
        self.register_buffer("alpha_byte_frozen", fr_b, persistent=True)
        self.register_buffer("alpha_length_frozen", fr_l, persistent=True)
        # Deeper α-only decode (hard-v2+): α_feat → hidden 4*d → logits.
        # Still conditioned ONLY on α (JL fingerprint); no byte_decoder/skip bypass.
        self.alpha_mlp_hidden = 4 * d_model
        self.alpha_byte_proj = AlphaOnlyMLP(
            self.alpha_only_dim,
            max_payload_bytes * 256,
            self.alpha_mlp_hidden,
            bias_out=False,
        )
        self.alpha_length_proj = AlphaOnlyMLP(
            self.alpha_only_dim,
            max_payload_bytes,
            self.alpha_mlp_hidden,
            bias_out=False,
        )
        # mix_w ∈ [obl_mix_floor, 1] applied to FROZEN branch — CE cannot shut off.
        self.obl_mix = nn.Parameter(torch.tensor(0.0))
        self.obl_mix_floor = 0.35
        self.obl_gate = nn.Parameter(torch.tensor(1.0))  # trainable additive scale
        self.obl_floor = 0.5
        self.field_obligatory_readout = False  # enabled via AtomNativeModel flag
        self.field_obligatory_hard = False  # hard-v2: floor=1.0 + freeze bypass + α-proj
        self._soft_obl_mix_floor = 0.35
        self.state_norm = nn.LayerNorm(d_model)
        self.byte_decoder = nn.Linear(d_model, max_payload_bytes * 256)
        self.length_decoder = nn.Linear(d_model, max_payload_bytes)
        self._init_field_readout()
        self._init_obligatory_readout()
        # Binary mask over byte alphabet (1=printable UTF-8 / Latin-1 friendly).
        # Used by decode bias and by printable_aux train loss on hard logits.
        _pm = torch.zeros(256)
        for value in self._PRINTABLE:
            _pm[value] = 1.0
        self.register_buffer("printable_mask", _pm, persistent=False)
        # Atom-payload production (hard path): per-atom α-local embed + copy-bias
        # (no mean-pool / global-hist). NOT byte_decoder bypass.
        self.byte_embed = nn.Embedding(256, d_model)
        self.payload_alpha_gate = nn.Linear(self.alpha_only_dim, d_model, bias=False)
        self.payload_to_byte = nn.Linear(d_model, max_payload_bytes * 256, bias=False)
        self.payload_to_length = nn.Linear(d_model, max_payload_bytes, bias=False)
        self.payload_copy_scale = nn.Parameter(torch.tensor(1.0))
        self.payload_enabled = True
        self._init_payload_production()
        self.last_atom_readout = bool(last_atom_readout)
        self.last_atom = None
        # CLI --last-atom-scale S (default 1.0): fixed gain on la_byte (hard mix).
        # S=1.0 = current behavior. Opt-in adaptive: see last_atom_scale_adaptive.
        self.last_atom_scale = 1.0
        self.last_atom_scale_adaptive = False
        self.last_atom_scale_min = 0.0
        self.last_atom_scale_max = 1.0
        # Opt-in: reinject full field feats (mean‖std‖persist‖atom.r) into hard
        # logits via the existing bias-free field_byte_skip (hard early-return
        # otherwise discards that path). Default 0 = identical to current hard.
        self.field_hard_couple_scale = 0.0
        # CLI --last-atom-dropout P (default 0.0 = exact current behavior, no RNG
        # draw). Training only (self.training AND grad enabled): each forward =
        # one byte-tick position; with prob P the la_byte contribution is zeroed
        # (no inverted scaling) so CE must be carried by the field path.
        # Eval / generation / no-grad predictions (efference own-byte pick) ignore P.
        self.last_atom_dropout = 0.0
        # Runtime diagnostic only (not persisted): force la_byte = 0 in any mode.
        self.last_atom_force_off = False
        if self.last_atom_readout:
            if max_payload_bytes != 1:
                raise ValueError("last_atom_readout requires max_payload_bytes=1 (byte-tick)")
            self.last_atom = LastAtomReadout(d_model)


    def effective_last_atom_scale(
        self,
        alpha_term: "torch.Tensor",
        la_raw: "torch.Tensor",
        eps: float = 1e-8,
    ) -> "torch.Tensor":
        """Return S_eff for hard-mix last_atom contribution.

        Fixed mode: scalar last_atom_scale (default 1.0).
        Adaptive mode (opt-in): S_eff = clamp(α_rms / (la_rms + eps), S_min, S_max)
        with RMS detached so the scale itself is not a gradient pathway.
        """
        if not bool(getattr(self, "last_atom_scale_adaptive", False)):
            return torch.tensor(
                float(getattr(self, "last_atom_scale", 1.0)),
                device=la_raw.device,
                dtype=la_raw.dtype,
            )
        alpha_rms = alpha_term.detach().pow(2).mean().sqrt()
        la_rms = la_raw.detach().pow(2).mean().sqrt()
        s_min = float(getattr(self, "last_atom_scale_min", 0.0))
        s_max = float(getattr(self, "last_atom_scale_max", 1.0))
        s_eff = alpha_rms / (la_rms + float(eps))
        return s_eff.clamp(min=s_min, max=s_max)

    def _init_payload_production(self) -> None:
        """Small init so payload branch is audible but does not swamp frozen α."""
        with torch.no_grad():
            nn.init.normal_(self.byte_embed.weight, mean=0.0, std=0.02)
            nn.init.zeros_(self.payload_alpha_gate.weight)
            # Slight identity-ish open: gate starts near 0.5 after sigmoid(0).
            nn.init.normal_(self.payload_to_byte.weight, mean=0.0, std=0.02)
            nn.init.normal_(self.payload_to_length.weight, mean=0.0, std=0.02)
            self.payload_copy_scale.fill_(1.0)

    def _init_field_readout(self) -> None:
        """Identity-on-mean + std residual; seeded Xavier skips. Safe for fresh + migrate."""
        with torch.no_grad():
            self.field_to_state.weight.zero_()
            eye = torch.eye(self.d_model, device=self.field_to_state.weight.device)
            self.field_to_state.weight[:, : self.d_model].copy_(eye)
            self.field_to_state.weight[:, self.d_model : 2 * self.d_model].copy_(0.5 * eye)
            self.field_to_state.bias.zero_()
            # Deterministic skip init so migrated checkpoints probe reproducibly.
            gen = torch.Generator(device="cpu")
            gen.manual_seed(20260917)
            for module in (self.field_byte_skip, self.field_length_skip):
                w = module.weight
                fan_in = w.shape[1]
                bound = (6.0 / (fan_in + w.shape[0])) ** 0.5
                w.copy_(
                    torch.empty_like(w, device="cpu")
                    .uniform_(-bound, bound, generator=gen)
                    .to(device=w.device, dtype=w.dtype)
                )
            self.field_gate.fill_(2.0)
            self.skip_gate.fill_(1.0)
            # Small open so S>0 is audible but does not swamp α-MLP at tip load.
            nn.init.normal_(self.field_hard_couple_byte.weight, mean=0.0, std=0.02)
            nn.init.normal_(self.field_hard_couple_length.weight, mean=0.0, std=0.02)

    def _init_obligatory_readout(self) -> None:
        """Seed JL/frozen buffers + trainable adapter (deterministic)."""
        with torch.no_grad():
            gen = torch.Generator(device="cpu")
            gen.manual_seed(20260918)
            jl = torch.empty_like(self.alpha_jl, device="cpu")
            jl.normal_(generator=gen, mean=0.0, std=(1.0 / self._alpha_flat_dim) ** 0.5)
            self.alpha_jl.copy_(jl.to(device=self.alpha_jl.device, dtype=self.alpha_jl.dtype))
            for buf_name, fan_out in (
                ("alpha_byte_frozen", self.alpha_byte_frozen.shape[0]),
                ("alpha_length_frozen", self.alpha_length_frozen.shape[0]),
            ):
                buf = getattr(self, buf_name)
                tmp = torch.empty_like(buf, device="cpu")
                tmp.normal_(
                    generator=gen,
                    mean=0.0,
                    std=(1.0 / self.alpha_only_dim) ** 0.5,
                )
                buf.copy_(tmp.to(device=buf.device, dtype=buf.dtype))
            self._init_alpha_trainable_mlp(generator=gen)
            self.obl_gate.fill_(1.0)
            self.obl_mix.fill_(0.0)

    def _init_alpha_trainable_mlp(
        self, generator: torch.Generator | None = None
    ) -> None:
        """Xavier-uniform init for α-only MLP adapters (does not touch frozen JL/maps)."""
        if generator is None:
            generator = torch.Generator(device="cpu")
            generator.manual_seed(20260918)
        with torch.no_grad():
            for module in (self.alpha_byte_proj, self.alpha_length_proj):
                for layer in (module.fc1, module.fc2):
                    w = layer.weight
                    fan_in = w.shape[1]
                    bound = (6.0 / (fan_in + w.shape[0])) ** 0.5
                    w.copy_(
                        torch.empty_like(w, device="cpu")
                        .uniform_(-bound, bound, generator=generator)
                        .to(device=w.device, dtype=w.dtype)
                    )
                    if layer.bias is not None:
                        layer.bias.zero_()

    def alpha_only_features(self, alpha: torch.Tensor) -> torch.Tensor:
        """JL fingerprint of full flattened α — no persist/atom_r (zeros stay zeros).

        mean‖std alone collapses inter-prompt cosine (~0.93) while full α sits
        near ~0.71; a fixed JL map preserves that separation without new attention.
        """
        flat = alpha.reshape(-1)
        # Match buffer device/dtype; pad/trim if shape drifts.
        if flat.numel() < self._alpha_flat_dim:
            flat = F.pad(flat, (0, self._alpha_flat_dim - flat.numel()))
        elif flat.numel() > self._alpha_flat_dim:
            flat = flat[: self._alpha_flat_dim]
        flat = flat.to(device=self.alpha_jl.device, dtype=self.alpha_jl.dtype)
        return self.alpha_jl @ flat

    def obligatory_scale(self) -> torch.Tensor:
        """softplus(obl_gate) + floor — additive residual cannot reach zero."""
        return F.softplus(self.obl_gate) + float(self.obl_floor)

    def obligatory_mix_weight(self) -> torch.Tensor:
        """Convex mix weight ∈ [obl_mix_floor, 1] — CE cannot shut α-branch off."""
        floor = float(self.obl_mix_floor)
        return floor + (1.0 - floor) * torch.sigmoid(self.obl_mix)

    def payload_produce(
        self,
        atom_payloads: list[bytes] | None,
        a_hat: torch.Tensor,
        a_only: torch.Tensor,
        *,
        atom_energies: list[float | torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Per-atom α-local payload embed + copy-bias (no mean-pool / global-hist).

        Each living atom contributes from its own payload, gated by similarity of
        that atom's embed with the α-derived gate and by atom energy. Softmax over
        atoms keeps the mixture local to α so distinct prompt atoms diversify
        surface logits. Returns (byte_logits [L,256], length_logits [L]).
        Vanishes when α is exactly zero or no non-empty payloads.
        No Transformer QKV; no byte_decoder reopen.
        """
        device = a_hat.device
        dtype = a_hat.dtype
        L = self.max_payload_bytes
        zero_b = torch.zeros(L, 256, device=device, dtype=dtype)
        zero_l = torch.zeros(L, device=device, dtype=dtype)
        raw = list(atom_payloads or [])
        energy_raw = list(atom_energies) if atom_energies is not None else [1.0] * len(raw)
        if len(energy_raw) < len(raw):
            energy_raw = energy_raw + [1.0] * (len(raw) - len(energy_raw))

        embeds: list[torch.Tensor] = []
        lengths: list[int] = []
        local_hists: list[torch.Tensor] = []
        energies: list[torch.Tensor] = []
        kept_payloads: list[bytes] = []
        for payload, energy in zip(raw, energy_raw):
            payload_b = bytes(payload) if payload else b""
            ids = list(payload_b[:L])
            if not ids:
                continue
            kept_payloads.append(payload_b)
            lengths.append(len(ids))
            idx = torch.tensor(ids, device=device, dtype=torch.long)
            embeds.append(self.byte_embed(idx).mean(dim=0))
            hist_i = torch.zeros(256, device=device, dtype=dtype)
            hist_i.scatter_add_(
                0,
                idx,
                torch.ones(idx.numel(), device=device, dtype=dtype),
            )
            local_hists.append(hist_i)
            if torch.is_tensor(energy):
                e_val = energy.detach().to(device=device, dtype=dtype).reshape(-1).mean()
            else:
                e_val = torch.tensor(float(energy), device=device, dtype=dtype)
            energies.append(F.softplus(e_val).clamp_min(1e-6))
        if not embeds:
            return zero_b, zero_l

        gate = torch.sigmoid(self.payload_alpha_gate(a_hat))
        # α-local scores: payload-embed alignment with α-gate × energy.
        scores = []
        for embed, energy in zip(embeds, energies):
            sim = (embed * gate).sum()
            scores.append(sim + torch.log(energy))
        score_t = torch.stack(scores, dim=0)
        # Soft pool for embed MLP path; hard α-local winner for copy-bias so one
        # living atom's payload owns surface bytes (no mean-pool / global-hist).
        pool_w = F.softmax(score_t, dim=0)
        copy_w = torch.zeros_like(score_t)
        copy_w[int(score_t.argmax().item())] = 1.0

        pooled = torch.stack(
            [pool_w[i] * embeds[i] for i in range(len(embeds))], dim=0
        ).sum(dim=0)
        gated = pooled * gate
        byte_logits = self.payload_to_byte(gated).view(L, 256)
        length_logits = self.payload_to_length(gated)

        # Architectural gain: single-atom copy must stay audible vs frozen α RMS (~1).
        copy = F.softplus(self.payload_copy_scale) * 4.0
        # Position-aligned copy only. No unigram broadcast across span slots.
        for i, payload in enumerate(kept_payloads):
            w = copy_w[i]
            for pos, b in enumerate(payload[:L]):
                byte_logits[pos, int(b)] = byte_logits[pos, int(b)] + copy * w
        if lengths:
            # Length prior from sharp α-local winner length (not uniform mean).
            mean_len = float(
                sum(copy_w[i].detach() * float(lengths[i]) for i in range(len(lengths)))
                / max(float(copy_w.detach().sum().item()), 1e-8)
            )
            li = max(0, min(L - 1, int(round(mean_len)) - 1))
            length_logits = length_logits.clone()
            length_logits[li] = length_logits[li] + copy

        # Presence gate: exact α=0 ⇒ off (hard zeroing); nonzero α ⇒ full-strength
        # payload branch (do not crush copy-bias by tiny JL RMS).
        alive = (a_only.detach().pow(2).sum() > 0).to(dtype=dtype)
        return alive * byte_logits, alive * length_logits

    def set_obligatory_hard(self, enabled: bool) -> None:
        """Hard-v2 obligatory: mix floor 1.0; freeze non-α bypass; α-proj + payload.

        Forward: logits = frozen_α_map(α) + cap * (scale*α_MLP(α) + payload_prod)
        with cap≤1 so trainable RMS cannot exceed frozen RMS.
        Decoder / field skip / gates stay zeroed+frozen; CE may train
        alpha_*_proj MLP (+ obl_gate) and atom-payload production. Field dynamics
        still train. Payload path is atom-native (not byte_decoder reopen).
        """
        enabled = bool(enabled)
        self.field_obligatory_hard = enabled
        if enabled:
            self.field_obligatory_readout = True
            self.obl_mix_floor = 1.0
            self._freeze_non_alpha_bypass(True)
        else:
            self.obl_mix_floor = float(getattr(self, "_soft_obl_mix_floor", 0.35))
            self._freeze_non_alpha_bypass(False)

    def _freeze_non_alpha_bypass(self, freeze: bool) -> None:
        """Zero + freeze decoder/skip/gates that can bypass α on the CE path."""
        modules = (
            self.field_to_state,
            self.field_byte_skip,
            self.field_length_skip,
            self.byte_decoder,
            self.length_decoder,
        )
        scalars = (self.field_gate, self.skip_gate, self.obl_mix)
        with torch.no_grad():
            if freeze:
                for module in modules:
                    for param in module.parameters():
                        param.zero_()
                        param.requires_grad_(False)
                for param in scalars:
                    param.zero_()
                    param.requires_grad_(False)
            else:
                for module in modules:
                    for param in module.parameters():
                        param.requires_grad_(True)
                for param in scalars:
                    param.requires_grad_(True)

    def field_features(
        self,
        alpha: torch.Tensor,
        persistent_state: torch.Tensor | None = None,
        atom_r: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Build spectral living-field features from α (+ consolidation, atom)."""
        if alpha.ndim == 1:
            alpha = alpha.unsqueeze(0)
        mean = alpha.mean(dim=0)
        std = alpha.std(dim=0, unbiased=False)
        if persistent_state is None:
            persist = torch.zeros_like(mean)
        else:
            persist = (
                persistent_state.mean(dim=0)
                if persistent_state.ndim > 1
                else persistent_state
            )
            persist = persist.to(device=mean.device, dtype=mean.dtype).reshape(-1)
            if persist.numel() != mean.numel():
                persist = persist[: mean.numel()]
                if persist.numel() < mean.numel():
                    persist = F.pad(persist, (0, mean.numel() - persist.numel()))
        if atom_r is None:
            atom = torch.zeros_like(mean)
        else:
            atom = atom_r.mean(dim=0) if atom_r.ndim > 1 else atom_r
            atom = atom.to(device=mean.device, dtype=mean.dtype).reshape(-1)
            if atom.numel() != mean.numel():
                atom = atom[: mean.numel()]
                if atom.numel() < mean.numel():
                    atom = F.pad(atom, (0, mean.numel() - atom.numel()))
        return torch.cat([mean, std, persist, atom], dim=-1)

    def field_conditioned_state(
        self,
        alpha: torch.Tensor,
        persistent_state: torch.Tensor | None = None,
        atom_r: torch.Tensor | None = None,
        legacy_state: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """d_model state where spectral α dominates (gated residual w/ legacy)."""
        feats = self.field_feat_norm(
            self.field_features(alpha, persistent_state, atom_r)
        )
        field_state = self.field_to_state(feats)
        gate = torch.sigmoid(self.field_gate)
        if legacy_state is None:
            return field_state
        legacy = legacy_state.mean(dim=0) if legacy_state.ndim > 1 else legacy_state
        legacy = legacy.to(device=field_state.device, dtype=field_state.dtype)
        return gate * field_state + (1.0 - gate) * legacy

    def forward(
        self,
        state: torch.Tensor | None = None,
        *,
        alpha: torch.Tensor | None = None,
        persistent_state: torch.Tensor | None = None,
        atom_r: torch.Tensor | None = None,
        atom_payloads: list[bytes] | None = None,
        atom_energies: list[float | torch.Tensor] | None = None,
        last_phi: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor]:
        """Decode from living field α when provided; else legacy mean-state path.

        When ``field_obligatory_readout`` is on, an α-only residual (mean‖std)
        is always added into the main CE logits with a non-zero floor scale so
        persist/atom_r / decoder bias cannot fully bypass living α.

        Under ``field_obligatory_hard``, living ``atom_payloads`` drive an
        α-local per-atom payload branch (no mean-pool / global-hist) alongside α-MLP.
        """
        if alpha is not None:
            feats = self.field_feat_norm(
                self.field_features(alpha, persistent_state, atom_r)
            )
            field_state = self.field_to_state(feats)
            gate = torch.sigmoid(self.field_gate)
            if state is None:
                combined = field_state
            else:
                legacy = state.mean(dim=0) if state.ndim > 1 else state
                legacy = legacy.to(device=field_state.device, dtype=field_state.dtype)
                combined = gate * field_state + (1.0 - gate) * legacy
            h = self.state_norm(combined)
            byte_logits = self.byte_decoder(h).view(self.max_payload_bytes, 256)
            length_logits = self.length_decoder(h)
            skip = F.softplus(self.skip_gate)
            # When obligatory: skip uses α-only feats (padded) so shared persist/
            # atom_r cannot substitute for α inside the skip path either.
            if self.field_obligatory_readout:
                a_only = self.alpha_only_features(alpha)
                a_rms = a_only.pow(2).mean().sqrt().clamp_min(1e-8)
                a_hat = a_only / a_rms
                skip_in = F.pad(a_hat, (0, self.field_feat_dim - self.alpha_only_dim))
                skip_in = self.field_feat_norm(skip_in)
            else:
                skip_in = feats
            byte_logits = byte_logits + skip * self.field_byte_skip(skip_in).view(
                self.max_payload_bytes, 256
            )
            length_logits = length_logits + skip * self.field_length_skip(skip_in)
            if self.field_obligatory_readout:
                a_only = self.alpha_only_features(alpha)
                a_rms = a_only.pow(2).mean().sqrt().clamp_min(1e-8)
                a_hat = a_only / a_rms
                # Frozen branch (CE cannot collapse) + trainable adapter.
                frozen_byte = (self.alpha_byte_frozen @ a_hat).view(
                    self.max_payload_bytes, 256
                )
                frozen_len = self.alpha_length_frozen @ a_hat
                if getattr(self, "field_obligatory_hard", False):
                    # Hard speech mix (ARCHITECTURE.md): trainable α-MLP +
                    # 0.3 frozen JL + optional payload branch. No RMS cap on
                    # the trainable branch: the cap pinned cap≈0.01 and crushed
                    # trainable gradients ~100x (measured), blocking learning.
                    # Payload stays outside any cap; α=0 still zeros payload
                    # via unclamped a_scale. No decoder/skip reopen.
                    train_byte = self.alpha_byte_proj(a_hat).view(
                        self.max_payload_bytes, 256
                    )
                    train_len = self.alpha_length_proj(a_hat)
                    if getattr(self, "payload_enabled", True):
                        pay_byte, pay_len = self.payload_produce(
                            atom_payloads,
                            a_hat,
                            a_only,
                            atom_energies=atom_energies,
                        )
                    else:
                        pay_byte = torch.zeros_like(frozen_byte)
                        pay_len = torch.zeros_like(frozen_len)
                    obl = self.obligatory_scale()
                    la_byte = torch.zeros_like(frozen_byte)
                    if getattr(self, "last_atom", None) is not None:
                        last_b, prev_b = _tail_bytes(atom_payloads)
                        phi_vec = last_phi
                        if phi_vec is None:
                            phi_vec = torch.zeros(
                                self.d_model,
                                device=frozen_byte.device,
                                dtype=frozen_byte.dtype,
                            )
                        la_raw = self.last_atom(last_b, prev_b, phi_vec).unsqueeze(0)
                        # Equalize last_atom vs α-MLP (obl*train_byte) when adaptive.
                        scale = self.effective_last_atom_scale(obl * train_byte, la_raw)
                        la_byte = scale * la_raw
                        if bool(getattr(self, "last_atom_force_off", False)):
                            la_byte = torch.zeros_like(la_byte)
                        else:
                            p_drop = float(getattr(self, "last_atom_dropout", 0.0) or 0.0)
                            if (
                                p_drop > 0.0
                                and self.training
                                and torch.is_grad_enabled()
                                and (p_drop >= 1.0 or float(torch.rand(())) < p_drop)
                            ):
                                la_byte = torch.zeros_like(la_byte)
                    # Field→hard-decode couple: reinject full spectral feats
                    # (incl. persist‖atom.r) that the hard early-return else drops.
                    # Uses dedicated W (field_byte_skip is zeroed+frozen under hard).
                    couple_byte = torch.zeros_like(frozen_byte)
                    couple_len = torch.zeros_like(frozen_len)
                    s_couple = float(getattr(self, "field_hard_couple_scale", 0.0) or 0.0)
                    if s_couple > 0.0:
                        # `feats` already = LN([mean‖std‖persist‖atom.r]) above.
                        couple_byte = s_couple * self.field_hard_couple_byte(feats).view(
                            self.max_payload_bytes, 256
                        )
                        couple_len = s_couple * self.field_hard_couple_length(feats)
                    return {
                        "byte_logits": obl * train_byte + 0.3 * frozen_byte + pay_byte + la_byte + couple_byte,
                        "length_logits": obl * train_len + 0.3 * frozen_len + pay_len + couple_len,
                    }
                train_byte = self.alpha_byte_proj(a_hat).view(self.max_payload_bytes, 256)
                train_len = self.alpha_length_proj(a_hat)
                mix = self.obligatory_mix_weight()
                # Hard mix uses FROZEN map so inter-prompt α structure reaches logits.
                byte_logits = (1.0 - mix) * byte_logits + mix * frozen_byte
                length_logits = (1.0 - mix) * length_logits + mix * frozen_len
                obl = self.obligatory_scale()
                byte_logits = byte_logits + obl * train_byte
                length_logits = length_logits + obl * train_len
            return {"byte_logits": byte_logits, "length_logits": length_logits}

        if state is None:
            raise ValueError("AtomSurfaceHead.forward requires state or alpha")
        if state.ndim > 1:
            state = state.mean(dim=0)
        state = self.state_norm(state)
        byte_logits = self.byte_decoder(state).view(self.max_payload_bytes, 256)
        length_logits = self.length_decoder(state)
        return {"byte_logits": byte_logits, "length_logits": length_logits}

    def _apply_printable_bias(self, byte_logits: torch.Tensor, strength: float = 2.0) -> torch.Tensor:
        """Softly discourage control/non-text bytes without hard-masking UTF-8."""
        if strength <= 0:
            return byte_logits
        mask = getattr(self, "printable_mask", None)
        if mask is None:
            penalty = torch.full(
                (256,),
                -strength,
                dtype=byte_logits.dtype,
                device=byte_logits.device,
            )
            for value in self._PRINTABLE:
                penalty[value] = 0.0
        else:
            # printable→0, non-printable→-strength (vectorized)
            penalty = (1.0 - mask.to(device=byte_logits.device, dtype=byte_logits.dtype)) * (-strength)
        # Keep NUL strongly suppressed; allow tab/LF/CR via printable set.
        penalty = penalty.clone()
        penalty[0] = -strength * 2.0
        return byte_logits + penalty

    def printable_mass(self, byte_logits: torch.Tensor) -> torch.Tensor:
        """Per-position softmax mass on printable UTF-8 / Latin-1 friendly bytes."""
        probs = byte_logits.softmax(dim=-1)
        mask = self.printable_mask.to(device=probs.device, dtype=probs.dtype)
        return (probs * mask).sum(dim=-1)

    def decode(
        self,
        output: dict[str, torch.Tensor],
        temperature: float = 1.0,
        top_k: int | None = None,
        deterministic: bool = True,
        prefer_printable: bool = True,
    ) -> bytes:
        if temperature <= 0:
            raise ValueError("temperature must be positive")
        length_logits = output["length_logits"] / temperature
        if deterministic:
            length = int(length_logits.argmax().item()) + 1
        else:
            k = length_logits.numel()
            if top_k and top_k > 0:
                k = min(top_k, k)
            length_values, length_indices = torch.topk(length_logits, k)
            # Guard against non-finite / zero-mass softmax from extreme logits.
            probs = length_values.softmax(dim=-1)
            if not torch.isfinite(probs).all() or float(probs.sum()) <= 0:
                length = int(length_logits.argmax().item()) + 1
            else:
                length = int(length_indices[torch.multinomial(probs, 1)].item()) + 1
        length = max(1, min(length, self.max_payload_bytes))
        byte_logits = output["byte_logits"][:length] / temperature
        # Hard α-MLP path: always prefer printable (chat_atom_native contract).
        if prefer_printable or getattr(self, "field_obligatory_hard", False):
            strength = 2.5 if getattr(self, "field_obligatory_hard", False) else 2.0
            byte_logits = self._apply_printable_bias(byte_logits, strength=strength)
        if deterministic:
            values = byte_logits.argmax(dim=-1)
        else:
            if top_k and top_k > 0 and top_k < byte_logits.shape[-1]:
                values, indices = torch.topk(byte_logits, top_k, dim=-1)
                probs = values.softmax(dim=-1)
                if not torch.isfinite(probs).all():
                    values = byte_logits.argmax(dim=-1)
                else:
                    picks = torch.multinomial(probs, 1)
                    values = indices.gather(-1, picks).squeeze(-1)
            else:
                probs = byte_logits.softmax(dim=-1)
                if not torch.isfinite(probs).all():
                    values = byte_logits.argmax(dim=-1)
                else:
                    values = torch.multinomial(probs, num_samples=1).squeeze(-1)
        return bytes(int(value) for value in values.tolist())


def field_leak_spectrum_rates(
    n_modes: int, lmin: float, lmax: float, slow_frac: float = 1.0
) -> torch.Tensor:
    """Per-mode leak rates λ_m for ``--field-leak-spectrum LMIN LMAX``.

    λ is a geometric ladder from LMIN (slow memory) to LMAX (fast / local byte
    context) over the ``n_modes`` modes.  The ladder is assigned by the *only*
    mode-dependent term of the linear toroidal dynamics, the rotation rate
    ``sin(mode_phase_m)`` (linear rate ``sin φ_m + energy_decay - 1``): the mode
    with the least intrinsic decay gets LMIN, the most-damped mode gets LMAX
    (stable argsort, ties -> lower index first).  Because sin φ is smooth on the
    ring, λ is a smooth V-shaped profile around the torus (slowest at m≈n/4,
    fastest at m≈3n/4), so the neighbour coupling (phase_sync Laplacian) only
    mixes modes of similar timescale; interleaving would put a fast mode next to
    every slow one and the Laplacian would drain the slow modes.
    ``slow_frac`` F in (0, 1]: only the k = max(2, round(F·n)) lowest-decay
    modes carry the ladder LMIN → LMAX; all other modes sit at LMAX.  F = 1
    (default) = the ladder spans every mode.
    """
    n = int(n_modes)
    slow_frac = float(slow_frac)
    if not 0.0 < slow_frac <= 1.0:
        raise ValueError("slow_frac must be in (0, 1]")
    k = n if slow_frac >= 1.0 else min(n, max(2, int(round(slow_frac * n))))
    lmin = float(lmin)
    lmax = float(lmax)
    phase = torch.arange(n, dtype=torch.float64) * (2.0 * torch.pi / n)
    growth = torch.sin(phase)
    # stable: descending growth, ties broken by mode index
    order = sorted(range(n), key=lambda m: (-round(float(growth[m]), 12), m))
    rank = torch.empty(n, dtype=torch.float64)
    for r, m in enumerate(order):
        rank[m] = float(r)
    frac = (rank / max(1, k - 1)).clamp(max=1.0)
    lam = lmin * (lmax / lmin) ** frac
    return lam.to(torch.float32)


def field_leak_slow_block_modes(n_modes: int, frac: float) -> list[int]:
    """Modes of ``--field-leak-slow-block F``: one CONTIGUOUS block of
    k = max(2, round(F·n)) modes centred on m = n/2 (sin φ = 0, the ring point where
    the intrinsic RK4 rate dt·(sin φ_m + energy_decay − 1) is ~0).  Contiguous so the
    phase-sync Laplacian only drains the block at its two edges (interleaved slow/fast
    modes would be drained every tick)."""
    n = int(n_modes)
    frac = float(frac)
    if not 0.0 < frac < 1.0:
        raise ValueError("field_leak_slow_block fraction must be in (0, 1)")
    k = min(n - 1, max(2, int(round(frac * n))))
    start = n // 2 - k // 2
    return [(start + i) % n for i in range(k)]


def slow_mode_indices(rates: torch.Tensor | None, tol: float = 1e-6) -> list[int]:
    """Modes that are slower than the fast leak LMAX (the ``--field-leak-spectrum``
    ladder minus its top rung).  Empty when no spectrum / every mode at LMAX."""
    if rates is None:
        return []
    r = rates.detach().reshape(-1).to(torch.float64)
    lmax = float(r.max())
    return [int(m) for m in range(r.numel()) if float(r[m]) < lmax - tol]


class SlowModeReadout(nn.Module):
    """Own readout for the slow modes of ``--field-leak-spectrum`` (opt-in).

    ``--slow-mode-readout`` (docs/SLOW_MODE_READOUT.md).  Input = the slow modes'
    state α[slow_idx] (n_slow × d_model), flattened and divided by its own RMS
    (so the slow content is read at unit scale whatever its amplitude, and never
    has to compete with the fast modes inside the 64-d a_hat JL fingerprint):

        x        = (ŝ − μ) / σ   per-dim running standardization (buffers)
        z        = tanh(W_t · x + b_t)                (rank-r trunk)
        Δlogits  = W_o · sg(z)      W_o zero-init  -> step-0 logits unchanged
        mem_K(z) = W_K · z + b_K    one 256-way head per memory lag K (aux only)

    Next-byte CE trains only W_o (z is a stop-gradient feature there); the trunk
    is trained only by the memory aux (``--slow-memory-loss``).
    μ/σ: ~92% of ŝ's energy is a static offset (measured on the 4996k weights), so
    unstandardized input let the head learn only that offset (the trunk saturated,
    z became a constant).  μ/σ are running per-dim statistics updated on training
    ticks only (cumulative average for the first ``STAT_WINDOW`` ticks, then an EMA
    with that window), saved with the checkpoint and frozen at eval.

    Still α-only (no attention, no token lookup).  Built under a forked RNG with a
    dedicated generator so enabling it never shifts the global RNG stream.
    """

    SEED = 20260929
    STAT_WINDOW = 2000
    VAR_FLOOR = 1e-4

    def __init__(self, n_slow: int, d_model: int, rank: int, out_dim: int,
                 lags: Iterable[int] = ()) -> None:
        super().__init__()
        self.n_slow = int(n_slow)
        self.d_model = int(d_model)
        self.rank = int(rank)
        self.out_dim = int(out_dim)
        if self.n_slow < 1 or self.rank < 1:
            raise ValueError("SlowModeReadout needs n_slow >= 1 and rank >= 1")
        self.in_dim = self.n_slow * self.d_model
        self.register_buffer("in_mean", torch.zeros(self.in_dim))
        self.register_buffer("in_var", torch.ones(self.in_dim))
        self.register_buffer("n_seen", torch.zeros((), dtype=torch.float64))
        with torch.random.fork_rng(devices=[]):
            self.trunk = nn.Linear(self.in_dim, self.rank)
            self.to_logits = nn.Linear(self.rank, self.out_dim, bias=False)
            self.mem_heads = nn.ModuleDict()
        gen = torch.Generator(device="cpu")
        gen.manual_seed(self.SEED)
        with torch.no_grad():
            self._xavier(self.trunk.weight, gen)
            self.trunk.bias.zero_()
            self.to_logits.weight.zero_()
        for k in lags:
            self.add_lag(int(k))

    @staticmethod
    def _xavier(w: torch.Tensor, gen: torch.Generator) -> None:
        bound = (6.0 / (w.shape[0] + w.shape[1])) ** 0.5
        w.copy_(torch.empty(w.shape).uniform_(-bound, bound, generator=gen).to(w))

    def add_lag(self, k: int) -> None:
        key = str(int(k))
        if key in self.mem_heads:
            return
        with torch.random.fork_rng(devices=[]):
            head = nn.Linear(self.rank, 256)
        gen = torch.Generator(device="cpu")
        gen.manual_seed(self.SEED + 1 + int(k))
        with torch.no_grad():
            self._xavier(head.weight, gen)
            head.bias.zero_()
        head.to(device=self.trunk.weight.device, dtype=self.trunk.weight.dtype)
        self.mem_heads[key] = head

    def set_lags(self, lags: Iterable[int]) -> None:
        want = {str(int(k)) for k in lags}
        for key in list(self.mem_heads.keys()):
            if key not in want:
                del self.mem_heads[key]
        for k in sorted(int(x) for x in want):
            self.add_lag(k)

    def features(self, slow: torch.Tensor, update_stats: bool = False) -> torch.Tensor:
        flat = slow.reshape(-1).to(device=self.trunk.weight.device, dtype=self.trunk.weight.dtype)
        x = flat / flat.pow(2).mean().sqrt().clamp_min(1e-8)
        if update_stats:
            with torch.no_grad():
                self.n_seen.add_(1.0)
                m = 1.0 / min(float(self.n_seen), float(self.STAT_WINDOW))
                xd = x.detach()
                delta = xd - self.in_mean
                self.in_mean.add_(m * delta)
                self.in_var.mul_(1.0 - m).add_(m * delta * (xd - self.in_mean))
                if float(self.n_seen) <= 1.0:
                    self.in_var.fill_(1.0)
        x = (x - self.in_mean) / self.in_var.clamp_min(self.VAR_FLOOR).sqrt()
        return torch.tanh(self.trunk(x))


class FieldAmplitudeController:
    """Bound field amplitude without changing the toroidal core equations.

    The controller is an opt-in safety policy around the existing adapter.  It
    rescales only when the field RMS exceeds ``max_rms``; directions and small
    amplitudes are left untouched.  Dynamics, interactions and the eight atom
    properties remain the same.
    """

    def __init__(self, max_rms: float | None = None, eps: float = 1e-6) -> None:
        if max_rms is not None and max_rms <= 0:
            raise ValueError("max_rms must be positive or None")
        self.max_rms = max_rms
        self.eps = eps

    def __call__(self, field: torch.Tensor) -> tuple[torch.Tensor, dict[str, float]]:
        rms_before = torch.sqrt(field.pow(2).mean() + self.eps)
        if self.max_rms is None:
            scale = torch.ones_like(rms_before)
            stabilized = field
        else:
            scale = torch.minimum(
                torch.ones_like(rms_before),
                field.new_tensor(self.max_rms) / (rms_before + self.eps),
            )
            stabilized = field * scale
        rms_after = torch.sqrt(stabilized.pow(2).mean() + self.eps)
        return stabilized, {
            "field_rms_before": float(rms_before.detach().item()),
            "field_rms_after": float(rms_after.detach().item()),
            "field_scale": float(scale.detach().item()),
        }


class AtomNativeModel(nn.Module):
    """Use existing ATOM dynamics with an atom-native input/output contract."""

    def __init__(
        self,
        d_model: int = 8,
        n_modes: int = 8,
        n_atoms_max: int = 128,
        max_payload_bytes: int = 32,
        atomizer: Atomizer | None = None,
        field_max_rms: float | None = None,
        energy_decay_bounds: tuple[float, float] | None = None,
        enable_merge: bool = True,
        merge_coherence_threshold: float = 0.45,
        merge_energy_floor: float = 0.08,
        slow_every: int = 1,
        field_loss_weight: float = 0.0,
        field_contrast_weight: float = 0.0,
        field_contrast_margin: float = 0.55,
        field_ignorance_weight: float = 0.0,
        field_ignorance_margin: float = 0.85,
        field_ignorance_ablate_shared: bool = False,
        field_ignorance_prompt_bank: bool = False,
        field_obligatory_readout: bool = False,
        field_obligatory_hard: bool = False,
        last_atom_readout: bool = False,
        printable_aux_weight: float = 0.0,
        field_next_packet_weight: float = 0.0,
        field_leak: float = 0.0,
        field_leak_spectrum: tuple[float, float] | None = None,
        field_leak_spectrum_gain: str = "none",
        field_leak_spectrum_slow_frac: float = 1.0,
        slow_rms_rel_tol: float = 0.15,
        slow_mode_readout: bool = False,
        slow_readout_rank: int = 64,
        slow_readout_ahat: str = "exclude",
        slow_memory_weight: float = 0.0,
        slow_memory_lags: tuple[int, ...] = (16, 32, 64),
        slow_memory_target: str = "bag",
        field_bptt: int = 0,
        field_leak_learnable: bool = False,
    ) -> None:
        super().__init__()
        self.atomizer = atomizer or Atomizer(max_span_bytes=max_payload_bytes, pack_mode="linguistic")
        self.core = ToroidalFractalIntelligence(
            vocab_size=256,
            d_model=d_model,
            n_modes=n_modes,
            n_atoms_max=n_atoms_max,
        )
        self.compiler = AtomCompiler(self.atomizer.feature_dim, d_model)
        self.surface = AtomSurfaceHead(
            d_model,
            max_payload_bytes=max_payload_bytes,
            n_modes=n_modes,
            last_atom_readout=bool(last_atom_readout),
        )
        self.max_payload_bytes = max_payload_bytes
        self.field_controller = FieldAmplitudeController(field_max_rms)
        self.field_max_rms = field_max_rms
        # Opt-in per-tick field leak λ (--field-leak). 0.0 = exact legacy path
        # (no op at all).  λ>0: α ← (1-λ)·α before each atom injection, so the
        # field is a contractive, recency-weighted memory instead of an
        # integrator whose unstable modes pin it at the field_max_rms clamp.
        self.field_leak = float(field_leak)
        if not 0.0 <= self.field_leak < 1.0:
            raise ValueError("field_leak must be in [0, 1)")
        # Opt-in multi-timescale leak (--field-leak-spectrum LMIN LMAX).  None =
        # off = exact current behavior (scalar --field-leak path, or nothing).
        # When set it REPLACES the scalar leak: α_m ← (1-λ_m)·α_m per mode.
        self.field_leak_spectrum = None
        self.field_leak_spectrum_gain = "none"
        self.field_leak_spectrum_slow_frac = 1.0
        self._field_leak_vec = None
        self._field_gain_vec = None
        self.set_field_leak_spectrum(
            field_leak_spectrum, gain=field_leak_spectrum_gain,
            slow_frac=field_leak_spectrum_slow_frac,
        )
        # Always-on physical band.  The weak prior floor (1e-3) let persist
        # training drift energy_decay to ~0.001 and collapse the field.
        if energy_decay_bounds is None:
            energy_decay_bounds = DEFAULT_ENERGY_DECAY_BOUNDS
        self.energy_decay_bounds = energy_decay_bounds
        # Continuum intelligence knobs (structure + field dependence, not Θ growth).
        self.enable_merge = bool(enable_merge)
        self.merge_coherence_threshold = float(merge_coherence_threshold)
        self.merge_energy_floor = float(merge_energy_floor)
        self.core.aggregation.phase_coherence_threshold = self.merge_coherence_threshold
        self.core.aggregation.merge_energy_floor = self.merge_energy_floor
        self.slow_every = max(1, int(slow_every))
        self.field_loss_weight = float(field_loss_weight)
        self.field_contrast_weight = float(field_contrast_weight)
        self.field_contrast_margin = float(field_contrast_margin)
        self.field_ignorance_weight = float(field_ignorance_weight)
        self.field_ignorance_margin = float(field_ignorance_margin)
        self.field_ignorance_ablate_shared = bool(field_ignorance_ablate_shared)
        self.field_ignorance_prompt_bank = bool(field_ignorance_prompt_bank)
        self.field_obligatory_hard = bool(field_obligatory_hard)
        self.last_atom_readout = bool(last_atom_readout)
        self.field_obligatory_readout = bool(field_obligatory_readout) or self.field_obligatory_hard
        # Light aux: push hard α-MLP logits toward printable UTF-8 mass (no decoder bypass).
        self.printable_aux_weight = float(printable_aux_weight)
        # Predict *next* packet features from α (+ atom_r in field_features).
        self.field_next_packet_weight = float(field_next_packet_weight)
        self.surface.field_obligatory_readout = self.field_obligatory_readout
        if self.field_obligatory_hard:
            self.surface.set_obligatory_hard(True)
        else:
            self.surface.field_obligatory_hard = False
        self.slow_rms_rel_tol = float(slow_rms_rel_tol)
        self.field_probe = nn.Linear(self.surface.field_feat_dim, self.atomizer.feature_dim)
        self.field_next_probe = nn.Linear(self.surface.field_feat_dim, self.atomizer.feature_dim)
        self.merge_count_total = 0
        self._tick = 0
        self._efference_every = 0  # set by trainer CLI (--efference-every)
        self._prev_field_rms: float | None = None
        # Detached rolling α snapshots for harder field-contrast negatives
        # (mode-shuffle alone is too weak vs CE; inter-prompt collapse returns).
        self._alpha_bank: list[torch.Tensor] = []
        self._alpha_bank_max = 8
        # Distinct-prompt α bank for L_ign (filled by _refresh_prompt_alpha_bank).
        self._prompt_alpha_bank: list[torch.Tensor] = []
        self._prompt_bank_refresh_every = 256
        self._ignorance_prompts: tuple[str, ...] = DEFAULT_IGNORANCE_PROMPTS

        # The legacy encoder and legacy 256-way production head are not used
        # by this adapter.  Keep them in the core checkpoint for compatibility,
        # but make clear that no GPT-style embedding participates in training.
        for parameter in self.core.encoder.parameters():
            parameter.requires_grad_(False)
        for parameter in self.core.production.parameters():
            parameter.requires_grad_(False)
        self.core.state.alpha.requires_grad_(False)

        # Opt-in own readout for the slow spectrum modes (--slow-mode-readout,
        # docs/SLOW_MODE_READOUT.md).  Off = no module, no state, no RNG draw:
        # every code path below is skipped and behaviour is exactly unchanged.
        self.slow_mode_readout = False
        self.slow_readout = None
        self.slow_readout_rank = int(slow_readout_rank)
        self.slow_readout_ahat = "exclude"
        self.slow_memory_weight = 0.0
        self.slow_memory_lags: tuple[int, ...] = tuple(int(k) for k in slow_memory_lags)
        self.slow_memory_target = "bag"
        self._slow_idx = None
        self._slow_fast_mask = None
        self._slow_hist = bytearray()
        if slow_mode_readout:
            self.set_slow_mode_readout(
                True, rank=slow_readout_rank, ahat=slow_readout_ahat,
                memory_weight=slow_memory_weight, memory_lags=slow_memory_lags,
                memory_target=slow_memory_target,
            )
        # Opt-in truncated BPTT through the field (--field-bptt K,
        # docs/DECODE_FIX_AND_BPTT.md).  0 = off = exact legacy path (α detached
        # every tick, injection under no_grad).  K>0: while training with grad
        # enabled, α_t keeps its autograd graph to α_{t-1} and the atom injection
        # is differentiable; the trainer backprops/steps every K ticks and calls
        # detach_field_graph() at the window boundary.  Forward values are
        # identical to the legacy path (same ops, same order).
        self.field_bptt = 0
        self._bptt_alpha: torch.Tensor | None = None
        self.set_field_bptt(field_bptt)
        # Opt-in learnable per-mode leak (--field-leak-learnable): λ_m =
        # sigmoid(field_leak_logit_m), init at --field-leak.  None = off (no
        # parameter, checkpoint layout unchanged).
        self.field_leak_logit: nn.Parameter | None = None
        if field_leak_learnable:
            self.set_field_leak_learnable(True)
        # --field-leak-slow-block (opt-in, needs the learnable leak): None = off.
        self.field_leak_slow_block: dict | None = None
        self._slow_block_mask: torch.Tensor | None = None
        self._slow_block_floor: torch.Tensor | None = None

    # ------------------------------------------------------------------
    # Truncated BPTT through the field (opt-in)
    # ------------------------------------------------------------------
    def set_field_bptt(self, k: int) -> None:
        k = int(k or 0)
        if k < 0:
            raise ValueError("field_bptt must be >= 0")
        self.field_bptt = k
        self._bptt_alpha = None

    def detach_field_graph(self) -> None:
        """Window boundary: keep α's value, drop its autograd history."""
        if self._bptt_alpha is not None:
            self._bptt_alpha = self._bptt_alpha.detach()

    def _bptt_prev_alpha(self) -> torch.Tensor:
        """α_{t-1} with its graph when it is still the live field value, else a
        detached copy (after reset/load/eval ticks the graph is stale)."""
        prev = self._bptt_alpha
        st = self.core.state.alpha
        if prev is None or prev.shape != st.shape or not torch.equal(prev.detach(), st.detach()):
            return st.detach().clone()
        return prev

    def _bptt_active(self) -> bool:
        return bool(self.field_bptt > 0 and self.training and torch.is_grad_enabled())

    def set_field_leak_learnable(self, enabled: bool, init: float | None = None) -> None:
        if getattr(self, "field_leak_slow_block", None) is not None:
            self.field_leak_slow_block = None
            self._slow_block_mask = None
            self._slow_block_floor = None
        if not enabled:
            self.field_leak_logit = None
            return
        if getattr(self, "field_leak_spectrum", None) is not None:
            raise ValueError("field_leak_learnable is exclusive with field_leak_spectrum")
        lam = float(self.field_leak if init is None else init)
        lam = min(max(lam, 1e-4), 1.0 - 1e-4)
        n_modes = self.core.state.n_modes
        logit = math.log(lam / (1.0 - lam))
        self.field_leak_logit = nn.Parameter(
            torch.full((n_modes, 1), logit, dtype=torch.float32,
                       device=self.core.state.alpha.device)
        )

    def learnable_leak_rates(self) -> torch.Tensor | None:
        if self.field_leak_logit is None:
            return None
        if getattr(self, "field_leak_slow_block", None) is not None:
            return self._slow_block_rates(self.field_leak_logit)
        return torch.sigmoid(self.field_leak_logit)

    # ------------------------------------------------------------------
    # --field-leak-slow-block F [cap C] [gain ema|none]  (opt-in long memory)
    # ------------------------------------------------------------------
    SLOW_BLOCK_INIT_SIGMA = 0.25   # slow-mode logit init: λ = floor + cap·0.25
    SLOW_BLOCK_EPS = 1e-3          # net per-tick decay floor of a slow mode

    def _slow_block_rates(self, logit: torch.Tensor) -> torch.Tensor:
        cfg = self.field_leak_slow_block
        mask = self._slow_block_mask.to(device=logit.device)
        floor = self._slow_block_floor.to(device=logit.device, dtype=logit.dtype)
        sig = torch.sigmoid(logit)
        return torch.where(mask, floor + float(cfg["cap"]) * sig, sig)

    def slow_block_input_gain(self, lam: torch.Tensor | None = None) -> torch.Tensor | None:
        """Per-mode input gain g_m for the slow block ('ema': g_m = λ_m / ref, so a slow
        mode is an EMA of its input with the amplitude of a λ=ref mode instead of a 1/λ_m
        integrator that would swamp the readout / hit the RMS cap).  None = no gain."""
        cfg = getattr(self, "field_leak_slow_block", None)
        if cfg is None or cfg.get("gain", "ema") == "none":
            return None
        if lam is None:
            lam = self.learnable_leak_rates()
        mask = self._slow_block_mask.to(device=lam.device)
        return torch.where(mask, lam / float(cfg["ref"]), torch.ones_like(lam))

    def set_field_leak_slow_block(self, frac: float | None, cap: float = 0.02,
                                  gain: str = "ema", ref: float | None = None,
                                  reinit: bool = True) -> None:
        """Enable/disable the slow block (needs --field-leak-learnable).

        Slow modes (``field_leak_slow_block_modes``): λ_m = floor_m + cap·sigmoid(logit_m),
        floor_m = max(0, dt·(sin φ_m + ed_max − 1)) + 1e-3, so the net per-tick retention
        (1−λ_m)·exp(dt·(sin φ_m + ed − 1)) stays < 1 for every allowed energy_decay (no
        runaway mode) while λ can go ~100× below the learned fast leak; the optimiser can
        never push a slow mode above floor+cap.  Other modes are unchanged (sigmoid).
        ``reinit``: slow logits start at sigmoid = 0.25.  Disabling keeps every λ_m value
        (logits re-expressed in the plain sigmoid form)."""
        if self.field_leak_logit is None:
            if frac:
                raise ValueError("--field-leak-slow-block needs --field-leak-learnable")
            self.field_leak_slow_block = None
            return
        if not frac:
            if getattr(self, "field_leak_slow_block", None) is not None:
                lam = self.learnable_leak_rates().detach().clamp(1e-6, 1 - 1e-6)
                self.field_leak_slow_block = None
                with torch.no_grad():
                    self.field_leak_logit.copy_(torch.log(lam / (1 - lam)))
            self.field_leak_slow_block = None
            self._slow_block_mask = None
            self._slow_block_floor = None
            return
        n = int(self.core.state.n_modes)
        modes = field_leak_slow_block_modes(n, frac)
        cap = float(cap)
        if not 0.0 < cap < 0.5:
            raise ValueError("slow block cap must be in (0, 0.5)")
        if gain not in ("ema", "none"):
            raise ValueError("slow block gain must be 'ema' or 'none'")
        dt = float(getattr(self.core.dynamics, "dt", 0.01))
        ed_hi = float((self.energy_decay_bounds or DEFAULT_ENERGY_DECAY_BOUNDS)[1])
        phase = torch.arange(n, dtype=torch.float64) * (2.0 * math.pi / n)
        floor = (dt * (torch.sin(phase) + ed_hi - 1.0)).clamp(min=0.0) + self.SLOW_BLOCK_EPS
        mask = torch.zeros(n, 1, dtype=torch.bool)
        mask[modes, 0] = True
        was_on = getattr(self, "field_leak_slow_block", None) is not None
        self._slow_block_mask = mask
        self._slow_block_floor = floor.to(torch.float32).view(n, 1)
        self.field_leak_slow_block = {"frac": float(frac), "cap": cap, "gain": gain,
                                      "ref": float(self.field_leak if ref is None else ref) or 0.05,
                                      "modes": modes}
        if reinit and not was_on:
            s0 = self.SLOW_BLOCK_INIT_SIGMA
            with torch.no_grad():
                self.field_leak_logit[mask] = math.log(s0 / (1 - s0))

    @property
    def trainable_parameters(self) -> Iterable[nn.Parameter]:
        return (parameter for parameter in self.parameters() if parameter.requires_grad)

    def trainable_parameter_count(self) -> int:
        return sum(parameter.numel() for parameter in self.trainable_parameters)

    def stabilize_dynamics_parameters(self) -> dict[str, float]:
        """Project shared dynamics params into a physical range every step.

        ``energy_decay`` is called a decay parameter by the core.  Without a
        strong floor it can drift near 0 and wipe the field each tick
        (``decay = (energy_decay - 1) * alpha``).  Bounds are always-on
        (default ``DEFAULT_ENERGY_DECAY_BOUNDS``).  Does not alter
        ``src/toroidal/dynamics.py``.
        """
        parameter = self.core.dynamics.dynamics
        bounds = self.energy_decay_bounds or DEFAULT_ENERGY_DECAY_BOUNDS
        low, high = bounds
        if not (0.0 < low <= high):
            raise ValueError("energy_decay_bounds must satisfy 0 < low <= high")
        with torch.no_grad():
            before = float(parameter.energy_decay.detach().item())
            parameter.energy_decay.clamp_(low, high)
            parameter.coupling_scale.clamp_(0.0, 2.0)
            parameter.phase_sync.clamp_(-1.0, 1.0)
            after = float(parameter.energy_decay.detach().item())
        return {
            "coupling_scale": float(parameter.coupling_scale.detach().item()),
            "energy_decay": after,
            "phase_sync": float(parameter.phase_sync.detach().item()),
            "energy_decay_before": before,
            "energy_decay_repaired": before != after,
            "energy_decay_bounds": (float(low), float(high)),
        }

    def reset_state(self, reset_atomizer: bool = True) -> None:
        """Start a new explicit episode while preserving learned weights."""
        if getattr(self, "_bptt_alpha", None) is not None:
            self._bptt_alpha = None
        with torch.no_grad():
            self.core.state.alpha.zero_()
            self.core.state.t.zero_()
            self.core.consolidation.persistence.zero_()
            self.core.consolidation.persistence_strength.zero_()
            self.core.abstraction.abstraction_memory.zero_()
            self.core.abstraction.abstraction_usage.zero_()
            self.core.abstraction.abstraction_age.zero_()
        self.core.atoms = type(self.core.atoms)()
        self.core.energy_history = []
        self.core.consolidation_count = 0
        self.core.abstraction_count = 0
        self._tick = 0
        self._prev_field_rms = None
        self._alpha_bank = []
        if getattr(self, "slow_readout", None) is not None:
            self._slow_hist = bytearray()
        if reset_atomizer:
            self.atomizer.reset()

    def _output_state(
        self,
        alpha: torch.Tensor,
        persistent_state: torch.Tensor | None,
        abstractions: list[dict],
        atom_r: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Field-dominated surface conditioning vector (also used by probes).

        Prefer the spectral living-field readout when available so probes and
        decode agree.  Abstractions remain a light residual on the legacy mean.
        """
        legacy = alpha.mean(dim=0) if alpha.ndim > 1 else alpha
        if persistent_state is not None:
            legacy = legacy + (
                persistent_state.mean(dim=0)
                if persistent_state.ndim > 1
                else persistent_state
            )
        if abstractions:
            patterns = [
                item["pattern"].to(device=legacy.device, dtype=legacy.dtype).reshape(-1)
                for item in abstractions
                if item["pattern"].numel() == self.core.encoder.d_model
            ]
            if patterns:
                legacy = legacy + 0.1 * torch.stack(patterns).mean(dim=0)
        # Always-on consolidation buffer when per-tick persist is absent.
        persist = persistent_state
        if persist is None:
            persist = self.core.consolidation.persistence
        return self.surface.field_conditioned_state(
            alpha,
            persistent_state=persist,
            atom_r=atom_r,
            legacy_state=legacy,
        )

    def set_merge_thresholds(
        self,
        coherence_threshold: float | None = None,
        energy_floor: float | None = None,
    ) -> None:
        """Apply MERGE gates onto AggregationEngine (CLI / resume overrides)."""
        if coherence_threshold is not None:
            self.merge_coherence_threshold = float(coherence_threshold)
            self.core.aggregation.phase_coherence_threshold = self.merge_coherence_threshold
        if energy_floor is not None:
            self.merge_energy_floor = float(energy_floor)
            self.core.aggregation.merge_energy_floor = self.merge_energy_floor

    @staticmethod
    def _atom_prop_batch1(t: torch.Tensor, *, is_rho: bool = False) -> torch.Tensor:
        """Normalize MERGE outputs to AtomCompiler shapes: (1, d) / rho (1, 1)."""
        x = t.detach().clone().reshape(-1)
        if is_rho:
            return x[:1].reshape(1, 1)
        return x.reshape(1, -1)

    def _maybe_merge_atoms(self, *, enabled: bool | None = None) -> tuple[int, dict]:
        """MERGE coherent structural atoms into heavier ones (detached memory).

        ``enabled`` overrides ``self.enable_merge`` for one call. Chat/probe ingest
        passes ``enabled=False`` so prompt packets stay distinct living atoms;
        training leaves it ``None`` and keeps MERGE thr=0.45.
        """
        empty_diag = {"max_phase_coherence": 0.0, "n_pairs_above_energy_floor": 0}
        use_merge = self.enable_merge if enabled is None else bool(enabled)
        if not use_merge or len(self.core.atoms) < 2:
            return 0, empty_diag
        atoms = self.core.atoms
        merged, remove_idx, merge_count, diag = self.core.aggregation.merge_coherent(
            atoms.r,
            atoms.phi,
            atoms.omega,
            atoms.E,
            atoms.kappa,
            atoms.M,
            atoms.tau,
            atoms.rho,
        )
        if merge_count <= 0 or not remove_idx:
            return 0, diag
        # Capture payloads before remove (MERGE concatenates constituent bytes).
        pre_payloads = [bytes(getattr(a, "payload", b"") or b"") for a in atoms.atoms]
        atoms.remove(remove_idx)
        # Aggregation squeezes the compiler's leading batch-1 dim for math; restore
        # (1, d) / (1, 1) so ToroidalAtomCollection can stack with live atoms.
        max_keep = max(64, int(self.max_payload_bytes) * 4)
        new_atoms = []
        for item in merged:
            constituents = item.get("constituents") or []
            merged_payload = b"".join(
                pre_payloads[i] for i in constituents if 0 <= i < len(pre_payloads)
            )
            if len(merged_payload) > max_keep:
                merged_payload = merged_payload[-max_keep:]
            new_atoms.append(
                ToroidalAtom(
                    r=self._atom_prop_batch1(item["r"]),
                    phi=self._atom_prop_batch1(item["phi"]),
                    omega=self._atom_prop_batch1(item["omega"]),
                    E=self._atom_prop_batch1(item["E"]),
                    kappa=self._atom_prop_batch1(item["kappa"]),
                    M=self._atom_prop_batch1(item["M"]),
                    tau=self._atom_prop_batch1(item["tau"]),
                    rho=self._atom_prop_batch1(item["rho"], is_rho=True),
                    payload=merged_payload,
                )
            )
        atoms.extend(new_atoms)
        self.merge_count_total += merge_count
        return merge_count, diag

    FIELD_LEAK_SPECTRUM_GAINS = ("none", "ema", "sqrt")

    def set_field_leak_spectrum(
        self, spectrum, gain: str | None = None, slow_frac: float | None = None
    ) -> None:
        """Set/clear the per-mode leak spectrum (None/empty = off).

        ``gain`` (only used when the spectrum is on) scales the per-tick *input*
        of mode m (this tick's evolve of the atom projection + token drive from a
        zero field) by g_m, so slow modes do not integrate the input to a much
        larger amplitude than fast ones and swamp the normalized a_hat readout:
          none : g_m = 1 (pure per-mode leak)
          ema  : g_m = λ_m / LMAX  (every mode ≈ an EMA of the input; the
                 LMAX modes are exactly the scalar-λ=LMAX field)
          sqrt : g_m = sqrt(λ_m / LMAX)
        ``slow_frac`` (see ``field_leak_spectrum_rates``): fraction of modes
        carrying the LMIN→LMAX ladder, the rest at LMAX.  ``None`` keeps the
        current setting for ``gain`` / ``slow_frac``.
        """
        if slow_frac is not None:
            slow_frac = float(slow_frac)
            if not 0.0 < slow_frac <= 1.0:
                raise ValueError("field_leak_spectrum_slow_frac must be in (0, 1]")
            self.field_leak_spectrum_slow_frac = slow_frac
        if gain is not None:
            gain = str(gain or "none")
            if gain not in self.FIELD_LEAK_SPECTRUM_GAINS:
                raise ValueError(f"field_leak_spectrum_gain must be one of {self.FIELD_LEAK_SPECTRUM_GAINS}")
            self.field_leak_spectrum_gain = gain
        if spectrum is None or len(spectrum) == 0:
            self.field_leak_spectrum = None
            self._field_leak_vec = None
            self._field_gain_vec = None
            self._refresh_slow_idx()
            return
        if len(spectrum) != 2:
            raise ValueError("field_leak_spectrum must be (LMIN, LMAX)")
        lmin, lmax = float(spectrum[0]), float(spectrum[1])
        if not (0.0 < lmin <= lmax < 1.0):
            raise ValueError("field_leak_spectrum needs 0 < LMIN <= LMAX < 1")
        self.field_leak_spectrum = (lmin, lmax)
        n_modes = int(self.core.state.n_modes)
        lam = field_leak_spectrum_rates(
            n_modes, lmin, lmax, getattr(self, "field_leak_spectrum_slow_frac", 1.0)
        )
        # [n_modes, 1] retention factor, broadcast over d_model.
        self._field_leak_vec = (1.0 - lam).view(n_modes, 1)
        g = getattr(self, "field_leak_spectrum_gain", "none")
        if g == "ema":
            self._field_gain_vec = (lam / lmax).view(n_modes, 1)
        elif g == "sqrt":
            self._field_gain_vec = (lam / lmax).sqrt().view(n_modes, 1)
        else:
            self._field_gain_vec = None
        self._refresh_slow_idx()

    SLOW_READOUT_AHAT = ("exclude", "mix")
    SLOW_MEMORY_TARGETS = ("bag", "byte")

    def _refresh_slow_idx(self) -> None:
        """Recompute slow-mode indices / fast mask when the readout exists.
        No spectrum (or no slow mode) -> readout inactive (idx None)."""
        if getattr(self, "slow_readout", None) is None:
            return
        idx = slow_mode_indices(self.field_leak_rates())
        if not idx or len(idx) != self.slow_readout.n_slow:
            self._slow_idx = None
            self._slow_fast_mask = None
            return
        n_modes = int(self.core.state.n_modes)
        self._slow_idx = torch.tensor(idx, dtype=torch.long)
        mask = torch.ones(n_modes, 1)
        mask[self._slow_idx] = 0.0
        self._slow_fast_mask = mask

    def slow_readout_active(self) -> bool:
        return getattr(self, "slow_readout", None) is not None and getattr(self, "_slow_idx", None) is not None

    def set_slow_mode_readout(
        self,
        enabled: bool,
        *,
        rank: int | None = None,
        ahat: str | None = None,
        memory_weight: float | None = None,
        memory_lags: Iterable[int] | None = None,
        memory_target: str | None = None,
    ) -> None:
        """Enable/disable the slow-mode readout (``--slow-mode-readout``).

        Needs ``--field-leak-spectrum`` (the slow modes are its λ < LMAX modes).
        An existing compatible head (same n_slow / rank) is kept, so a resumed
        checkpoint keeps its trained weights; memory heads are added/removed to
        match ``memory_lags``.  ``enabled=False`` drops the module entirely.
        """
        if not enabled:
            self.slow_mode_readout = False
            self.slow_readout = None
            self._slow_idx = None
            self._slow_fast_mask = None
            self.slow_memory_weight = 0.0
            self._slow_hist = bytearray()
            return
        if rank is not None:
            self.slow_readout_rank = int(rank)
        if ahat is not None:
            if ahat not in self.SLOW_READOUT_AHAT:
                raise ValueError(f"slow_readout_ahat must be one of {self.SLOW_READOUT_AHAT}")
            self.slow_readout_ahat = str(ahat)
        if memory_weight is not None:
            if float(memory_weight) < 0:
                raise ValueError("slow_memory_weight must be >= 0")
            self.slow_memory_weight = float(memory_weight)
        if memory_lags is not None:
            lags = tuple(sorted({int(k) for k in memory_lags}))
            if any(k < 1 for k in lags):
                raise ValueError("slow_memory_lags must be >= 1")
            self.slow_memory_lags = lags
        if memory_target is not None:
            if memory_target not in self.SLOW_MEMORY_TARGETS:
                raise ValueError(f"slow_memory_target must be one of {self.SLOW_MEMORY_TARGETS}")
            self.slow_memory_target = str(memory_target)
        if getattr(self, "field_leak_spectrum", None) is None:
            raise ValueError("slow_mode_readout requires field_leak_spectrum (the slow modes)")
        idx = slow_mode_indices(self.field_leak_rates())
        if not idx:
            raise ValueError("slow_mode_readout: the spectrum has no mode slower than LMAX")
        d_model = int(self.core.encoder.d_model)
        sr = getattr(self, "slow_readout", None)
        if sr is None or sr.n_slow != len(idx) or sr.d_model != d_model or sr.rank != self.slow_readout_rank:
            sr = SlowModeReadout(
                len(idx), d_model, self.slow_readout_rank,
                self.max_payload_bytes * 256, self.slow_memory_lags,
            ).to(device=self.core.state.alpha.device)
            self.slow_readout = sr
        else:
            sr.set_lags(self.slow_memory_lags)
        self.slow_mode_readout = True
        self._slow_hist = bytearray()
        self._refresh_slow_idx()

    def readout_alpha(self, alpha: torch.Tensor) -> torch.Tensor:
        """α as the existing (a_hat) readout sees it: slow modes zeroed when the
        slow readout is active with ``slow_readout_ahat='exclude'``; else α."""
        if self.slow_readout_active() and self.slow_readout_ahat == "exclude":
            return alpha * self._slow_fast_mask.to(device=alpha.device, dtype=alpha.dtype)
        return alpha

    def _slow_memory_loss(self, output: dict) -> tuple[torch.Tensor, dict[str, float]]:
        """Aux: predict the recent past from the slow modes through their own head.

        target 'bag'  : for each lag K, the byte histogram of the bytes at lags
                        [K//2, K) (lag 0 = the current byte) -> soft-target CE.
        target 'byte' : the single byte at lag K -> CE.
        The slow state is detached here: with 1-tick truncated gradients the aux
        can only make the modes *readable* (train the head), it cannot change
        what earlier ticks wrote, and it must not steer the shared dynamics.
        """
        sr = self.slow_readout
        slow = output["slow_state"].detach()
        z = sr.features(slow)
        hist = self._slow_hist
        n = len(hist)
        losses = []
        info: dict[str, float] = {}
        for k in self.slow_memory_lags:
            head = sr.mem_heads[str(int(k))]
            if self.slow_memory_target == "byte":
                if n < k + 1:
                    continue
                tgt = torch.tensor([hist[n - 1 - k]], dtype=torch.long, device=z.device)
                lk = F.cross_entropy(head(z).unsqueeze(0), tgt)
            else:
                lo, hi = max(0, n - k), n - k // 2
                if hi <= lo:
                    continue
                ids = torch.tensor(list(hist[lo:hi]), dtype=torch.long, device=z.device)
                p = torch.zeros(256, device=z.device, dtype=z.dtype)
                p.scatter_add_(0, ids, torch.ones(ids.numel(), device=z.device, dtype=z.dtype))
                p = p / p.sum()
                lk = -(p * F.log_softmax(head(z), dim=-1)).sum()
            losses.append(lk)
            info[f"slow_memory_loss_{int(k)}"] = float(lk.detach().item())
        if not losses:
            return torch.zeros((), device=z.device), info
        loss = torch.stack(losses).mean()
        info["slow_memory_loss"] = float(loss.detach().item())
        return loss, info

    def field_leak_rates(self) -> torch.Tensor | None:
        """Per-mode λ actually applied each tick (None = no per-mode leak)."""
        if getattr(self, "_field_leak_vec", None) is None:
            return None
        return 1.0 - self._field_leak_vec.view(-1)

    def _field_rms_stable(self, rms: float) -> bool:
        prev = self._prev_field_rms
        self._prev_field_rms = rms
        if prev is None:
            return False
        return abs(rms - prev) / (abs(prev) + 1e-6) <= self.slow_rms_rel_tol

    def _advance(
        self,
        atom: ToroidalAtom,
        operation: torch.Tensor,
        confidence: torch.Tensor,
        *,
        merge_enabled: bool | None = None,
    ) -> dict:
        """Advance the unchanged toroidal core by one compiled atom.

        Fast path (every tick): inject + evolve + surface.
        Slow path (``slow_every`` + RMS-stable): abstraction + consolidation.
        Optional MERGE densifies coherent structural memory without GPU farms.
        ``merge_enabled`` overrides training MERGE for chat/probe ingest only.
        """
        self._tick += 1
        keep = getattr(self, "_field_leak_vec", None)
        gain_vec = getattr(self, "_field_gain_vec", None) if keep is not None else None
        injected = None
        bptt = self._bptt_active() if hasattr(self, "field_bptt") else False
        learn_lam = self.learnable_leak_rates() if getattr(self, "field_leak_logit", None) is not None else None
        alpha_graph = None
        if bptt:
            # --field-bptt: same ops/order as the legacy block below, but out of
            # place and with autograd (leak, then differentiable injection).
            alpha_graph = self._bptt_prev_alpha()
            if keep is not None:
                alpha_graph = alpha_graph * keep.to(device=alpha_graph.device, dtype=alpha_graph.dtype)
            elif learn_lam is not None:
                alpha_graph = alpha_graph * (1.0 - learn_lam)
            else:
                leak = float(getattr(self, "field_leak", 0.0) or 0.0)
                if leak > 0.0:
                    alpha_graph = alpha_graph * (1.0 - leak)
            injected = self.core.state._projection(
                atom.r, atom.phi, atom.omega, atom.E, atom.kappa
            )
            slow_gain = self.slow_block_input_gain(learn_lam) if learn_lam is not None else None
            if slow_gain is not None:
                alpha_graph = alpha_graph + injected * slow_gain
            else:
                alpha_graph = alpha_graph + injected
            with torch.no_grad():
                self.core.state.alpha.copy_(alpha_graph.detach())
        elif learn_lam is not None:
            with torch.no_grad():
                self.core.state.alpha.mul_(1.0 - learn_lam)
                slow_gain = self.slow_block_input_gain(learn_lam)
                if slow_gain is not None:
                    self.core.state.alpha.add_(
                        self.core.state._projection(atom.r, atom.phi, atom.omega, atom.E, atom.kappa) * slow_gain
                    )
                else:
                    self.core.state.add_atom_contribution(atom.r, atom.phi, atom.omega, atom.E, atom.kappa)
        elif keep is not None:
            # Multi-timescale leak: per-mode α_m ← (1-λ_m)·α_m (replaces scalar λ).
            with torch.no_grad():
                alpha_state = self.core.state.alpha
                alpha_state.mul_(keep.to(device=alpha_state.device, dtype=alpha_state.dtype))
        else:
            leak = float(getattr(self, "field_leak", 0.0) or 0.0)
            if leak > 0.0:
                with torch.no_grad():
                    self.core.state.alpha.mul_(1.0 - leak)
        if alpha_graph is not None or learn_lam is not None:
            pass  # --field-bptt / --field-leak-learnable: already injected above
        elif gain_vec is not None:
            # Same op as add_atom_contribution, but keep the projection for the
            # per-mode input gain below.
            with torch.no_grad():
                injected = self.core.state._projection(
                    atom.r, atom.phi, atom.omega, atom.E, atom.kappa
                )
                self.core.state.alpha.add_(injected)
        else:
            self.core.state.add_atom_contribution(atom.r, atom.phi, atom.omega, atom.E, atom.kappa)

        # Structural memory is intentionally detached from the current loss
        # graph.  The current atom still drives dynamics and receives gradients.
        memory_atom = ToroidalAtom(
            r=atom.r.detach().clone(),
            phi=atom.phi.detach().clone(),
            omega=atom.omega.detach().clone(),
            E=atom.E.detach().clone(),
            kappa=atom.kappa.detach().clone(),
            M=atom.M.detach().clone(),
            tau=atom.tau.detach().clone(),
            rho=atom.rho.detach().clone(),
            payload=bytes(getattr(atom, "payload", b"") or b""),
        )
        self.core.atoms.add(memory_atom)
        merge_count, merge_diag = self._maybe_merge_atoms(enabled=merge_enabled)

        if alpha_graph is not None:
            alpha = alpha_graph
        else:
            alpha = self.core.state.get_field().detach().clone()
        alpha_new = self.core.dynamics.evolve(alpha, input_token=atom.r)
        if gain_vec is not None:
            # The tick is affine in α: evolve(x + p) = Lin(x) + evolve(p) with
            # evolve(p) = this tick's input (projection + token drive).  Scale
            # only the input part per mode: + (g_m - 1)·evolve(p).
            tick_input = self.core.dynamics.evolve(injected, input_token=atom.r)
            g = gain_vec.to(device=alpha_new.device, dtype=alpha_new.dtype)
            alpha_new = alpha_new + (g - 1.0) * tick_input
        interaction_energy = self.core.interaction.field_interaction(alpha_new)
        left = torch.roll(alpha_new, 1, 0)
        right = torch.roll(alpha_new, -1, 0)
        alpha_new = alpha_new + 0.1 * self.core.dynamics.dynamics.phase_sync * (
            left + right - 2.0 * alpha_new
        )
        alpha_new, amplitude = self.field_controller(alpha_new)

        rms_after = float(amplitude["field_rms_after"])
        do_slow = self.slow_every <= 1 or (
            self._tick % self.slow_every == 0 and self._field_rms_stable(rms_after)
        )
        if self.slow_every > 1 and self._tick % self.slow_every != 0:
            self._prev_field_rms = rms_after

        aggregates: list[dict] = []
        abstractions: list[dict] = []
        persistent_state = None
        if do_slow:
            if len(self.core.atoms) > 10:
                aggregates, _ = self.core.aggregation.aggregate(
                    self.core.atoms.r,
                    self.core.atoms.phi,
                    self.core.atoms.omega,
                    self.core.atoms.E,
                    self.core.atoms.kappa,
                    self.core.atoms.M,
                    self.core.atoms.tau,
                    self.core.atoms.rho,
                )
            for aggregate in aggregates:
                index = self.core.abstraction.create_or_update_abstraction([aggregate])
                if index >= 0:
                    abstractions.append(
                        {
                            "id": index,
                            "pattern": self.core.abstraction.abstraction_memory[index].detach(),
                        }
                    )
                    self.core.abstraction_count += 1
            self.core.abstraction.step_age()

            if len(self.core.energy_history) > 5:
                energy = self.core.atoms.E.mean(dim=0)
                stability = self.core.consolidation.compute_stability(
                    energy, self.core.energy_history
                )
                persistent_state, mask = self.core.consolidation.consolidate(
                    alpha_new,
                    energy.unsqueeze(0),
                    stability.unsqueeze(0),
                )
                if mask.any():
                    self.core.consolidation_count += int(mask.sum().item())
        self.core.energy_history.append(self.core.atoms.E.mean(dim=0).detach())
        self.core.energy_history = self.core.energy_history[-20:]

        # Living consolidation buffer (EMA) — not only the per-tick consolidate().
        persist_for_readout = persistent_state
        if persist_for_readout is None:
            persist_for_readout = self.core.consolidation.persistence
        slow_on = self.slow_readout_active()
        # --slow-mode-readout with ahat='exclude': the existing readout sees α
        # with the slow modes zeroed; they are read only by their own head.
        alpha_read = self.readout_alpha(alpha_new) if slow_on else alpha_new
        output_state = self._output_state(
            alpha_read, persist_for_readout, abstractions, atom_r=atom.r
        )
        living_atoms = list(self.core.atoms.atoms)
        living_payloads = [
            bytes(getattr(a, "payload", b"") or b"") for a in living_atoms
        ]
        living_energies = [
            a.E.detach().float().mean() for a in living_atoms
        ]
        la_last_b, la_prev_b = _tail_bytes(living_payloads)
        surface = self.surface(
            output_state,
            alpha=alpha_read,
            persistent_state=persist_for_readout,
            atom_r=atom.r,
            atom_payloads=living_payloads,
            atom_energies=living_energies,
            last_phi=atom.phi.reshape(-1),
        )
        slow_extra: dict = {}
        if slow_on:
            # Own slow-mode readout: Δlogits = W_o·z, z = tanh(W_t·ŝ), W_o zero-init.
            # A field term, so it stays on under last_atom_force_off.  z is a
            # stop-gradient feature for next-byte CE: CE trains W_o only.  The
            # trunk W_t is shaped only by --slow-memory-loss (first sweep: when CE
            # also trained W_t, it drove the trunk bias into tanh saturation,
            # z became a constant and the head degenerated into an output bias).
            slow_state = alpha_new[self._slow_idx.to(alpha_new.device)]
            upd = bool(self.training and torch.is_grad_enabled())
            with torch.no_grad():
                z = self.slow_readout.features(slow_state, update_stats=upd)
            add = self.slow_readout.to_logits(z).view_as(surface["byte_logits"])
            surface = {**surface, "byte_logits": surface["byte_logits"] + add}
            self._slow_hist.extend(bytes(getattr(atom, "payload", b"") or b""))
            keep = max(self.slow_memory_lags or (1,)) + 1
            if len(self._slow_hist) > 4 * keep:
                del self._slow_hist[: len(self._slow_hist) - keep]
            slow_extra = {"slow_state": slow_state, "slow_z": z}
        with torch.no_grad():
            self.core.state.alpha.copy_(alpha_new.detach())
            self.core.state.t.add_(self.core.dynamics.dt)
        if hasattr(self, "field_bptt"):
            self._bptt_alpha = alpha_new if bptt else None
        return {
            **slow_extra,
            "surface": surface,
            "field": alpha_new,
            "confidence": confidence,
            "operation": operation,
            "interaction_energy": interaction_energy,
            "aggregates": aggregates,
            "abstractions": abstractions,
            "persistent_state": persistent_state,
            "n_atoms": len(self.core.atoms),
            "merge_count": merge_count,
            "merge_count_total": self.merge_count_total,
            "max_phase_coherence": float(merge_diag.get("max_phase_coherence", 0.0)),
            "n_pairs_above_energy_floor": int(merge_diag.get("n_pairs_above_energy_floor", 0)),
            "slow_tick": do_slow,
            "atom_r": atom.r,
            "la_bytes": (la_last_b, la_prev_b),
            "la_phi": atom.phi.reshape(-1).detach(),
            **amplitude,
        }

    def forward_packet(
        self, packet: AtomPacket, *, merge_enabled: bool | None = None
    ) -> dict:
        features = packet.features.to(self.core.state.alpha.device)
        atom, operation, confidence = self.compiler(features, atom_count=len(self.core.atoms))
        # Living atom carries the packet's surface bytes for payload production.
        atom.payload = bytes(packet.payload or b"")
        return self._advance(atom, operation, confidence, merge_enabled=merge_enabled)

    @torch.no_grad()
    def _refresh_prompt_alpha_bank(self) -> None:
        """Fill ``_prompt_alpha_bank`` from fixed distinct prompts without lasting
        mutation of the live train field / atoms / atomizer / train α-bank.

        Snapshots core + atomizer + tick bookkeeping, runs each prompt from a
        clean episode, stores detached α, then restores. Used only when
        ``field_ignorance_prompt_bank`` is on.
        """
        device = self.core.state.alpha.device
        snap = {
            "state": {k: v.detach().clone() if torch.is_tensor(v) else v
                      for k, v in self.core.state.state_dict().items()},
            "atoms": self.core.atoms.state_dict(),
            "consolidation": {
                k: v.detach().clone() if torch.is_tensor(v) else v
                for k, v in self.core.consolidation.state_dict().items()
            },
            "abstraction": {
                k: v.detach().clone() if torch.is_tensor(v) else v
                for k, v in self.core.abstraction.state_dict().items()
            },
            "atomizer": self.atomizer.state_dict(),
            "energy_history": list(self.core.energy_history),
            "consolidation_count": self.core.consolidation_count,
            "abstraction_count": self.core.abstraction_count,
            "tick": self._tick,
            "prev_field_rms": self._prev_field_rms,
            "alpha_bank": [a.detach().clone() for a in self._alpha_bank],
            "training": self.training,
        }
        was_training = self.training
        self.eval()
        bank: list[torch.Tensor] = []
        try:
            for text in self._ignorance_prompts:
                self.reset_state(reset_atomizer=True)
                # restore train α-bank after reset_state cleared it (temp)
                packets = self.atomizer.encode(text, reset=True)
                if not packets:
                    continue
                out = None
                for packet in packets:
                    out = self.forward_packet(packet)
                if out is not None:
                    bank.append(out["field"].detach().clone().to(device=device))
        finally:
            self.core.state.load_state_dict(snap["state"], strict=False)
            self.core.atoms.load_state_dict(snap["atoms"])
            self.core.consolidation.load_state_dict(snap["consolidation"])
            self.core.abstraction.load_state_dict(snap["abstraction"])
            self.atomizer.load_state_dict(snap["atomizer"])
            self.core.energy_history = snap["energy_history"]
            self.core.consolidation_count = snap["consolidation_count"]
            self.core.abstraction_count = snap["abstraction_count"]
            self._tick = snap["tick"]
            self._prev_field_rms = snap["prev_field_rms"]
            self._alpha_bank = snap["alpha_bank"]
            if was_training:
                self.train()
            else:
                self.eval()
        self._prompt_alpha_bank = bank

    def _field_auxiliary_losses(
        self,
        output: dict,
        current: AtomPacket,
        target: AtomPacket | None = None,
    ) -> tuple[torch.Tensor, dict[str, float]]:
        """Keep surface conditioned on α; probe persistence + next-packet from field.

        Contrastive hinge: logits from true α must differ from zeroed / shuffled α
        (cosine below margin). Field-ignorance hinge: logits from true α must differ
        from logits under stopgrad mean/bank of *other* prompts' α (not zero).
        Persistence probe: reconstruct *current* packet features from spectral field
        feats. Next-packet production: predict *target* packet features from the same
        α (+ atom_r) features — small weight under hard. Weighted small vs CE.
        """
        device = output["field"].device
        zero = torch.zeros((), device=device)
        info: dict[str, float] = {
            "field_contrast_loss": 0.0,
            "field_persist_loss": 0.0,
            "field_next_packet_loss": 0.0,
            "field_ignorance_loss": 0.0,
            "printable_aux_loss": 0.0,
            "printable_mass_mean": float("nan"),
            "field_logit_cos_zero": float("nan"),
            "field_logit_cos_shuf": float("nan"),
            "field_logit_cos_ignorance": float("nan"),
        }
        contrast = zero
        persist = zero
        next_pkt = zero
        ignorance = zero
        printable = zero
        alpha = output["field"]
        atom_r = output.get("atom_r")
        persist_state = output.get("persistent_state")
        if persist_state is None:
            persist_state = self.core.consolidation.persistence
        true_flat = output["surface"]["byte_logits"].reshape(-1)
        # Last-atom readout: aux hinges compare same-bytes/different-alpha so
        # they keep isolating alpha (bytes held constant, no shape change).
        la_extra: dict = {}
        if getattr(self.surface, "last_atom", None) is not None:
            la_pair = output.get("la_bytes") or (0, 0)
            la_extra = {
                "atom_payloads": [bytes([la_pair[1]]), bytes([la_pair[0]])],
                "last_phi": output.get("la_phi"),
            }

        if self.field_contrast_weight > 0:
            surf_zero = self.surface(
                None,
                alpha=torch.zeros_like(alpha),
                persistent_state=persist_state,
                atom_r=atom_r,
                **la_extra,
            )
            cos_z = F.cosine_similarity(
                true_flat.unsqueeze(0),
                surf_zero["byte_logits"].reshape(-1).unsqueeze(0),
            ).squeeze()
            if alpha.shape[0] > 1:
                shuf_alpha = alpha[torch.randperm(alpha.shape[0], device=alpha.device)]
            else:
                shuf_alpha = -alpha
            surf_shuf = self.surface(
                None, alpha=shuf_alpha, persistent_state=persist_state, atom_r=atom_r,
                **la_extra,
            )
            cos_s = F.cosine_similarity(
                true_flat.unsqueeze(0),
                surf_shuf["byte_logits"].reshape(-1).unsqueeze(0),
            ).squeeze()
            # Harder negative: a prior living field from the episode bank (if any).
            cos_b = true_flat.new_zeros(())
            if self._alpha_bank:
                bank_alpha = self._alpha_bank[int(torch.randint(len(self._alpha_bank), (1,)).item())]
                bank_alpha = bank_alpha.to(device=alpha.device, dtype=alpha.dtype)
                if bank_alpha.shape == alpha.shape and not torch.allclose(bank_alpha, alpha, atol=1e-5):
                    surf_bank = self.surface(
                        None, alpha=bank_alpha, persistent_state=persist_state, atom_r=atom_r,
                        **la_extra,
                    )
                    cos_b = F.cosine_similarity(
                        true_flat.unsqueeze(0),
                        surf_bank["byte_logits"].reshape(-1).unsqueeze(0),
                    ).squeeze()
            margin = self.field_contrast_margin
            # Tighter margin on bank negatives so inter-prompt logits must separate.
            contrast = (
                F.relu(cos_z - margin)
                + F.relu(cos_s - margin)
                + F.relu(cos_b - (margin - 0.15))
            )
            info["field_logit_cos_zero"] = float(cos_z.detach().item())
            info["field_logit_cos_shuf"] = float(cos_s.detach().item())
            info["field_logit_cos_bank"] = float(cos_b.detach().item()) if cos_b.ndim == 0 or cos_b.numel()==1 else float("nan")
            info["field_contrast_loss"] = float(contrast.detach().item())

        feats = None
        if self.field_loss_weight > 0 or self.field_next_packet_weight > 0:
            feats = self.surface.field_feat_norm(
                self.surface.field_features(alpha, persist_state, atom_r)
            )
        if self.field_loss_weight > 0:
            pred = self.field_probe(feats)
            cur_feat = current.features.to(device=pred.device, dtype=pred.dtype)
            cos = F.cosine_similarity(pred.unsqueeze(0), cur_feat.unsqueeze(0)).squeeze()
            persist = 1.0 - cos
            info["field_persist_loss"] = float(persist.detach().item())

        # Field → next-packet production: predict *upcoming* packet features from α+atoms.
        if self.field_next_packet_weight > 0 and target is not None:
            pred_next = self.field_next_probe(feats)
            next_feat = target.features.to(device=pred_next.device, dtype=pred_next.dtype)
            cos_n = F.cosine_similarity(pred_next.unsqueeze(0), next_feat.unsqueeze(0)).squeeze()
            next_pkt = 1.0 - cos_n
            info["field_next_packet_loss"] = float(next_pkt.detach().item())

        # Field-ignorance: surface must not map any non-zero α to one shared
        # logit template. Compare ℓ(α) vs ℓ(sg[α_bar]) where α_bar is the
        # sliding mean (or bank) of *other* prompts — not zero, not same-prompt.
        if self.field_ignorance_weight > 0:
            if self.field_ignorance_prompt_bank:
                if not self._prompt_alpha_bank:
                    self._refresh_prompt_alpha_bank()
                source = self._prompt_alpha_bank
            else:
                # Bank already includes current snap (appended in transition_loss);
                # use other entries only.
                source = self._alpha_bank[:-1]
            others = [
                a for a in source
                if a.shape == alpha.shape and not torch.allclose(a, alpha.detach(), atol=1e-5)
            ]
            if others:
                alpha_bar = torch.stack(
                    [a.to(device=alpha.device, dtype=alpha.dtype) for a in others]
                ).mean(dim=0)
                alpha_bar = alpha_bar.detach()  # sg[α_bar]
                # Ablate shared atom_r / persistence on the α_bar call only so
                # the hinge isolates α-driven readout. ℓ(α) keeps live state.
                if self.field_ignorance_ablate_shared:
                    bar_persist = None
                    bar_atom_r = None
                else:
                    bar_persist = persist_state
                    bar_atom_r = atom_r
                surf_bar = self.surface(
                    None, alpha=alpha_bar, persistent_state=bar_persist, atom_r=bar_atom_r,
                    **la_extra,
                )
                cos_ign = F.cosine_similarity(
                    true_flat.unsqueeze(0),
                    surf_bar["byte_logits"].reshape(-1).unsqueeze(0),
                ).squeeze()
                ignorance = F.relu(cos_ign - self.field_ignorance_margin)
                info["field_logit_cos_ignorance"] = float(cos_ign.detach().item())
                info["field_ignorance_loss"] = float(ignorance.detach().item())

        # Printable UTF-8 aux on hard (or any) surface logits: reward mass on
        # printable bytes so CE under frozen+MLP stops favoring binary garbage.
        # Gradients flow into α-MLP only under hard (decoder frozen).
        if self.printable_aux_weight > 0:
            byte_logits = output["surface"]["byte_logits"]
            mass = self.surface.printable_mass(byte_logits)
            # 1 - mass ∈ [0,1]; also lightly penalize via -log(mass) for sharp push.
            printable = (1.0 - mass).mean() + 0.25 * (-mass.clamp_min(1e-6).log()).mean()
            info["printable_aux_loss"] = float(printable.detach().item())
            info["printable_mass_mean"] = float(mass.detach().mean().item())

        total_aux = (
            self.field_contrast_weight * contrast
            + self.field_loss_weight * persist
            + self.field_next_packet_weight * next_pkt
            + self.field_ignorance_weight * ignorance
            + self.printable_aux_weight * printable
        )
        return total_aux, info

    def transition_loss(self, current: AtomPacket, target: AtomPacket) -> tuple[torch.Tensor, dict]:
        """Predict the complete next atom surface packet (+ optional field aux)."""
        output = self.forward_packet(current)
        # Update α bank (detached) for cross-field contrast on later steps.
        with torch.no_grad():
            snap = output["field"].detach().clone()
            self._alpha_bank.append(snap)
            if len(self._alpha_bank) > self._alpha_bank_max:
                self._alpha_bank.pop(0)
            # Refresh distinct-prompt bank for L_ign (does not alter train bank).
            if (
                self.field_ignorance_prompt_bank
                and self.field_ignorance_weight > 0
                and (
                    not self._prompt_alpha_bank
                    or (self._tick > 0 and self._tick % self._prompt_bank_refresh_every == 0)
                )
            ):
                self._refresh_prompt_alpha_bank()
        surface = output["surface"]
        payload = target.payload[: self.max_payload_bytes]
        target_length = torch.tensor(
            [len(payload) - 1], dtype=torch.long, device=surface["length_logits"].device
        )
        length_loss = F.cross_entropy(surface["length_logits"].unsqueeze(0), target_length)
        target_bytes = torch.tensor(list(payload), dtype=torch.long, device=surface["byte_logits"].device)
        byte_logits = surface["byte_logits"][: len(payload)]
        byte_loss = F.cross_entropy(byte_logits, target_bytes)
        ce_loss = length_loss + byte_loss
        aux_loss, aux_info = self._field_auxiliary_losses(output, current, target)
        loss = ce_loss + aux_loss
        if "slow_state" in output and float(getattr(self, "slow_memory_weight", 0.0)) > 0.0:
            mem_loss, mem_info = self._slow_memory_loss(output)
            loss = loss + self.slow_memory_weight * mem_loss
            aux_info = {**aux_info, **mem_info}
        return loss, {
            "loss": float(loss.detach().item()),
            "ce_loss": float(ce_loss.detach().item()),
            "length_loss": float(length_loss.detach().item()),
            "byte_loss": float(byte_loss.detach().item()),
            "n_atoms": len(self.core.atoms),
            "surface": surface,
            "field_norm": float(output["field"].detach().norm().item()),
            "field_rms_before": output["field_rms_before"],
            "field_rms_after": output["field_rms_after"],
            "field_scale": output["field_scale"],
            "merge_count": int(output.get("merge_count", 0)),
            "merge_count_total": int(output.get("merge_count_total", self.merge_count_total)),
            "max_phase_coherence": float(output.get("max_phase_coherence", 0.0)),
            "n_pairs_above_energy_floor": int(output.get("n_pairs_above_energy_floor", 0)),
            "slow_tick": bool(output.get("slow_tick", True)),
            **aux_info,
        }

    def efference_loss(self, current: AtomPacket, target: AtomPacket) -> tuple[torch.Tensor, dict]:
        """One efference tick (bridge #9): predict from gold state, commit the
        model's OWN byte (generate-like), then CE on the gold target from that
        state. Trains recovery from free-run drift (train~=generate).

        Falls back to a normal transition when the atomizer holds a pending
        span (only possible off byte-tick; byte-tick buffer is always empty).
        """
        if len(self.atomizer._buffer) == 0:
            with torch.no_grad():
                pred_bytes = self.predict_packet(current, deterministic=True, merge_enabled=False)
            pred_packet = self.atomizer.packet_from_payload(bytes(pred_bytes))
            return self.transition_loss(pred_packet, target)
        return self.transition_loss(current, target)

    def free_run_aux_loss(
        self,
        start: AtomPacket,
        golds: list[AtomPacket],
        *,
        weight: float = 1.0,
    ) -> tuple[torch.Tensor, dict]:
        """Multi-tick free-run CE using the SAME commit path as generate_packets.

        Math (L_roll): for horizon H = len(golds),
          L_roll = (1/H) Σ_h CE(logits(s_h), gold_h)
        where s_0 = start, and s_{h+1} = packet_from_payload(argmax(decode(s_h)))
        with the discrete commit stop-grad (scheduled sampling). Gradients flow
        through each step's surface logits (α-MLP + last_atom + …) under the
        free-run field trajectory — closing teacher-forced ↔ chat mismatch.

        Mutates living field like generate (no restore): continuum after this
        step is free-run-drifted. Falls back to transition_loss when golds is
        empty or atomizer has a pending span.
        """
        if not golds:
            raise ValueError("free_run_aux_loss requires at least one gold packet")
        if len(self.atomizer._buffer) != 0:
            return self.transition_loss(start, golds[0])

        device = self.core.state.alpha.device
        total_ce = torch.zeros((), device=device)
        byte_losses: list[float] = []
        length_losses: list[float] = []
        current = start
        n_steps = 0
        last_output: dict | None = None
        for gold in golds:
            output = self.forward_packet(current, merge_enabled=False)
            last_output = output
            surface = output["surface"]
            payload = gold.payload[: self.max_payload_bytes]
            if not payload:
                # Skip empty gold; still need a commit to advance free-run.
                with torch.no_grad():
                    pred_bytes = self.surface.decode(
                        surface, deterministic=True, prefer_printable=True
                    )
                current = self.atomizer.packet_from_payload(bytes(pred_bytes))
                continue
            target_length = torch.tensor(
                [len(payload) - 1], dtype=torch.long, device=surface["length_logits"].device
            )
            length_loss = F.cross_entropy(
                surface["length_logits"].unsqueeze(0), target_length
            )
            target_bytes = torch.tensor(
                list(payload), dtype=torch.long, device=surface["byte_logits"].device
            )
            byte_logits = surface["byte_logits"][: len(payload)]
            byte_loss = F.cross_entropy(byte_logits, target_bytes)
            step_ce = length_loss + byte_loss
            total_ce = total_ce + step_ce
            byte_losses.append(float(byte_loss.detach().item()))
            length_losses.append(float(length_loss.detach().item()))
            n_steps += 1
            # Discrete free-run commit (stop-grad) — same decode as chat.
            with torch.no_grad():
                pred_bytes = self.surface.decode(
                    surface, deterministic=True, prefer_printable=True
                )
            current = self.atomizer.packet_from_payload(bytes(pred_bytes))

        if n_steps == 0:
            return self.transition_loss(start, golds[0])

        mean_ce = total_ce / float(n_steps)
        w = float(weight)
        loss = w * mean_ce
        mean_byte = sum(byte_losses) / len(byte_losses)
        mean_len = sum(length_losses) / len(length_losses)
        out = last_output or {}
        return loss, {
            "loss": float(loss.detach().item()),
            "ce_loss": float(mean_ce.detach().item()),
            "byte_loss": mean_byte,
            "length_loss": mean_len,
            "n_atoms": len(self.core.atoms),
            "field_norm": float(out.get("field", self.core.state.alpha).detach().norm().item())
                if "field" in out else float(self.core.state.alpha.detach().norm().item()),
            "field_rms_before": float(out.get("field_rms_before", 0.0) or 0.0),
            "field_rms_after": float(out.get("field_rms_after", 0.0) or 0.0),
            "field_scale": float(out.get("field_scale", 1.0) or 1.0),
            "merge_count": int(out.get("merge_count", 0) or 0),
            "merge_count_total": int(out.get("merge_count_total", self.merge_count_total) or self.merge_count_total),
            "max_phase_coherence": float(out.get("max_phase_coherence", 0.0) or 0.0),
            "n_pairs_above_energy_floor": int(out.get("n_pairs_above_energy_floor", 0) or 0),
            "slow_tick": bool(out.get("slow_tick", True)),
            "free_run_aux": True,
            "free_run_horizon": int(n_steps),
            "free_run_aux_weight": w,
            "field_contrast_loss": 0.0,
            "field_persist_loss": 0.0,
            "field_ignorance_loss": 0.0,
            "printable_aux_loss": 0.0,
            "printable_mass_mean": float("nan"),
            "field_next_packet_loss": 0.0,
        }

    @torch.no_grad()
    def predict_packet(
        self,
        packet: AtomPacket,
        temperature: float = 1.0,
        top_k: int | None = None,
        deterministic: bool = True,
        prefer_printable: bool = True,
        *,
        merge_enabled: bool | None = None,
    ) -> bytes:
        output = self.forward_packet(packet, merge_enabled=merge_enabled)
        return self.surface.decode(
            output["surface"],
            temperature=temperature,
            top_k=top_k,
            deterministic=deterministic,
            prefer_printable=prefer_printable,
        )

    def generate_packets(
        self,
        prompt: str,
        max_packets: int = 8,
        temperature: float = 1.0,
        top_k: int | None = None,
        deterministic: bool = True,
        max_length: int | None = None,
        prefer_printable: bool = True,
        reset: bool = True,
        merge_enabled: bool = False,
        payload_copy: bool = False,
        dialogue_wrap: bool = True,
        speech_gate: bool = True,
        anti_repeat: bool = True,
    ) -> bytes:
        """Prime on a prompt and generate bounded atom surface packets."""
        self.eval()
        if self.field_obligatory_hard:
            self.surface.set_obligatory_hard(True)
            prefer_printable = True
        prev_payload_flag = bool(getattr(self.surface, "payload_enabled", True))
        self.surface.payload_enabled = bool(payload_copy)
        try:
            if reset:
                self.reset_state(reset_atomizer=True)
            if dialogue_wrap:
                prompt = format_dialogue_prompt(prompt)
            prompt_packets = self.atomizer.encode(prompt, reset=reset)
            if not prompt_packets:
                raise ValueError("prompt must produce at least one atom packet")
            for packet in prompt_packets[:-1]:
                self.forward_packet(packet, merge_enabled=merge_enabled)
            current = prompt_packets[-1]
            generated = bytearray()
            for _ in range(max_packets):
                payload = self.predict_packet(
                    current,
                    temperature=temperature,
                    top_k=top_k,
                    deterministic=deterministic,
                    prefer_printable=prefer_printable,
                    merge_enabled=merge_enabled,
                )
                if anti_repeat and payload and _byte_cycle(generated + payload):
                    # N-gram anti-repeat (ARCHITECTURE.md): the new bytes close
                    # a 2-6 byte cycle -> draw once more. RNG state is
                    # saved/restored and the redraw is seeded by position so
                    # deterministic probes stay reproducible.
                    rng_state = torch.get_rng_state()
                    try:
                        torch.manual_seed(0xA700 + len(generated))
                        payload = self.predict_packet(
                            current,
                            temperature=0.7,
                            top_k=8,
                            deterministic=False,
                            prefer_printable=prefer_printable,
                            merge_enabled=merge_enabled,
                        )
                    finally:
                        torch.set_rng_state(rng_state)
                if not payload:
                    break
                # Byte-tick contract (ARCHITECTURE.md): the speech gate judges
                # the accumulated phrase, never a single (1-byte) packet.
                if (
                    speech_gate
                    and len(generated) >= 8
                    and not speech_ok(bytes(generated))
                ):
                    break
                generated.extend(payload)
                current = self.atomizer.packet_from_payload(payload)
                if max_length is not None and len(generated) >= max_length:
                    break
                if generated.endswith(b"\n\n") and len(generated) > 2:
                    break
            if max_length is not None:
                return bytes(generated[:max_length])
            return bytes(generated)
        finally:
            self.surface.payload_enabled = prev_payload_flag


    def _checkpoint(self) -> dict:
        core = self.core
        return {
            "compiler": self.compiler.state_dict(),
            "surface": self.surface.state_dict(),
            "core": {
                "encoder": core.encoder.state_dict(),
                "state": core.state.get_state_dict(),
                "dynamics": core.dynamics.state_dict(),
                "interaction": core.interaction.state_dict(),
                "aggregation": core.aggregation.state_dict(),
                "abstraction": core.abstraction.state_dict(),
                "consolidation": core.consolidation.state_dict(),
                "production": core.production.state_dict(),
                "atoms": core.atoms.state_dict(),
                "energy_history": core.energy_history,
                "training_loss": core.training_loss,
                "consolidation_count": core.consolidation_count,
                "abstraction_count": core.abstraction_count,
            },
            "atomizer": self.atomizer.state_dict(),
            "config": {
                "d_model": core.encoder.d_model,
                "n_modes": core.state.n_modes,
                "n_atoms_max": core.encoder.n_atoms_max,
                "max_payload_bytes": self.max_payload_bytes,
                "atomizer_version": self.atomizer.VERSION,
                "field_max_rms": self.field_max_rms,
                "field_leak": float(getattr(self, "field_leak", 0.0) or 0.0),
                "field_leak_spectrum": (
                    list(self.field_leak_spectrum)
                    if getattr(self, "field_leak_spectrum", None) is not None
                    else None
                ),
                "field_leak_spectrum_gain": str(getattr(self, "field_leak_spectrum_gain", "none")),
                "field_leak_spectrum_slow_frac": float(
                    getattr(self, "field_leak_spectrum_slow_frac", 1.0)
                ),
                "energy_decay_bounds": self.energy_decay_bounds,
                "enable_merge": self.enable_merge,
                "merge_coherence_threshold": float(self.merge_coherence_threshold),
                "merge_energy_floor": float(self.merge_energy_floor),
                "slow_every": self.slow_every,
                "field_loss_weight": self.field_loss_weight,
                "field_contrast_weight": self.field_contrast_weight,
                "field_contrast_margin": self.field_contrast_margin,
                "field_ignorance_weight": self.field_ignorance_weight,
                "field_ignorance_margin": self.field_ignorance_margin,
                "field_ignorance_ablate_shared": self.field_ignorance_ablate_shared,
                "field_ignorance_prompt_bank": self.field_ignorance_prompt_bank,
                "field_obligatory_readout": bool(self.field_obligatory_readout),
                "field_obligatory_hard": bool(self.field_obligatory_hard),
                "payload_copy": bool(getattr(self.surface, "payload_enabled", True)),
                "last_atom_readout": bool(self.last_atom_readout),
                "last_atom_scale": float(getattr(self.surface, "last_atom_scale", 1.0)),
                "last_atom_scale_adaptive": bool(
                    getattr(self.surface, "last_atom_scale_adaptive", False)
                ),
                "last_atom_scale_min": float(
                    getattr(self.surface, "last_atom_scale_min", 0.0)
                ),
                "last_atom_scale_max": float(
                    getattr(self.surface, "last_atom_scale_max", 1.0)
                ),
                "field_hard_couple_scale": float(
                    getattr(self.surface, "field_hard_couple_scale", 0.0)
                ),
                "last_atom_dropout": float(
                    getattr(self.surface, "last_atom_dropout", 0.0)
                ),
                "printable_aux_weight": float(self.printable_aux_weight),
                "field_next_packet_weight": float(self.field_next_packet_weight),
                "merge_count_total": self.merge_count_total,
            },
            "field_probe": self.field_probe.state_dict(),
            "field_next_probe": self.field_next_probe.state_dict(),
            **self._slow_readout_checkpoint(),
            **self._bptt_checkpoint(),
        }

    def _bptt_checkpoint(self) -> dict:
        """Extra keys only when --field-bptt / --field-leak-learnable are on
        (off = checkpoint layout exactly as before)."""
        out: dict = {}
        if int(getattr(self, "field_bptt", 0) or 0) > 0:
            out["field_bptt"] = int(self.field_bptt)
        if getattr(self, "field_leak_logit", None) is not None:
            out["field_leak_logit"] = self.field_leak_logit.detach().cpu().clone()
        if getattr(self, "field_leak_slow_block", None) is not None:
            out["field_leak_slow_block"] = {k: v for k, v in self.field_leak_slow_block.items() if k != "modes"}
        return out

    def _slow_readout_checkpoint(self) -> dict:
        """Extra checkpoint keys, present only when the slow readout exists
        (off = checkpoint layout exactly as before)."""
        sr = getattr(self, "slow_readout", None)
        if sr is None:
            return {}
        return {
            "slow_readout": sr.state_dict(),
            "slow_readout_config": {
                "slow_mode_readout": True,
                "n_slow": int(sr.n_slow),
                "slow_idx": [int(i) for i in (self._slow_idx.tolist() if self._slow_idx is not None else [])],
                "rank": int(sr.rank),
                "ahat": str(self.slow_readout_ahat),
                "memory_weight": float(self.slow_memory_weight),
                "memory_lags": [int(k) for k in self.slow_memory_lags],
                "memory_target": str(self.slow_memory_target),
            },
        }

    def save(self, path: str | Path, extra_state: dict | None = None) -> str:
        checkpoint = self._checkpoint()
        if extra_state is not None:
            checkpoint["training"] = extra_state
        path = str(path)
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        torch.save(checkpoint, path)
        return path


    @staticmethod
    def _pad_copy_leading(dst: torch.Tensor, src: torch.Tensor) -> int:
        """Copy leading rows/elems from src into dst; return rows copied on dim0."""
        if dst.ndim == 0 or src.ndim == 0:
            return 0
        n = min(int(dst.shape[0]), int(src.shape[0]))
        if n <= 0:
            return 0
        if dst.ndim == 1:
            dst[:n].copy_(src[:n].to(device=dst.device, dtype=dst.dtype))
        else:
            # Match trailing dims; truncate if src wider (should not happen on grow).
            if dst.shape[1:] != src.shape[1:]:
                # Best-effort: copy overlapping trailing slice.
                slices = [slice(0, n)]
                for d_dim, s_dim in zip(dst.shape[1:], src.shape[1:]):
                    slices.append(slice(0, min(d_dim, s_dim)))
                dst[tuple(slices)].copy_(src[tuple(slices)].to(device=dst.device, dtype=dst.dtype))
            else:
                dst[:n].copy_(src[:n].to(device=dst.device, dtype=dst.dtype))
        return n

    def _migrate_span_head_weights(self, surface_state: dict) -> dict:
        """Grow max_payload surface heads: keep prefix weights; init only new rows.

        Does **not** touch field/core. Used when resuming 16-byte hard ckpt into
        a larger max_payload (e.g. 32) under hard-v2 α-MLP.
        """
        info = {
            "span_head_migrated": False,
            "padded_keys": [],
            "old_max_payload": None,
            "new_max_payload": int(self.max_payload_bytes),
        }
        new_max = int(self.max_payload_bytes)
        # Infer old max from length_decoder bias/weight if present.
        old_len = surface_state.get("length_decoder.bias")
        if old_len is None:
            old_len = surface_state.get("length_decoder.weight")
        if old_len is None:
            old_len = surface_state.get("alpha_length_frozen")
        if old_len is None or not hasattr(old_len, "shape"):
            return info
        old_max = int(old_len.shape[0])
        info["old_max_payload"] = old_max
        if old_max >= new_max:
            return info  # shrink or equal — filtered load handles equal; shrink rare

        current = self.surface.state_dict()
        padded = []
        with torch.no_grad():
            for key, dst in current.items():
                if key not in surface_state:
                    continue
                src = surface_state[key]
                if not hasattr(src, "shape") or not hasattr(dst, "shape"):
                    continue
                if tuple(src.shape) == tuple(dst.shape):
                    continue
                # Only pad along output (dim0) growth for known span heads.
                if src.shape[0] >= dst.shape[0]:
                    continue
                if key.endswith(".fc1.weight") or key.endswith(".fc1.bias"):
                    continue  # MLP input layer not payload-sized
                n = self._pad_copy_leading(dst, src)
                if n > 0:
                    padded.append(key)
            # Re-apply padded tensors onto the module.
            if padded:
                patch = {k: current[k] for k in padded}
                self.surface.load_state_dict(patch, strict=False)
        info["span_head_migrated"] = bool(padded)
        info["padded_keys"] = padded
        return info


    def load(self, path: str | Path) -> dict | None:
        checkpoint = torch.load(path, map_location="cpu", weights_only=False)
        self.compiler.load_state_dict(checkpoint["compiler"])
        surface_state = checkpoint["surface"]
        # Legacy checkpoints predate ``state_norm`` and/or field-α readout.
        # Load matching tensors; activate new field path; damp biases that drown α.
        legacy_surface = "state_norm.weight" not in surface_state
        field_readout_missing = "field_to_state.weight" not in surface_state
        has_alpha_proj = any(
            key.startswith("alpha_byte_proj") for key in surface_state
        )
        obligatory_missing = (not has_alpha_proj) or ("alpha_jl" not in surface_state)
        current = self.surface.state_dict()
        filtered = {key: value for key, value in surface_state.items() if key in current and current[key].shape == value.shape}
        missing = [key for key in current if key not in filtered]
        has_la_keys = any(key.startswith("last_atom") for key in surface_state)
        want_la = getattr(self.surface, "last_atom", None) is not None
        if has_la_keys and not want_la:
            print("warning: checkpoint has last-atom weights but model built without last_atom_readout (keys dropped)")
        if want_la and not has_la_keys:
            print("info: last_atom_readout on but checkpoint lacks last-atom weights (fresh init kept)")
        self.surface.load_state_dict(filtered, strict=False)
        # If max_payload grew (e.g. 16→32), pad-copy span heads before any reinit wipe.
        span_migrate = self._migrate_span_head_weights(surface_state)
        if span_migrate.get("span_head_migrated"):
            # Recompute missing after pad: only truly-absent keys remain "missing".
            current = self.surface.state_dict()
            filtered_keys = set(filtered) | set(span_migrate.get("padded_keys") or [])
            missing = [key for key in current if key not in filtered_keys]
        mlp_missing = any(
            key.startswith("alpha_byte_proj") or key.startswith("alpha_length_proj")
            for key in missing
        )
        def _is_payload_key(key: str) -> bool:
            return (
                key.startswith("byte_embed")
                or key.startswith("payload_")
            )
        payload_missing = any(_is_payload_key(key) for key in missing)
        def _is_hard_couple_key(key: str) -> bool:
            # New opt-in head: must NOT trigger full field-readout reinit
            # (that would wipe tip field_feat_norm / field_to_state).
            return key.startswith("field_hard_couple_")

        non_mlp_missing = [
            key
            for key in missing
            if not (
                key.startswith("alpha_byte_proj")
                or key.startswith("alpha_length_proj")
                or key.startswith("last_atom")
                or _is_payload_key(key)
                or _is_hard_couple_key(key)
            )
        ]
        if payload_missing:
            with torch.no_grad():
                self.surface._init_payload_production()
        couple_missing = any(_is_hard_couple_key(key) for key in missing)
        if couple_missing:
            with torch.no_grad():
                nn.init.normal_(
                    self.surface.field_hard_couple_byte.weight, mean=0.0, std=0.02
                )
                nn.init.normal_(
                    self.surface.field_hard_couple_length.weight, mean=0.0, std=0.02
                )
        if legacy_surface or non_mlp_missing or field_readout_missing:
            with torch.no_grad():
                if legacy_surface or "state_norm.weight" in missing:
                    self.surface.state_norm.reset_parameters()
                # Shrink decoder bias so field-conditioned Wx / skip can compete.
                self.surface.byte_decoder.bias.mul_(0.05)
                self.surface.length_decoder.bias.mul_(0.05)
                if field_readout_missing or any(
                    (
                        (key.startswith("field_") and not _is_hard_couple_key(key))
                        or key in {"skip_gate"}
                    )
                    for key in missing
                ):
                    self.surface._init_field_readout()
                    self.surface.field_feat_norm.reset_parameters()
                    # Open skip strongly so migrated ckpts separate logits before retrain.
                    self.surface.skip_gate.fill_(4.0)
                    # _init_field_readout already seeds couple weights; keep them.

        frozen_alpha_missing = any(
            key in missing
            for key in (
                "alpha_jl",
                "alpha_byte_frozen",
                "alpha_length_frozen",
                "obl_gate",
                "obl_mix",
            )
        )
        if obligatory_missing or frozen_alpha_missing:
            with torch.no_grad():
                self.surface._init_obligatory_readout()
                # Slight bias damp so fresh α-only residual is audible on load.
                if not (legacy_surface or field_readout_missing):
                    self.surface.byte_decoder.bias.mul_(0.5)
                    self.surface.length_decoder.bias.mul_(0.5)
        elif mlp_missing:
            # Linear→MLP deepen: keep loaded frozen JL / maps; reinit trainable MLP only.
            with torch.no_grad():
                self.surface._init_alpha_trainable_mlp()
        core_state = checkpoint["core"]
        core = self.core
        core.encoder.load_state_dict(core_state["encoder"])
        core.state.load_state_dict(core_state["state"])
        core.dynamics.load_state_dict(core_state["dynamics"])
        core.interaction.load_state_dict(core_state["interaction"])
        core.aggregation.load_state_dict(core_state["aggregation"])
        core.abstraction.load_state_dict(core_state["abstraction"])
        core.consolidation.load_state_dict(core_state["consolidation"])
        core.production.load_state_dict(core_state["production"])
        core.atoms.load_state_dict(core_state["atoms"])
        core.energy_history = core_state.get("energy_history", [])
        core.training_loss = core_state.get("training_loss", 0.0)
        core.consolidation_count = core_state.get("consolidation_count", 0)
        core.abstraction_count = core_state.get("abstraction_count", 0)
        self.atomizer.load_state_dict(checkpoint["atomizer"])
        if "field_probe" in checkpoint:
            try:
                self.field_probe.load_state_dict(checkpoint["field_probe"])
            except Exception:
                pass  # feature_dim / d_model drift — keep fresh probe
        if "field_next_probe" in checkpoint:
            try:
                self.field_next_probe.load_state_dict(checkpoint["field_next_probe"])
            except Exception:
                pass
        cfg = checkpoint.get("config") or {}
        if "enable_merge" in cfg:
            self.enable_merge = bool(cfg["enable_merge"])
        if "merge_coherence_threshold" in cfg:
            self.merge_coherence_threshold = float(cfg["merge_coherence_threshold"])
            self.core.aggregation.phase_coherence_threshold = self.merge_coherence_threshold
        if "merge_energy_floor" in cfg:
            self.merge_energy_floor = float(cfg["merge_energy_floor"])
            self.core.aggregation.merge_energy_floor = self.merge_energy_floor
        if "slow_every" in cfg:
            self.slow_every = max(1, int(cfg["slow_every"]))
        if "field_loss_weight" in cfg:
            self.field_loss_weight = float(cfg["field_loss_weight"])
        if "field_contrast_weight" in cfg:
            self.field_contrast_weight = float(cfg["field_contrast_weight"])
        if "field_contrast_margin" in cfg:
            self.field_contrast_margin = float(cfg["field_contrast_margin"])
        if "field_ignorance_weight" in cfg:
            self.field_ignorance_weight = float(cfg["field_ignorance_weight"])
        if "field_ignorance_margin" in cfg:
            self.field_ignorance_margin = float(cfg["field_ignorance_margin"])
        if "field_ignorance_ablate_shared" in cfg:
            self.field_ignorance_ablate_shared = bool(cfg["field_ignorance_ablate_shared"])
        if "field_ignorance_prompt_bank" in cfg:
            self.field_ignorance_prompt_bank = bool(cfg["field_ignorance_prompt_bank"])
        if "field_obligatory_readout" in cfg:
            self.field_obligatory_readout = bool(cfg["field_obligatory_readout"])
        if "printable_aux_weight" in cfg:
            self.printable_aux_weight = float(cfg["printable_aux_weight"])
        if "field_next_packet_weight" in cfg:
            self.field_next_packet_weight = float(cfg["field_next_packet_weight"])
        if "field_obligatory_hard" in cfg and bool(cfg["field_obligatory_hard"]):
            self.field_obligatory_hard = True
            self.field_obligatory_readout = True
            self.surface.set_obligatory_hard(True)
        else:
            self.field_obligatory_hard = False
            self.surface.field_obligatory_hard = False
            # Constructor / CLI may set obligatory after load; always sync surface.
            self.surface.field_obligatory_readout = bool(self.field_obligatory_readout)
        self.merge_count_total = int(cfg.get("merge_count_total", 0) or 0)
        if "payload_copy" in cfg:
            self.surface.payload_enabled = bool(cfg["payload_copy"])
        # Provenance scale; CLI --last-atom-scale / --last-atom-scale-adaptive
        # win after load in run_atom_native.
        if "last_atom_scale" in cfg:
            self.surface.last_atom_scale = float(cfg["last_atom_scale"])
        if "last_atom_scale_adaptive" in cfg:
            self.surface.last_atom_scale_adaptive = bool(cfg["last_atom_scale_adaptive"])
        if "last_atom_scale_min" in cfg:
            self.surface.last_atom_scale_min = float(cfg["last_atom_scale_min"])
        if "last_atom_scale_max" in cfg:
            self.surface.last_atom_scale_max = float(cfg["last_atom_scale_max"])
        if "field_hard_couple_scale" in cfg:
            self.surface.field_hard_couple_scale = float(cfg["field_hard_couple_scale"])
        # Provenance only; CLI --last-atom-dropout wins after load in run_atom_native
        # (and P is ignored at eval/generation regardless).
        if "last_atom_dropout" in cfg:
            self.surface.last_atom_dropout = float(cfg["last_atom_dropout"])
        # Field leak changes inference dynamics, so it is restored from the
        # checkpoint (chat/eval see what was trained); trainer CLI --field-leak
        # wins after load.  Absent key (older ckpts) = 0.0 = legacy.
        if "field_leak" in cfg:
            self.field_leak = float(cfg["field_leak"] or 0.0)
        # Multi-timescale leak spectrum: same restore rule (absent/None = off).
        if "field_leak_spectrum" in cfg:
            self.set_field_leak_spectrum(
                cfg["field_leak_spectrum"],
                gain=cfg.get("field_leak_spectrum_gain", "none"),
                slow_frac=cfg.get("field_leak_spectrum_slow_frac", 1.0),
            )
        # Slow-mode readout (after the spectrum): restored with its weights so
        # chat/eval see the trained head.  Absent key = off (trainer CLI wins).
        src_cfg = checkpoint.get("slow_readout_config")
        if src_cfg and src_cfg.get("slow_mode_readout") and self.field_leak_spectrum is not None:
            self.set_slow_mode_readout(
                True, rank=int(src_cfg.get("rank", 64)), ahat=src_cfg.get("ahat", "exclude"),
                memory_weight=float(src_cfg.get("memory_weight", 0.0)),
                memory_lags=src_cfg.get("memory_lags", (16, 32, 64)),
                memory_target=src_cfg.get("memory_target", "bag"),
            )
            self.slow_readout.load_state_dict(checkpoint["slow_readout"])
        elif getattr(self, "slow_readout", None) is not None:
            self.set_slow_mode_readout(False)
        # --field-bptt / --field-leak-learnable: restored when present (absent =
        # off); the trainer CLI wins after load.
        self.set_field_bptt(int(checkpoint.get("field_bptt", 0) or 0))
        if "field_leak_logit" in checkpoint:
            self.set_field_leak_learnable(True)
            sb = checkpoint.get("field_leak_slow_block")
            if sb is not None:
                self.set_field_leak_slow_block(sb["frac"], cap=sb["cap"], gain=sb.get("gain", "ema"),
                                               ref=sb.get("ref"), reinit=False)
            elif getattr(self, "field_leak_slow_block", None) is not None:
                self.field_leak_slow_block = None
                self._slow_block_mask = None
                self._slow_block_floor = None
            with torch.no_grad():
                self.field_leak_logit.copy_(checkpoint["field_leak_logit"].to(self.field_leak_logit.device))
        else:
            self.set_field_leak_learnable(False)
            self.field_leak_slow_block = None
        # Always repair drifted dynamics (e.g. energy_decay~0.001 from 1.05M persist).
        dynamics_state = self.stabilize_dynamics_parameters()
        training = checkpoint.get("training")
        if training is None:
            training = {}
        training = dict(training)
        training["legacy_surface_migrated"] = bool(
            legacy_surface or bool(non_mlp_missing) or field_readout_missing
            or bool(span_migrate.get("span_head_migrated"))
        )
        training["alpha_mlp_migrated"] = bool(mlp_missing and not obligatory_missing)
        training["field_readout_migrated"] = bool(field_readout_missing)
        training["field_obligatory_migrated"] = bool(obligatory_missing)
        training["payload_prod_migrated"] = bool(payload_missing)
        training["span_head_migrated"] = bool(span_migrate.get("span_head_migrated"))
        training["span_head_migrate_info"] = span_migrate
        training["dynamics_on_load"] = dynamics_state
        return training