Blame for tools/run_atom_native.py

"""Run and document the atom-native ATOM experiment.

Usage:
    PYTHONPATH=. .venv/bin/python tools/run_atom_native.py \
      --data /tmp/wiki.train.raw --steps 500 \
      --output-dir checkpoints/atom_native_500

Streaming full-dataset ingest (no giant packet list in RAM):
    PYTHONPATH=. .venv/bin/python tools/run_atom_native.py \
      --stream --data-glob 'data/shards/*.txt' --steps 50000 \
      --output-dir checkpoints/atom_native_stream

The script intentionally uses the raw local corpus and never calls Hugging
Face or the legacy GPT-2 tokenizer.  It trains one transition per atom packet.
Episodes still pick random corpus starts in non-stream mode, but by default the
toroidal field is NOT wiped every short window (see ``--episode-length`` /
``--no-episode-reset``) so dialogue-scale persistence can accumulate.  Use
``--episode-reset`` to restore the old hard reset-every-episode behaviour.

In ``--stream`` mode the continuum is preferred: shards are read online and
``--no-episode-reset`` keeps the field alive across the infinite shard loop.
"""

from __future__ import annotations

import argparse
import atexit
import fcntl
import json
import os
import random
import time
from pathlib import Path

import numpy as np
import torch

from src.atom_native import DEFAULT_ENERGY_DECAY_BOUNDS, AtomNativeModel
from src.io.atomizer import Atomizer
from src.io.stream_corpus import StreamingPacketSource, resolve_shard_paths

# Advisory flock so two run_atom_native.py processes do not hammer one box by accident.
# Default path (cwd-relative): checkpoints/atom-ai-train.lock
# Fallback if checkpoints/ cannot be created: /tmp/atom-ai-train.lock
DEFAULT_TRAIN_LOCK_PATH = Path("checkpoints") / "atom-ai-train.lock"
FALLBACK_TRAIN_LOCK_PATH = Path("/tmp/atom-ai-train.lock")


def acquire_train_lock(lock_path: Path | None = None):
    """Non-blocking exclusive flock; fail-fast if another train holds it.

    Returns an open file object that must stay open for the lock lifetime.
    Writes this process PID into the lock file for diagnostics.
    """
    candidates: list[Path] = []
    if lock_path is not None:
        candidates.append(Path(lock_path))
    else:
        candidates.append(DEFAULT_TRAIN_LOCK_PATH)
        candidates.append(FALLBACK_TRAIN_LOCK_PATH)

    last_error: Exception | None = None
    for path in candidates:
        try:
            path.parent.mkdir(parents=True, exist_ok=True)
            fh = open(path, "a+", encoding="utf-8")
        except OSError as exc:
            last_error = exc
            continue
        try:
            fcntl.flock(fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            fh.seek(0)
            holder = fh.read().strip() or "unknown"
            fh.close()
            raise SystemExit(
                f"error: another train holds the advisory lock\n"
                f"  lock path: {path.resolve()}\n"
                f"  holder:    {holder}\n"
                f"  this PID:  {os.getpid()}\n"
                f"Pass --allow-parallel-train to opt out (not recommended on one box)."
            ) from None
        except OSError as exc:
            fh.close()
            last_error = exc
            continue
        fh.seek(0)
        fh.truncate()
        fh.write(f"pid={os.getpid()}\n")
        fh.flush()
        os.fsync(fh.fileno())

        def _release(handle=fh, lock=path) -> None:
            try:
                if not handle.closed:
                    fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
                    handle.close()
            except (OSError, ValueError):
                pass

        atexit.register(_release)
        print(f"train lock acquired: {path.resolve()} (pid={os.getpid()})")
        return fh

    raise SystemExit(
        f"error: could not create/acquire train lock "
        f"(tried {[str(p) for p in candidates]}): {last_error}"
    )


PROMPTS = ["The future of AI is", "Bonjour, je m'appelle", "Valkyria Chronicles III is"]


def seed_all(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def finite_model(model: AtomNativeModel) -> bool:
    return all(torch.isfinite(parameter).all().item() for parameter in model.parameters())


def transition_metrics(
    model: AtomNativeModel,
    current,
    target,
    train: bool,
    efference: bool = False,
    free_run_golds=None,
    free_run_weight: float = 1.0,
) -> tuple[float, dict]:
    if train:
        if free_run_golds is not None:
            loss, info = model.free_run_aux_loss(
                current, list(free_run_golds), weight=float(free_run_weight)
            )
        elif efference:
            loss, info = model.efference_loss(current, target)
        else:
            loss, info = model.transition_loss(current, target)
        if not torch.isfinite(loss):
            raise FloatingPointError("non-finite atom-native loss")
        loss.backward()
        slow_head = getattr(model, "slow_readout", None)
        if slow_head is None:
            grad_norm = torch.nn.utils.clip_grad_norm_(list(model.trainable_parameters), 1.0)
        else:
            # --slow-mode-readout: clip the new head on its own so its (fresh,
            # large) gradients never rescale the control's parameter updates.
            head_ids = {id(p) for p in slow_head.parameters()}
            grad_norm = torch.nn.utils.clip_grad_norm_(
                [p for p in model.trainable_parameters if id(p) not in head_ids], 1.0
            )
            head_norm = torch.nn.utils.clip_grad_norm_(
                [p for p in slow_head.parameters() if p.requires_grad], 1.0
            )
            info = {**info, "slow_head_grad_norm": float(head_norm)}
        if not torch.isfinite(torch.as_tensor(grad_norm)):
            raise FloatingPointError("non-finite atom-native gradient")
        return float(loss.detach().item()), {**info, "grad_norm": float(grad_norm)}
    with torch.no_grad():
        loss, info = model.transition_loss(current, target)
    return float(loss.item()), info


def _bptt_meta(args) -> dict:
    """Config keys for the opt-in levers, present only when on (off = files as before)."""
    out = {}
    if int(getattr(args, "field_bptt", 0) or 0) > 0:
        out["field_bptt"] = int(args.field_bptt)
    if getattr(args, "field_leak_learnable", False):
        out["field_leak_learnable"] = True
    if float(getattr(args, "field_leak_slow_block", 0.0) or 0.0) > 0:
        out["field_leak_slow_block"] = {"frac": float(args.field_leak_slow_block),
                                        "cap": float(args.field_leak_slow_cap), "gain": args.field_leak_slow_gain}
    if getattr(args, "legacy_gen_labels", False):
        out["legacy_gen_labels"] = True
    if int(getattr(args, "stream_skip_bytes", 0) or 0) > 0:
        out["stream_skip_bytes"] = int(args.stream_skip_bytes)
    return out


def bptt_window_step(model, optimizer, acc_loss, n_in_window):
    """Backprop the accumulated window loss (mean over its ticks), clip, step,
    and cut the field graph at the window boundary."""
    loss = acc_loss / float(max(n_in_window, 1))
    if not torch.isfinite(loss):
        raise FloatingPointError("non-finite atom-native loss (bptt window)")
    loss.backward()
    grad_norm = torch.nn.utils.clip_grad_norm_(list(model.trainable_parameters), 1.0)
    if not torch.isfinite(torch.as_tensor(grad_norm)):
        raise FloatingPointError("non-finite atom-native gradient (bptt window)")
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)
    model.detach_field_graph()
    return float(grad_norm)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--data",
        default=None,
        help="single corpus file (optional if --data-glob is set)",
    )
    parser.add_argument(
        "--data-glob",
        default=None,
        help="glob of shard files, e.g. 'data/shards/*.txt' (ordered lexicographically)",
    )
    parser.add_argument(
        "--stream",
        action="store_true",
        help="stream shards online (required for large corpora / globs); "
             "do not encode the whole file into one packet list",
    )
    parser.add_argument(
        "--chunk-bytes",
        type=int,
        default=1 << 20,
        help="UTF-8-safe read chunk size in stream mode (default 1048576)",
    )
    parser.add_argument(
        "--loop-shards",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="in stream mode, loop shards forever for an infinite continuum (default: loop)",
    )
    parser.add_argument(
        "--stream-buffer",
        type=int,
        default=2048,
        help="ring-buffer size of recent packets for stream-mode validation (default 2048)",
    )
    parser.add_argument(
        "--stream-prefetch",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="in --stream mode, prefetch next byte chunk on a background thread "
             "(default: on; use --no-stream-prefetch to disable)",
    )
    parser.add_argument(
        "--stream-skip-packets",
        type=int,
        default=0,
        help="in --stream mode, skip N packet transitions before training "
             "(default 0; use k*steps when chaining resumed chunks so each "
             "chunk sees fresh corpus instead of re-training the prefix)",
    )
    parser.add_argument(
        "--stream-skip-bytes",
        type=int,
        default=0,
        help="opt-in (default 0 = off): in --stream mode start the FIRST pass at this byte offset "
             "of the concatenated shards (seek, no replay; next codepoint start). In byte-tick "
             "(--max-span-bytes 1) 1 byte = 1 packet, so this replaces a long --stream-skip-packets "
             "replay on large corpora (atomizer state starts fresh at the offset)",
    )
    parser.add_argument(
        "--allow-parallel-train",
        action="store_true",
        default=False,
        help="opt out of the default advisory train lock "
             f"(default lock: {DEFAULT_TRAIN_LOCK_PATH})",
    )
    parser.add_argument(
        "--train-lock-path",
        default=None,
        help="override advisory lock path "
             f"(default: {DEFAULT_TRAIN_LOCK_PATH}; fallback {FALLBACK_TRAIN_LOCK_PATH})",
    )
    parser.add_argument("--output-dir", default="checkpoints/atom_native_chat")
    parser.add_argument("--checkpoint-name", default="atom_native.pt")
    parser.add_argument("--resume", default=None, help="optional checkpoint to resume")
    parser.add_argument("--steps", type=int, default=500)
    parser.add_argument(
        "--episode-length",
        type=int,
        default=512,
        help="packets per random corpus window (default 512; was 64)",
    )
    parser.add_argument(
        "--episode-reset",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="if set, wipe toroidal/atomizer state at each episode boundary "
             "(default: keep field across windows for dialogue continuity)",
    )
    parser.add_argument("--d-model", type=int, default=8)
    parser.add_argument("--n-modes", type=int, default=8)
    parser.add_argument("--n-atoms-max", type=int, default=128)
    parser.add_argument("--max-span-bytes", type=int, default=32,
                        help="Atomizer/surface max span bytes (default 32; was 16)")
    parser.add_argument(
        "--atomizer-pack",
        choices=("linguistic", "eager"),
        default="linguistic",
        help="Atomizer packing: linguistic=longer WS/punct-aligned spans (FR dialogue); eager=flush every boundary (legacy)",
    )
    parser.add_argument("--learning-rate", type=float, default=3e-4)
    parser.add_argument(
        "--surface-learning-rate",
        type=float,
        default=None,
        help="optional higher LR for AtomSurfaceHead only (param group); "
             "default: same as --learning-rate",
    )
    parser.add_argument(
        "--field-leak",
        type=float,
        default=0.0,
        help="per-tick field leak λ in [0,1): α ← (1-λ)·α before each atom "
             "injection (contractive recency memory; keeps α off the RMS cap). "
             "Default 0.0 = exact legacy behavior. CLI wins after --resume.",
    )
    parser.add_argument(
        "--field-leak-spectrum",
        type=float,
        nargs=2,
        default=None,
        metavar=("LMIN", "LMAX"),
        help="opt-in multi-timescale leak: per-mode λ_m geometric from LMIN (slow, "
             "modes with the least intrinsic decay, sin(mode_phase) high) to LMAX "
             "(fast); α_m ← (1-λ_m)·α_m before each injection. REPLACES --field-leak "
             "when set. Default off = exact current behavior. CLI wins after --resume "
             "(omit = off even if the checkpoint had one). See docs/FIELD_LEAK_SPECTRUM.md.",
    )
    parser.add_argument(
        "--field-leak-spectrum-gain",
        choices=["none", "ema", "sqrt"],
        default="none",
        help="per-mode input gain with --field-leak-spectrum: none (g=1), ema "
             "(g=λ_m/LMAX: every mode is an EMA of the tick input, LMAX modes = "
             "scalar λ=LMAX), sqrt (g=sqrt(λ_m/LMAX)). Ignored when the spectrum is off.",
    )
    parser.add_argument(
        "--field-leak-spectrum-slow-frac",
        type=float,
        default=1.0,
        help="with --field-leak-spectrum: fraction F in (0,1] of modes (those with the "
             "least intrinsic decay) that carry the LMIN->LMAX geometric ladder; the "
             "rest sit at LMAX. Default 1.0 = ladder over all modes.",
    )
    parser.add_argument(
        "--slow-mode-readout",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="opt-in: give the slow modes of --field-leak-spectrum (λ < LMAX) their own "
             "readout: Δlogits = W_o·tanh(W_t·ŝ), W_o zero-init (step-0 logits unchanged). "
             "Requires --field-leak-spectrum. Default off = exact current behavior. "
             "CLI wins after --resume. See docs/SLOW_MODE_READOUT.md.",
    )
    parser.add_argument("--slow-readout-rank", type=int, default=64,
                        help="rank r of the slow-mode readout trunk (default 64)")
    parser.add_argument(
        "--slow-readout-ahat", choices=["exclude", "mix"], default="exclude",
        help="with --slow-mode-readout: exclude (default) = the existing a_hat readout sees "
             "α with the slow modes zeroed; mix = slow modes stay in a_hat as well",
    )
    parser.add_argument(
        "--slow-readout-lr", type=float, default=None,
        help="LR of the slow-mode readout param group (default: --surface-learning-rate)",
    )
    parser.add_argument(
        "--slow-memory-loss", type=float, default=0.0, metavar="W",
        help="weight W of the slow-memory aux loss (needs --slow-mode-readout; default 0 = off)",
    )
    parser.add_argument(
        "--slow-memory-lags", type=int, nargs="+", default=[16, 32, 64], metavar="K",
        help="memory lags K for --slow-memory-loss (default 16 32 64)",
    )
    parser.add_argument(
        "--slow-memory-target", choices=["bag", "byte"], default="bag",
        help="bag (default): byte histogram of lags [K/2, K) from the slow head; "
             "byte: the single byte at lag K",
    )
    parser.add_argument(
        "--field-max-rms",
        type=float,
        default=3.0,
        help="RMS saturate cap (default 3.0; prior chat runs used 4.0). "
             "Tighter cap reduces saturate wipe while keeping signal.",
    )
    parser.add_argument(
        "--energy-decay-min",
        type=float,
        default=DEFAULT_ENERGY_DECAY_BOUNDS[0],
        help=f"always-on energy_decay floor (default {DEFAULT_ENERGY_DECAY_BOUNDS[0]})",
    )
    parser.add_argument(
        "--energy-decay-max",
        type=float,
        default=DEFAULT_ENERGY_DECAY_BOUNDS[1],
        help=f"always-on energy_decay ceiling (default {DEFAULT_ENERGY_DECAY_BOUNDS[1]})",
    )
    parser.add_argument("--seed", type=int, default=20260913)
    parser.add_argument(
        "--atom-flush-every",
        type=int,
        default=0,
        help="clear structural atom list every N steps (0=never, default). "
             "0 retains the multi-atom payload ring (MERGE keeps last ~128B) so "
             "α-gated copy-bias sees prompt-distinct bytes; set >0 only to bound cost.",
    )
    parser.add_argument(
        "--log-every",
        type=int,
        default=25,
        help="print/record trajectory every N steps",
    )
    parser.add_argument(
        "--enable-merge",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="MERGE phase-coherent energetic atoms into heavier structure (default: on)",
    )
    parser.add_argument(
        "--merge-coherence-threshold",
        type=float,
        default=0.45,
        help="phase*energy coherence gate for MERGE (default 0.45; mph~0.5 regime)",
    )
    parser.add_argument(
        "--merge-energy-floor",
        type=float,
        default=0.08,
        help="min mean-E / geo-mean energy floor for MERGE pairs (default 0.08)",
    )
    parser.add_argument(
        "--field-loss-weight",
        type=float,
        default=0.05,
        help="weight for persistence probe (reconstruct packet features from α); 0 disables",
    )
    parser.add_argument(
        "--field-contrast-weight",
        type=float,
        default=0.1,
        help="weight for α contrastive hinge on surface logits; 0 disables",
    )
    parser.add_argument(
        "--field-contrast-margin",
        type=float,
        default=0.55,
        help="max allowed cosine(true_α_logits, null/shuffled_α_logits) before hinge",
    )
    parser.add_argument(
        "--field-ignorance-weight",
        type=float,
        default=0.08,
        help="weight for field-ignorance hinge L_ign=ReLU(cos(l(a),l(sg[a_bar]))-m); 0 disables",
    )
    parser.add_argument(
        "--field-ignorance-margin",
        type=float,
        default=0.85,
        help="margin m for field-ignorance cosine hinge (default 0.85)",
    )
    parser.add_argument(
        "--field-ignorance-ablate-shared",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="for L_ign ℓ(α_bar) only: pass atom_r/persistence as None (zeros); "
             "ℓ(α) keeps live shared state (default: off)",
    )
    parser.add_argument(
        "--field-ignorance-prompt-bank",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="build α_bar from fixed distinct prompts (probe-style), not train ticks "
             "(default: off)",
    )
    parser.add_argument(
        "--field-obligatory-hard",
        action="store_true",
        help="hard-v2 obligatory: mix floor=1.0 + freeze non-alpha bypass; logits=frozen_α+scale*α_proj (implies readout)",
    )
    parser.add_argument(
        "--payload-copy",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="allow the α-local payload copy-bias branch on the hard CE path "
             "(default: on; use --no-payload-copy for the speech line so "
             "train matches chat payload_copy=False)",
    )
    parser.add_argument(
        "--last-atom-readout",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="byte-tick last-atom readout + dentate 2-gram on the hard path "
             "(requires --max-span-bytes 1; default: off)",
    )
    parser.add_argument(
        "--last-atom-scale",
        type=float,
        default=1.0,
        help="fixed scale on last_atom logits (hard mix): logits += S * la_byte "
             "(default 1.0 = current behavior; ignored when --last-atom-scale-adaptive)",
    )
    parser.add_argument(
        "--last-atom-scale-adaptive",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="opt-in: S_eff = clamp(α_rms/(la_rms+eps), S_min, S_max) each forward "
             "(detached RMS; overrides fixed --last-atom-scale; default OFF)",
    )
    parser.add_argument(
        "--last-atom-scale-min",
        type=float,
        default=0.0,
        help="adaptive clamp floor for S_eff (default 0.0 = true RMS match; "
             "0.05 floor falsified under free-run-aux)",
    )
    parser.add_argument(
        "--last-atom-scale-max",
        type=float,
        default=1.0,
        help="adaptive clamp ceiling for S_eff (default 1.0)",
    )
    parser.add_argument(
        "--last-atom-dropout",
        type=float,
        default=0.0,
        help="opt-in: training only, per byte-tick position zero the last_atom "
             "logit contribution with prob P (no inverted scaling) so CE must come "
             "from the field path; ignored at eval/generation (default 0 = current)",
    )
    parser.add_argument(
        "--field-hard-couple-scale",
        type=float,
        default=0.0,
        help="opt-in: under hard obligatory, add S * field_hard_couple_byte("
             "LN([mean‖std‖persist‖atom.r])) into hard logits "
             "(default 0 = OFF; dedicated W, not frozen field_byte_skip)",
    )
    parser.add_argument(
        "--legacy-gen-labels",
        action="store_true",
        help="byte-tick decode fix escape hatch (docs/DECODE_FIX_AND_BPTT.md): "
             "commit generated/efference bytes with the pre-fix labels "
             "(whitespace/punctuation). Default off = train-time labels.",
    )
    parser.add_argument(
        "--field-bptt",
        type=int,
        default=0,
        help="truncated BPTT through the toroidal field: keep the α graph and a "
             "differentiable atom injection within windows of K ticks; one "
             "backward + optimizer step per window (0=off = exact legacy "
             "per-tick detach + per-step update)",
    )
    parser.add_argument(
        "--field-leak-learnable",
        action="store_true",
        help="learnable per-mode leak λ_m = sigmoid(logit_m), init --field-leak "
             "(needs --field-bptt to receive gradient; off = scalar --field-leak)",
    )
    parser.add_argument(
        "--field-leak-slow-block",
        type=float,
        default=0.0,
        metavar="F",
        help="opt-in long memory (default 0 = off; needs --field-leak-learnable): a contiguous "
             "block of round(F*n_modes) modes around m=n/2 whose learnable leak is "
             "lambda = floor + CAP*sigmoid(logit) (floor keeps every mode stable), init sigmoid=.25, "
             "with EMA input gain lambda/ref (docs/MIX_DATA_LONG_MEMORY.md)",
    )
    parser.add_argument("--field-leak-slow-cap", type=float, default=0.02,
                        help="max learnable leak above the stability floor for slow-block modes (default .02)")
    parser.add_argument("--field-leak-slow-gain", choices=("ema", "none"), default="ema",
                        help="slow-block input gain: ema = lambda_m/ref (default), none = 1")
    parser.add_argument(
        "--efference-every",
        type=int,
        default=0,
        help="efference copy (bridge #9): every N train steps, commit the "
             "model's own predicted byte then CE on gold target (0=off; "
             "sandbox speech line used 10)",
    )
    parser.add_argument(
        "--free-run-aux-every",
        type=int,
        default=0,
        help="multi-tick free-run aux (L_roll): every N steps replace CE with "
             "H-step free-run CE using generate commit path (0=OFF default; "
             "does not change live recipe until opted in)",
    )
    parser.add_argument(
        "--free-run-aux-horizon",
        type=int,
        default=4,
        help="horizon H for --free-run-aux-every (default 4)",
    )
    parser.add_argument(
        "--free-run-aux-weight",
        type=float,
        default=1.0,
        help="scale on mean free-run CE when free-run aux fires (default 1.0)",
    )
    parser.add_argument(
        "--printable-aux-weight",
        type=float,
        default=0.0,
        help="aux weight rewarding printable UTF-8 mass on hard logits (0 disables; try ~0.08)",
    )
    parser.add_argument(
        "--field-next-packet-weight",
        type=float,
        default=0.0,
        help="aux weight: predict next packet features from α+atoms (0 disables; try ~0.05)",
    )
    parser.add_argument(
        "--field-obligatory-readout",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="main CE path: α-only residual into surface logits with non-zero floor "
             "(cannot bypass α via persist/atom_r alone; default: off)",
    )
    parser.add_argument(
        "--slow-every",
        type=int,
        default=1,
        help="run abstraction/consolidation every N ticks if field RMS stable (1=always)",
    )
    parser.add_argument(
        "--slow-rms-rel-tol",
        type=float,
        default=0.15,
        help="relative RMS change threshold for dual-clock slow path",
    )
    args = parser.parse_args()
    if not args.allow_parallel_train:
        acquire_train_lock(
            Path(args.train_lock_path) if args.train_lock_path else None
        )
    else:
        print("train lock skipped (--allow-parallel-train)")
    if not args.data and not args.data_glob:
        raise SystemExit("error: provide --data and/or --data-glob")
    if args.data_glob and not args.stream:
        raise SystemExit("error: --data-glob requires --stream (refuse full in-RAM encode of a shard glob)")
    if (args.energy_decay_min is None) != (args.energy_decay_max is None):
        raise ValueError("provide both --energy-decay-min and --energy-decay-max")
    energy_decay_bounds = (
        (float(args.energy_decay_min), float(args.energy_decay_max))
        if args.energy_decay_min is not None
        else DEFAULT_ENERGY_DECAY_BOUNDS
    )

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    stream_mode = bool(args.stream)
    shard_paths = resolve_shard_paths(args.data, args.data_glob)
    train_packets = None
    validation_packets = None
    packets = None
    raw = b""
    text = ""
    stream_source: StreamingPacketSource | None = None
    atomizer = Atomizer(max_span_bytes=args.max_span_bytes, pack_mode=args.atomizer_pack)

    if stream_mode:
        stream_source = StreamingPacketSource(
            atomizer,
            shard_paths,
            chunk_bytes=args.chunk_bytes,
            loop=args.loop_shards,
            buffer_size=args.stream_buffer,
            reset_atomizer=True,
            prefetch=bool(args.stream_prefetch),
        )
        # Peek one transition to fail fast on tiny/empty shards.
        try:
            first_pair = stream_source.next_transition()
        except StopIteration as exc:
            raise ValueError("stream source produced no packet transitions") from exc
        # Re-create source so training sees a clean continuum from the start.
        stream_source = StreamingPacketSource(
            Atomizer(max_span_bytes=args.max_span_bytes, pack_mode=args.atomizer_pack),
            shard_paths,
            chunk_bytes=args.chunk_bytes,
            loop=args.loop_shards,
            buffer_size=args.stream_buffer,
            reset_atomizer=True,
            prefetch=bool(args.stream_prefetch),
            **({"start_offset": int(args.stream_skip_bytes)} if int(args.stream_skip_bytes) > 0 else {}),
        )
        del first_pair
        if int(args.stream_skip_bytes) > 0:
            print(f"stream skip bytes: first pass starts at byte {int(args.stream_skip_bytes)}")
        skip = max(0, int(args.stream_skip_packets))
        if skip:
            print(f"stream skip: advancing {skip} packet transitions...")
            for _ in range(skip):
                try:
                    next(stream_source)
                except StopIteration as exc:
                    raise RuntimeError(
                        f"stream exhausted while skipping {skip} packets; "
                        "enable --loop-shards or reduce --stream-skip-packets"
                    ) from exc
            print(f"stream skip done: packets_seen={stream_source.packets_seen}")
        data_bytes = sum(p.stat().st_size for p in shard_paths)
        metadata_lengths_placeholder = True
    else:
        if len(shard_paths) != 1 or args.data_glob:
            raise SystemExit("non-stream mode supports a single --data file only; use --stream for globs")
        data_path = shard_paths[0]
        raw = data_path.read_bytes()
        text = raw.decode("utf-8")
        packets = atomizer.encode_bytes(raw)
        if len(packets) < 4:
            raise ValueError("the corpus must yield at least four atom packets")
        if b"".join(packet.payload for packet in packets) != raw:
            raise AssertionError("atomization is not lossless")
        train_packets = packets[:-64] if len(packets) > 96 else packets[:-4]
        validation_packets = packets[len(train_packets) :]
        if len(train_packets) < 3 or len(validation_packets) < 2:
            raise ValueError("not enough packets for train/validation split")
        data_bytes = len(raw)
        metadata_lengths_placeholder = False

    seed_all(args.seed)
    model = AtomNativeModel(
        d_model=args.d_model,
        n_modes=args.n_modes,
        n_atoms_max=args.n_atoms_max,
        max_payload_bytes=args.max_span_bytes,
        atomizer=Atomizer(max_span_bytes=args.max_span_bytes, pack_mode=args.atomizer_pack),
        field_max_rms=args.field_max_rms,
        field_leak=float(args.field_leak),
        energy_decay_bounds=energy_decay_bounds,
        enable_merge=bool(args.enable_merge),
        merge_coherence_threshold=float(args.merge_coherence_threshold),
        merge_energy_floor=float(args.merge_energy_floor),
        slow_every=int(args.slow_every),
        field_loss_weight=float(args.field_loss_weight),
        field_contrast_weight=float(args.field_contrast_weight),
        field_contrast_margin=float(args.field_contrast_margin),
        field_ignorance_weight=float(args.field_ignorance_weight),
        field_ignorance_margin=float(args.field_ignorance_margin),
        field_ignorance_ablate_shared=bool(args.field_ignorance_ablate_shared),
        field_ignorance_prompt_bank=bool(args.field_ignorance_prompt_bank),
        field_obligatory_readout=bool(args.field_obligatory_readout) or bool(args.field_obligatory_hard),
        field_obligatory_hard=bool(args.field_obligatory_hard),
        last_atom_readout=bool(args.last_atom_readout),
        printable_aux_weight=float(args.printable_aux_weight),
        field_next_packet_weight=float(args.field_next_packet_weight),
        slow_rms_rel_tol=float(args.slow_rms_rel_tol),
    )
    training_state = None
    start_step = 0
    if args.resume:
        training_state = model.load(args.resume)
        start_step = int((training_state or {}).get("step", 0) or 0)
        migrated = (training_state or {}).get("legacy_surface_migrated", False)
        span_mig = (training_state or {}).get("span_head_migrated", False)
        span_info = (training_state or {}).get("span_head_migrate_info") or {}
        dyn_load = (training_state or {}).get("dynamics_on_load") or {}
        print(
            f"resumed weights from {args.resume} "
            f"(recorded step={start_step}, legacy_surface_migrated={migrated}, "
            f"span_head_migrated={span_mig})"
        )
        if span_mig:
            print(
                f"span-head pad-migrate: old_max={span_info.get('old_max_payload')} -> "
                f"new_max={span_info.get('new_max_payload')} "
                f"padded={len(span_info.get('padded_keys') or [])} tensors "
                f"(field/core kept; new rows only)"
            )
        print(
            f"atomizer pack_mode={args.atomizer_pack} max_span_bytes={args.max_span_bytes}"
        )
        if dyn_load:
            print(
                f"dynamics on load: energy_decay="
                f"{dyn_load.get('energy_decay_before')} -> {dyn_load.get('energy_decay')} "
                f"(repaired={dyn_load.get('energy_decay_repaired')}, "
                f"bounds={dyn_load.get('energy_decay_bounds')})"
            )
        # CLI wins over ckpt config for obligatory readout (smoke / migrate).
        model.field_obligatory_readout = bool(args.field_obligatory_readout)
        model.surface.field_obligatory_readout = bool(args.field_obligatory_readout)
        # CLI wins for printable aux (mechanism under test).
        model.printable_aux_weight = float(args.printable_aux_weight)
        model.field_next_packet_weight = float(args.field_next_packet_weight)
        # CLI wins for field persistence / contrast weights (same pattern as
        # printable_aux / next_packet). Without this, model.load() restores
        # ckpt 0.0 and drops construct CLI e.g. --field-loss-weight 0.05.
        model.field_loss_weight = float(args.field_loss_weight)
        model.field_contrast_weight = float(args.field_contrast_weight)
        if float(args.field_next_packet_weight) > 0:
            print(
                f"field_next_packet_weight={float(args.field_next_packet_weight)} "
                "(predict next packet features from α+atoms)"
            )
        # CLI wins for MERGE gates (operational threshold fix under hard-v2+).
        model.enable_merge = bool(args.enable_merge)
        model.set_merge_thresholds(
            coherence_threshold=float(args.merge_coherence_threshold),
            energy_floor=float(args.merge_energy_floor),
        )
        print(
            f"MERGE gates: enable={bool(args.enable_merge)} "
            f"coherence_thr={float(args.merge_coherence_threshold):.3f} "
            f"energy_floor={float(args.merge_energy_floor):.3f}"
        )
        if getattr(args, "field_obligatory_hard", False):
            model.field_obligatory_hard = True
            model.field_obligatory_readout = True
            model.surface.set_obligatory_hard(True)
            print(
                "field_obligatory_hard=ON hard-v2+ "
                "(mix_floor=1.0, freeze bypass, logits=frozen+cap*(α_MLP+payload_prod))"
            )
            if float(args.printable_aux_weight) > 0:
                print(f"printable_aux_weight={float(args.printable_aux_weight)} (reward printable mass on hard logits)")
        if args.field_obligatory_readout:
            print(
                "field_obligatory_readout=ON "
                f"(migrated={bool((training_state or {}).get('field_obligatory_migrated'))})"
            )
    elif args.field_obligatory_readout:
        model.field_obligatory_readout = True
        model.surface.field_obligatory_readout = True
        print("field_obligatory_readout=ON (fresh)")
    # CLI wins for train-time payload copy-bias (after load, which restores
    # the checkpoint value for provenance). Speech line: --no-payload-copy.
    model.surface.payload_enabled = bool(args.payload_copy)
    print(f"payload_copy={bool(args.payload_copy)} (train-time copy-bias branch)")
    print(f"last_atom_readout={bool(args.last_atom_readout)} (byte-tick 2-gram + dentate)")
    # CLI wins for last_atom scale (default 1.0 / adaptive OFF = no behavior change).
    model.surface.last_atom_scale = float(args.last_atom_scale)
    model.surface.last_atom_scale_adaptive = bool(args.last_atom_scale_adaptive)
    model.surface.last_atom_scale_min = float(args.last_atom_scale_min)
    model.surface.last_atom_scale_max = float(args.last_atom_scale_max)
    # CLI wins for last_atom dropout (default 0 = no behavior change; train only).
    if not 0.0 <= float(args.last_atom_dropout) <= 1.0:
        raise SystemExit("error: --last-atom-dropout must be in [0, 1]")
    model.surface.last_atom_dropout = float(args.last_atom_dropout)
    print(
        f"last_atom_dropout={float(args.last_atom_dropout)} "
        "(train-only per-position zeroing of la_byte; 0=off/current)"
    )
    # CLI wins for field leak (default 0 = no-op = exact legacy dynamics).
    if not 0.0 <= float(args.field_leak) < 1.0:
        raise SystemExit("error: --field-leak must be in [0, 1)")
    model.field_leak = float(args.field_leak)
    print(
        f"field_leak={float(args.field_leak)} "
        "(per-tick α*=(1-λ) before injection; 0=off/current)"
    )
    # CLI wins for the multi-timescale leak spectrum too (omitted = off).
    try:
        model.set_field_leak_spectrum(
            tuple(args.field_leak_spectrum) if args.field_leak_spectrum else None,
            gain=args.field_leak_spectrum_gain,
            slow_frac=args.field_leak_spectrum_slow_frac,
        )
    except ValueError as exc:
        raise SystemExit(f"error: --field-leak-spectrum: {exc}")
    if model.field_leak_spectrum is not None:
        _lam = model.field_leak_rates()
        print(
            f"field_leak_spectrum={model.field_leak_spectrum} gain={model.field_leak_spectrum_gain} "
            f"slow_frac={model.field_leak_spectrum_slow_frac} "
            f"(per-mode λ min={float(_lam.min()):.4f} max={float(_lam.max()):.4f}; "
            f"replaces field_leak={float(args.field_leak)})"
        )
    # CLI wins for the slow-mode readout too (omitted = off, head dropped).
    try:
        if args.slow_mode_readout:
            model.set_slow_mode_readout(
                True, rank=int(args.slow_readout_rank), ahat=args.slow_readout_ahat,
                memory_weight=float(args.slow_memory_loss),
                memory_lags=tuple(args.slow_memory_lags),
                memory_target=args.slow_memory_target,
            )
        else:
            if float(args.slow_memory_loss) > 0:
                raise ValueError("--slow-memory-loss needs --slow-mode-readout")
            model.set_slow_mode_readout(False)
    except ValueError as exc:
        raise SystemExit(f"error: --slow-mode-readout: {exc}")
    if model.slow_readout is not None:
        print(
            f"slow_mode_readout=ON slow_modes={model._slow_idx.tolist()} rank={model.slow_readout_rank} "
            f"ahat={model.slow_readout_ahat} memory_weight={model.slow_memory_weight} "
            f"lags={list(model.slow_memory_lags)} target={model.slow_memory_target} "
            f"({sum(p.numel() for p in model.slow_readout.parameters())} params)"
        )
    # CLI wins for field→hard-decode couple (default 0 = no behavior change).
    model.surface.field_hard_couple_scale = float(args.field_hard_couple_scale)
    if float(args.field_hard_couple_scale) > 0:
        print(
            f"field_hard_couple_scale={float(args.field_hard_couple_scale)} "
            "(hard logits += S * field_hard_couple_byte(full feats))"
        )
    else:
        print("field_hard_couple_scale=0 (OFF — hard path unchanged)")
    if bool(args.last_atom_scale_adaptive):
        print(
            f"last_atom_scale=adaptive "
            f"(S_eff=clamp(α_rms/(la_rms+eps), "
            f"{float(args.last_atom_scale_min)}, {float(args.last_atom_scale_max)}); "
            f"fixed S={float(args.last_atom_scale)} ignored)"
        )
    else:
        print(
            f"last_atom_scale={float(args.last_atom_scale)} "
            "(1.0=off/current; hard-path la_byte *= S)"
        )
    model._efference_every = max(0, int(args.efference_every))
    print(f"efference_every={model._efference_every} (0=off)")
    # CLI wins for the byte-tick decode fix escape hatch (runtime-only flag).
    model.atomizer.legacy_gen_labels = bool(args.legacy_gen_labels)
    if args.legacy_gen_labels:
        print("legacy_gen_labels=ON (pre-fix whitespace/punctuation commit labels)")
    # CLI wins for truncated BPTT / learnable leak (after load, both directions).
    if int(args.field_bptt) < 0:
        raise SystemExit("error: --field-bptt must be >= 0")
    model.set_field_bptt(int(args.field_bptt))
    if args.field_leak_learnable:
        if model.field_leak_logit is None:
            try:
                model.set_field_leak_learnable(True, init=float(args.field_leak))
            except ValueError as exc:
                raise SystemExit(f"error: --field-leak-learnable: {exc}")
        _ll = model.learnable_leak_rates().detach()
        print(f"field_leak_learnable=ON per-mode λ min={float(_ll.min()):.4f} max={float(_ll.max()):.4f}")
    else:
        model.set_field_leak_learnable(False)
    # CLI wins for the slow block too (after load, both directions; same config = kept).
    _sb_f = float(args.field_leak_slow_block or 0.0)
    _sb_cur = getattr(model, "field_leak_slow_block", None)
    if _sb_f > 0:
        if not args.field_leak_learnable:
            raise SystemExit("error: --field-leak-slow-block needs --field-leak-learnable")
        _same = _sb_cur is not None and (_sb_cur["frac"], _sb_cur["cap"], _sb_cur["gain"]) == (
            _sb_f, float(args.field_leak_slow_cap), args.field_leak_slow_gain)
        if not _same:
            model.set_field_leak_slow_block(None)
            model.set_field_leak_slow_block(_sb_f, cap=float(args.field_leak_slow_cap),
                                            gain=args.field_leak_slow_gain, ref=float(args.field_leak))
        _ll = model.learnable_leak_rates().detach().reshape(-1)
        print(f"field_leak_slow_block=ON modes={model.field_leak_slow_block['modes']} "
              f"cap={model.field_leak_slow_block['cap']} gain={model.field_leak_slow_block['gain']} "
              f"slow λ={[round(float(_ll[i]), 4) for i in model.field_leak_slow_block['modes']]}")
    elif _sb_cur is not None:
        model.set_field_leak_slow_block(None)
    if model.field_bptt > 0:
        print(f"field_bptt={model.field_bptt} (truncated BPTT window; one optimizer step per window)")
    model._free_run_aux_every = max(0, int(args.free_run_aux_every))
    model._free_run_aux_horizon = max(1, int(args.free_run_aux_horizon))
    model._free_run_aux_weight = float(args.free_run_aux_weight)
    print(
        f"free_run_aux_every={model._free_run_aux_every} "
        f"horizon={model._free_run_aux_horizon} "
        f"weight={model._free_run_aux_weight} (0=off)"
    )
    surface_lr = (
        float(args.surface_learning_rate)
        if args.surface_learning_rate is not None
        else float(args.learning_rate)
    )
    surface_ids = {id(p) for p in model.surface.parameters() if p.requires_grad}
    slow_ids = (
        {id(p) for p in model.slow_readout.parameters()} if model.slow_readout is not None else set()
    )
    surface_params = [p for p in model.trainable_parameters if id(p) in surface_ids]
    other_params = [
        p for p in model.trainable_parameters if id(p) not in surface_ids and id(p) not in slow_ids
    ]
    param_groups = [
        {"params": other_params, "lr": float(args.learning_rate)},
        {"params": surface_params, "lr": surface_lr},
    ]
    if slow_ids:
        slow_lr = float(args.slow_readout_lr) if args.slow_readout_lr is not None else surface_lr
        param_groups.append(
            {"params": [p for p in model.slow_readout.parameters() if p.requires_grad], "lr": slow_lr}
        )
        print(f"optimizer group: slow_readout lr={slow_lr}")
    optimizer = torch.optim.AdamW(param_groups, weight_decay=0.01)
    print(
        f"optimizer groups: field/other lr={args.learning_rate} "
        f"({len(other_params)} tensors), surface lr={surface_lr} "
        f"({len(surface_params)} tensors); "
        f"energy_decay_bounds={energy_decay_bounds}; field_max_rms={args.field_max_rms}"
    )
    if stream_mode:
        print(
            f"stream mode: {len(shard_paths)} shard(s), chunk_bytes={args.chunk_bytes}, "
            f"loop_shards={args.loop_shards}, stream_buffer={args.stream_buffer}, "
            f"stream_prefetch={bool(args.stream_prefetch)}"
        )
    if training_state and "optimizer" in training_state:
        if training_state.get("legacy_surface_migrated"):
            print("optimizer restore skipped: surface LayerNorm migration changed bias scale")
        elif args.surface_learning_rate is not None:
            print("optimizer restore skipped: surface LR param groups differ from checkpoint")
        else:
            try:
                optimizer.load_state_dict(training_state["optimizer"])
                print("restored optimizer state")
            except Exception as exc:
                print(f"optimizer restore skipped: {exc}")
    model.train()

    if metadata_lengths_placeholder:
        metadata = {
            "stream": True,
            "data_path": str(args.data) if args.data else None,
            "data_glob": args.data_glob,
            "shard_paths": [str(p) for p in shard_paths],
            "data_bytes": data_bytes,
            "data_characters": None,
            "atomizer_version": atomizer.VERSION,
            "feature_dim": atomizer.feature_dim,
            "packet_count": None,
            "train_packet_count": None,
            "validation_packet_count": None,
            "mean_packet_bytes": None,
            "max_packet_bytes": None,
            "min_packet_bytes": None,
            "bytes_seen": 0,
            "config": {
                "d_model": args.d_model,
                "n_modes": args.n_modes,
                "n_atoms_max": args.n_atoms_max,
                "max_span_bytes": args.max_span_bytes,
                "episode_length": args.episode_length,
                "episode_reset": bool(args.episode_reset),
                "atom_flush_every": args.atom_flush_every,
                "learning_rate": args.learning_rate,
                "surface_learning_rate": surface_lr,
                "steps": args.steps,
                "start_step": start_step,
                "seed": args.seed,
                "surface_alphabet": 256,
                "field_max_rms": args.field_max_rms,
                "energy_decay_bounds": energy_decay_bounds,
                "stream": True,
                "chunk_bytes": args.chunk_bytes,
                "loop_shards": bool(args.loop_shards),
                "stream_buffer": args.stream_buffer,
                "stream_prefetch": bool(args.stream_prefetch),
                "stream_skip_packets": max(0, int(args.stream_skip_packets)),
                "last_atom_readout": bool(args.last_atom_readout),
                "last_atom_scale": float(args.last_atom_scale),
                "last_atom_scale_adaptive": bool(args.last_atom_scale_adaptive),
                "last_atom_scale_min": float(args.last_atom_scale_min),
                "last_atom_scale_max": float(args.last_atom_scale_max),
                "field_hard_couple_scale": float(args.field_hard_couple_scale),
                "last_atom_dropout": float(args.last_atom_dropout),
                "field_leak": float(args.field_leak),
                "field_leak_spectrum": list(args.field_leak_spectrum) if args.field_leak_spectrum else None,
                "field_leak_spectrum_gain": (
                    args.field_leak_spectrum_gain if args.field_leak_spectrum else "none"
                ),
                "field_leak_spectrum_slow_frac": (
                    float(args.field_leak_spectrum_slow_frac) if args.field_leak_spectrum else 1.0
                ),
                "slow_mode_readout": bool(args.slow_mode_readout),
                "slow_readout_rank": int(args.slow_readout_rank),
                "slow_readout_ahat": args.slow_readout_ahat,
                "slow_readout_lr": args.slow_readout_lr,
                "slow_memory_loss": float(args.slow_memory_loss),
                "slow_memory_lags": list(args.slow_memory_lags),
                "slow_memory_target": args.slow_memory_target,
                "efference_every": max(0, int(args.efference_every)),
                "free_run_aux_every": max(0, int(args.free_run_aux_every)),
                "free_run_aux_horizon": max(1, int(args.free_run_aux_horizon)),
                "free_run_aux_weight": float(args.free_run_aux_weight),
                "enable_merge": bool(args.enable_merge),
                "payload_copy": bool(args.payload_copy),
                "merge_coherence_threshold": float(args.merge_coherence_threshold),
                "merge_energy_floor": float(args.merge_energy_floor),
                "field_loss_weight": float(args.field_loss_weight),
                "field_contrast_weight": float(args.field_contrast_weight),
                "field_ignorance_weight": float(args.field_ignorance_weight),
                "field_ignorance_ablate_shared": bool(args.field_ignorance_ablate_shared),
                "field_ignorance_prompt_bank": bool(args.field_ignorance_prompt_bank),
                "field_obligatory_readout": bool(args.field_obligatory_readout) or bool(getattr(args, "field_obligatory_hard", False)),
                "field_obligatory_hard": bool(getattr(args, "field_obligatory_hard", False)),
                "printable_aux_weight": float(args.printable_aux_weight),
                "field_next_packet_weight": float(args.field_next_packet_weight),
                "slow_every": int(args.slow_every),
                **_bptt_meta(args),
            },
        }
    else:
        lengths = [len(packet.payload) for packet in packets]
        metadata = {
            "stream": False,
            "data_path": str(args.data),
            "data_glob": None,
            "shard_paths": [str(p) for p in shard_paths],
            "data_bytes": data_bytes,
            "data_characters": len(text),
            "atomizer_version": atomizer.VERSION,
            "feature_dim": atomizer.feature_dim,
            "packet_count": len(packets),
            "train_packet_count": len(train_packets),
            "validation_packet_count": len(validation_packets),
            "mean_packet_bytes": sum(lengths) / len(lengths),
            "max_packet_bytes": max(lengths),
            "min_packet_bytes": min(lengths),
            "config": {
                "d_model": args.d_model,
                "n_modes": args.n_modes,
                "n_atoms_max": args.n_atoms_max,
                "max_span_bytes": args.max_span_bytes,
                "episode_length": args.episode_length,
                "episode_reset": bool(args.episode_reset),
                "atom_flush_every": args.atom_flush_every,
                "learning_rate": args.learning_rate,
                "surface_learning_rate": surface_lr,
                "steps": args.steps,
                "start_step": start_step,
                "seed": args.seed,
                "surface_alphabet": 256,
                "field_max_rms": args.field_max_rms,
                "energy_decay_bounds": energy_decay_bounds,
                "stream": False,
                "enable_merge": bool(args.enable_merge),
                "payload_copy": bool(args.payload_copy),
                "last_atom_readout": bool(args.last_atom_readout),
                "last_atom_scale": float(args.last_atom_scale),
                "last_atom_scale_adaptive": bool(args.last_atom_scale_adaptive),
                "last_atom_scale_min": float(args.last_atom_scale_min),
                "last_atom_scale_max": float(args.last_atom_scale_max),
                "field_hard_couple_scale": float(args.field_hard_couple_scale),
                "last_atom_dropout": float(args.last_atom_dropout),
                "field_leak": float(args.field_leak),
                "field_leak_spectrum": list(args.field_leak_spectrum) if args.field_leak_spectrum else None,
                "field_leak_spectrum_gain": (
                    args.field_leak_spectrum_gain if args.field_leak_spectrum else "none"
                ),
                "field_leak_spectrum_slow_frac": (
                    float(args.field_leak_spectrum_slow_frac) if args.field_leak_spectrum else 1.0
                ),
                "slow_mode_readout": bool(args.slow_mode_readout),
                "slow_readout_rank": int(args.slow_readout_rank),
                "slow_readout_ahat": args.slow_readout_ahat,
                "slow_readout_lr": args.slow_readout_lr,
                "slow_memory_loss": float(args.slow_memory_loss),
                "slow_memory_lags": list(args.slow_memory_lags),
                "slow_memory_target": args.slow_memory_target,
                "efference_every": max(0, int(args.efference_every)),
                "free_run_aux_every": max(0, int(args.free_run_aux_every)),
                "free_run_aux_horizon": max(1, int(args.free_run_aux_horizon)),
                "free_run_aux_weight": float(args.free_run_aux_weight),
                "merge_coherence_threshold": float(args.merge_coherence_threshold),
                "merge_energy_floor": float(args.merge_energy_floor),
                "field_loss_weight": float(args.field_loss_weight),
                "field_contrast_weight": float(args.field_contrast_weight),
                "field_ignorance_weight": float(args.field_ignorance_weight),
                "field_ignorance_ablate_shared": bool(args.field_ignorance_ablate_shared),
                "field_ignorance_prompt_bank": bool(args.field_ignorance_prompt_bank),
                "field_obligatory_readout": bool(args.field_obligatory_readout) or bool(getattr(args, "field_obligatory_hard", False)),
                "field_obligatory_hard": bool(getattr(args, "field_obligatory_hard", False)),
                "printable_aux_weight": float(args.printable_aux_weight),
                "field_next_packet_weight": float(args.field_next_packet_weight),
                "slow_every": int(args.slow_every),
                **_bptt_meta(args),
            },
        }
    (output_dir / "atomization_stats.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8")

    trajectory: list[dict] = []
    all_losses: list[float] = []
    efference_steps = 0
    free_run_aux_steps = 0
    field_rms_before_values: list[float] = []
    field_rms_after_values: list[float] = []
    field_scales: list[float] = []
    field_norm_values: list[float] = []
    slow_mem_values: list[float] = []
    start_time = time.perf_counter()
    episode = -1
    episode_start = 0
    # Soft cap: if atoms approach n_atoms_max, reset once to avoid unbounded growth
    # when --no-episode-reset keeps the field alive across windows.
    atom_soft_cap = max(16, int(args.n_atoms_max * 0.9))
    bptt_acc = None
    bptt_n = 0
    bptt_windows = 0
    bptt_last_grad = float("nan")
    dynamics_state = model.stabilize_dynamics_parameters() if int(getattr(model, "field_bptt", 0) or 0) else None
    for step in range(1, args.steps + 1):
        current_episode = (step - 1) // args.episode_length
        free_run_golds = None
        do_free_run = bool(args.free_run_aux_every) and step % int(args.free_run_aux_every) == 0
        horizon = max(1, int(args.free_run_aux_horizon)) if do_free_run else 1

        if stream_mode:
            # Prefer continuous stream. Episode windows only force a field wipe
            # when --episode-reset is set; otherwise the field persists.
            if args.episode_reset and current_episode != episode:
                episode = current_episode
                model.reset_state(reset_atomizer=True)
                model.train()
            elif step == 1:
                episode = current_episode
                model.reset_state(reset_atomizer=True)
                model.train()
            else:
                episode = current_episode
            try:
                # Prefer iterator protocol (avoids extra Python method frame).
                current_packet, target_packet = next(stream_source)
                if do_free_run:
                    golds = [target_packet]
                    for _ in range(horizon - 1):
                        _, g = next(stream_source)
                        golds.append(g)
                    free_run_golds = golds
            except StopIteration as exc:
                raise RuntimeError(
                    f"stream exhausted at step {step}/{args.steps}; "
                    "enable --loop-shards for infinite continuum"
                ) from exc
        else:
            if current_episode != episode:
                episode = current_episode
                if args.episode_reset or step == 1:
                    model.reset_state(reset_atomizer=True)
                else:
                    # Continuity-preserving compromise: at each window boundary, clear
                    # the structural atom list (keeps step cost bounded) but KEEP the
                    # toroidal field alpha + consolidation/abstraction persistence so
                    # dialogue-scale memory is not wiped every 64/512 ticks.
                    model.core.atoms = type(model.core.atoms)()
                    model.core.energy_history = []
                    if len(model.core.atoms) >= atom_soft_cap:
                        pass  # already cleared
                model.train()
                max_start = max(0, len(train_packets) - args.episode_length - 2)
                episode_start = random.randint(0, max_start) if max_start > 0 else 0
            index = episode_start + ((step - 1) % args.episode_length)
            if index >= len(train_packets) - 1:
                index = index % (len(train_packets) - 1)
            current_packet = train_packets[index]
            target_packet = train_packets[index + 1]
            if do_free_run:
                golds = []
                for h in range(horizon):
                    gi = index + 1 + h
                    if gi >= len(train_packets):
                        break
                    golds.append(train_packets[gi])
                if golds:
                    free_run_golds = golds

        bptt_k = int(getattr(model, "field_bptt", 0) or 0)
        if not bptt_k:
            optimizer.zero_grad(set_to_none=True)
        # Free-run aux supersedes 1-tick efference when both would fire.
        do_eff = (
            free_run_golds is None
            and bool(args.efference_every)
            and step % int(args.efference_every) == 0
        )
        if do_eff:
            efference_steps += 1
        if do_free_run and free_run_golds is not None:
            free_run_aux_steps += 1
        if bptt_k:
            # --field-bptt K: accumulate K tick losses, one backward through the
            # field graph, one optimizer step, cut the graph (window boundary).
            if free_run_golds is not None:
                loss_t, info = model.free_run_aux_loss(
                    current_packet, list(free_run_golds), weight=float(args.free_run_aux_weight)
                )
            elif do_eff:
                loss_t, info = model.efference_loss(current_packet, target_packet)
            else:
                loss_t, info = model.transition_loss(current_packet, target_packet)
            if not torch.isfinite(loss_t):
                raise FloatingPointError("non-finite atom-native loss")
            bptt_acc = loss_t if bptt_n == 0 else bptt_acc + loss_t
            bptt_n += 1
            loss = float(loss_t.detach().item())
            episode_ends = bool(args.episode_reset) and step % args.episode_length == 0
            if bptt_n >= bptt_k or step == args.steps or episode_ends:
                bptt_last_grad = bptt_window_step(model, optimizer, bptt_acc, bptt_n)
                bptt_acc = None
                bptt_n = 0
                bptt_windows += 1
                dynamics_state = model.stabilize_dynamics_parameters()
            info = {**info, "grad_norm": bptt_last_grad}
        else:
            loss, info = transition_metrics(
                model,
                current_packet,
                target_packet,
                train=True,
                efference=do_eff,
                free_run_golds=free_run_golds,
                free_run_weight=float(args.free_run_aux_weight),
            )
            optimizer.step()
        if args.atom_flush_every and step % args.atom_flush_every == 0:
            model.core.atoms = type(model.core.atoms)()
            model.core.energy_history = []
        if not bptt_k:
            dynamics_state = model.stabilize_dynamics_parameters()
        model.core.training_loss = loss
        all_losses.append(loss)
        field_rms_before_values.append(info["field_rms_before"])
        field_rms_after_values.append(info["field_rms_after"])
        field_scales.append(info["field_scale"])
        field_norm_values.append(info["field_norm"])
        if "slow_memory_loss" in info:
            slow_mem_values.append(float(info["slow_memory_loss"]))
        global_step = start_step + step
        if step == 1 or step % args.log_every == 0 or step == args.steps:
            record = {
                "step": global_step,
                "run_step": step,
                "episode": episode,
                "loss": loss,
                "total_perplexity": float(np.exp(min(loss, 30.0))),
                "byte_loss": info["byte_loss"],
                "byte_perplexity": float(np.exp(min(info["byte_loss"], 30.0))),
                "length_loss": info["length_loss"],
                "field_norm": info["field_norm"],
                "field_rms_before": info["field_rms_before"],
                "field_rms_after": info["field_rms_after"],
                "field_scale": info["field_scale"],
                "grad_norm": info.get("grad_norm"),
                "n_atoms": info["n_atoms"],
                "efference": do_eff,
                "free_run_aux": bool(info.get("free_run_aux", False)),
                "free_run_horizon": int(info.get("free_run_horizon", 0) or 0),
                "merge_count": info.get("merge_count", 0),
                "merge_count_total": info.get("merge_count_total", 0),
                "slow_tick": info.get("slow_tick", True),
                "field_contrast_loss": info.get("field_contrast_loss", 0.0),
                "field_persist_loss": info.get("field_persist_loss", 0.0),
                "field_ignorance_loss": info.get("field_ignorance_loss", 0.0),
                "printable_aux_loss": info.get("printable_aux_loss", 0.0),
                "printable_mass_mean": info.get("printable_mass_mean", float("nan")),
                "field_next_packet_loss": info.get("field_next_packet_loss", 0.0),
                "max_phase_coherence": info.get("max_phase_coherence", 0.0),
                "n_pairs_above_energy_floor": info.get("n_pairs_above_energy_floor", 0),
                "energy_decay": dynamics_state["energy_decay"],
                "coupling_scale": dynamics_state["coupling_scale"],
                "phase_sync": dynamics_state["phase_sync"],
            }
            if "slow_memory_loss" in info:
                record["slow_memory_loss"] = info["slow_memory_loss"]
                record["slow_head_grad_norm"] = info.get("slow_head_grad_norm")
            if stream_mode and stream_source is not None:
                record["bytes_seen"] = stream_source.bytes_seen
                record["packets_seen"] = stream_source.packets_seen
                record["shard_index"] = stream_source.shard_index
                record["stream_epoch"] = stream_source.epoch
            trajectory.append(record)
            extra = ""
            if stream_mode and stream_source is not None:
                extra = (
                    f" bytes={stream_source.bytes_seen} "
                    f"pkts={stream_source.packets_seen} "
                    f"shard={stream_source.shard_index}"
                )
            print(
                f"step={global_step:07d} (run {step}/{args.steps}) "
                f"loss={loss:.5f} byte_ppl={record['byte_perplexity']:.3f} "
                f"atoms={record['n_atoms']} merges={record['merge_count_total']} "
                f"field={record['field_norm']:.4f} "
                f"rms={record['field_rms_after']:.4f} scale={record['field_scale']:.4f} "
                f"fcos={info.get('field_logit_cos_zero', float('nan')):.3f} "
                f"ign={info.get('field_ignorance_loss', 0.0):.4f} "
                f"prn={info.get('printable_mass_mean', float('nan')):.3f} "
                f"nxp={info.get('field_next_packet_loss', 0.0):.3f} "
                f"mph={info.get('max_phase_coherence', 0.0):.3f}"
                f"{extra}"
            )
        if not finite_model(model):
            raise FloatingPointError(f"non-finite parameter at step {step}")

    elapsed = time.perf_counter() - start_time
    if stream_mode and stream_source is not None:
        metadata["bytes_seen"] = stream_source.bytes_seen
        metadata["packets_seen"] = stream_source.packets_seen
        metadata["stream_epoch"] = stream_source.epoch
        metadata["validation_packet_count"] = len(stream_source.validation_packets())
        (output_dir / "atomization_stats.json").write_text(
            json.dumps(metadata, indent=2), encoding="utf-8"
        )

    checkpoint_path = output_dir / args.checkpoint_name
    model.save(
        checkpoint_path,
        extra_state={
            "step": start_step + args.steps,
            "run_steps": args.steps,
            "optimizer": optimizer.state_dict(),
            "trajectory": trajectory,
            "seed": args.seed,
            "episode_length": args.episode_length,
            "episode_reset": bool(args.episode_reset),
            "atom_flush_every": args.atom_flush_every,
            "stream": stream_mode,
            "bytes_seen": (stream_source.bytes_seen if stream_source else data_bytes),
        },
    )

    # Validation: non-stream uses held-out tail; stream uses ring-buffer samples.
    model.reset_state(reset_atomizer=True)
    model.eval()
    validation_records = []
    if stream_mode and stream_source is not None:
        val_pairs = stream_source.validation_transitions()
        # Cap validation work on huge rings.
        if len(val_pairs) > 256:
            val_pairs = val_pairs[-256:]
    else:
        val_pairs = list(zip(validation_packets[:-1], validation_packets[1:]))
    with torch.no_grad():
        for current, target in val_pairs:
            loss, info = transition_metrics(model, current, target, train=False)
            validation_records.append({"loss": loss, **info})
    if validation_records:
        validation_loss = float(np.mean([item["loss"] for item in validation_records]))
        validation_byte_loss = float(np.mean([item["byte_loss"] for item in validation_records]))
    else:
        validation_loss = float("nan")
        validation_byte_loss = float("nan")

    # Reload the final checkpoint before generation, then compare deterministic
    # samples so a restart cannot silently change the atom-native state path.
    reloaded = AtomNativeModel(
        d_model=args.d_model,
        n_modes=args.n_modes,
        n_atoms_max=args.n_atoms_max,
        max_payload_bytes=args.max_span_bytes,
        atomizer=Atomizer(max_span_bytes=args.max_span_bytes, pack_mode=args.atomizer_pack),
        field_max_rms=args.field_max_rms,
        energy_decay_bounds=energy_decay_bounds,
        enable_merge=bool(args.enable_merge),
        slow_every=int(args.slow_every),
        field_loss_weight=float(args.field_loss_weight),
        field_contrast_weight=float(args.field_contrast_weight),
        field_contrast_margin=float(args.field_contrast_margin),
        field_ignorance_weight=float(args.field_ignorance_weight),
        field_ignorance_margin=float(args.field_ignorance_margin),
        field_ignorance_ablate_shared=bool(args.field_ignorance_ablate_shared),
        field_ignorance_prompt_bank=bool(args.field_ignorance_prompt_bank),
        field_obligatory_readout=bool(args.field_obligatory_readout) or bool(args.field_obligatory_hard),
        field_obligatory_hard=bool(args.field_obligatory_hard),
        last_atom_readout=bool(args.last_atom_readout),
        printable_aux_weight=float(args.printable_aux_weight),
        slow_rms_rel_tol=float(args.slow_rms_rel_tol),
    )
    reloaded.load(checkpoint_path)
    reload_state = {
        "n_atoms": len(reloaded.core.atoms),
        "field_norm": float(reloaded.core.state.alpha.detach().norm().item()),
        "time": float(reloaded.core.state.t.item()),
        "finite": bool(torch.isfinite(reloaded.core.state.alpha).all().item()),
    }
    generation = {}
    for prompt in PROMPTS:
        torch.manual_seed(args.seed + 1)
        source_bytes = model.generate_packets(
            prompt,
            max_packets=8,
            max_length=64,
            temperature=0.8,
            top_k=5,
            deterministic=False,
        )
        torch.manual_seed(args.seed + 1)
        reload_bytes = reloaded.generate_packets(
            prompt,
            max_packets=8,
            max_length=64,
            temperature=0.8,
            top_k=5,
            deterministic=False,
        )
        generation[prompt] = {
            "temperature": 0.8,
            "top_k": 5,
            "max_length": 20,
            "bytes_hex": source_bytes.hex(),
            "text": source_bytes.decode("utf-8", errors="replace"),
            "reload_text": reload_bytes.decode("utf-8", errors="replace"),
            "reload_match": source_bytes == reload_bytes,
        }

    metrics = {
        **metadata,
        "train": {
            "steps": args.steps,
            "efference_every": max(0, int(args.efference_every)),
                "free_run_aux_every": max(0, int(args.free_run_aux_every)),
                "free_run_aux_horizon": max(1, int(args.free_run_aux_horizon)),
                "free_run_aux_weight": float(args.free_run_aux_weight),
            "efference_steps": efference_steps,
            **({"field_bptt": int(model.field_bptt), "bptt_windows": bptt_windows}
               if int(getattr(model, "field_bptt", 0) or 0) else {}),
            **({"field_leak_learnable_rates": [float(x) for x in model.learnable_leak_rates().detach().reshape(-1)]}
               if getattr(model, "field_leak_logit", None) is not None else {}),
            "free_run_aux_steps": free_run_aux_steps,
            "free_run_aux_every": max(0, int(args.free_run_aux_every)),
            "free_run_aux_horizon": max(1, int(args.free_run_aux_horizon)),
            "free_run_aux_weight": float(args.free_run_aux_weight),
            "last_atom_scale": float(args.last_atom_scale),
                "last_atom_scale_adaptive": bool(args.last_atom_scale_adaptive),
                "last_atom_scale_min": float(args.last_atom_scale_min),
                "last_atom_scale_max": float(args.last_atom_scale_max),
                "field_hard_couple_scale": float(args.field_hard_couple_scale),
                "last_atom_dropout": float(args.last_atom_dropout),
                "field_leak": float(args.field_leak),
                "field_leak_spectrum": list(args.field_leak_spectrum) if args.field_leak_spectrum else None,
                "field_leak_spectrum_gain": (
                    args.field_leak_spectrum_gain if args.field_leak_spectrum else "none"
                ),
                "field_leak_spectrum_slow_frac": (
                    float(args.field_leak_spectrum_slow_frac) if args.field_leak_spectrum else 1.0
                ),
            "slow_mode_readout": bool(args.slow_mode_readout),
            "slow_readout_ahat": args.slow_readout_ahat,
            "slow_readout_lr": args.slow_readout_lr,
            "slow_memory_loss": float(args.slow_memory_loss),
            "slow_memory_lags": list(args.slow_memory_lags),
            "slow_memory_target": args.slow_memory_target,
            "slow_memory_loss_mean": (
                float(np.mean(slow_mem_values)) if slow_mem_values else None
            ),
            "slow_memory_loss_last_1000": (
                float(np.mean(slow_mem_values[-1000:])) if slow_mem_values else None
            ),
            "start_step": start_step,
            "total_steps": start_step + args.steps,
            "episode_length": args.episode_length,
            "episode_reset": bool(args.episode_reset),
            "atom_flush_every": args.atom_flush_every,
            "elapsed_seconds": elapsed,
            "transitions_per_second": args.steps / max(elapsed, 1e-9),
            "first_loss": all_losses[0],
            "final_loss": all_losses[-1],
            "final_total_perplexity": float(np.exp(min(all_losses[-1], 30.0))),
            "finite_losses": all(np.isfinite(all_losses)),
            "finite_parameters": finite_model(model),
            "nan_or_inf": not all(np.isfinite(all_losses)) or not finite_model(model),
            "field_rms_before_max": max(field_rms_before_values),
            "field_rms_after_max": max(field_rms_after_values),
            "field_scale_min": min(field_scales),
            "field_norm_max": max(field_norm_values),
            "field_limited_steps": sum(scale < 0.999999 for scale in field_scales),
            "field_limited_fraction": sum(scale < 0.999999 for scale in field_scales) / len(field_scales),
            "final_dynamics": dynamics_state,
            "final_atoms_before_validation": len(reloaded.core.atoms),
            "stream": stream_mode,
            "bytes_seen": (stream_source.bytes_seen if stream_source else data_bytes),
        },
        "validation": {
            "transitions": len(validation_records),
            "loss": validation_loss,
            "total_perplexity": float(np.exp(min(validation_loss, 30.0))) if np.isfinite(validation_loss) else None,
            "byte_loss": validation_byte_loss,
            "byte_perplexity": float(np.exp(min(validation_byte_loss, 30.0))) if np.isfinite(validation_byte_loss) else None,
            "source": "stream_ring" if stream_mode else "held_out_tail",
        },
        "checkpoint": str(checkpoint_path),
        "reload": reload_state,
        "generation": generation,
    }
    (output_dir / "loss_trajectory.json").write_text(json.dumps(trajectory, indent=2), encoding="utf-8")
    (output_dir / "run_metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    print(json.dumps(metrics, indent=2))


if __name__ == "__main__":
    main()