Blame for src/atom_batched.py

"""Batched byte-tick runtime for ``AtomNativeModel`` (B independent byte streams).

docs/BATCHED_TRAINING.md.  This module adds **no parameter** and does not touch
``AtomNativeModel``: it re-implements the per-tick path of the byte-tick recipe

    --max-span-bytes 1 --field-obligatory-hard --no-enable-merge --no-payload-copy
    --last-atom-readout --efference-every 0 --slow-every 1 [--field-bptt K]
    [--field-leak L] [--field-leak-learnable] [--last-atom-dropout P]
    [--field-loss-weight W] --atom-flush-every N

with a leading batch dimension, reading the weights of a wrapped single-stream
model.  Only the *runtime* state (field α, consolidation persistence, energy
history, last byte) is batched, so a checkpoint written with ``model.save`` is
the ordinary single-stream checkpoint (same parameter names and shapes) and is
loaded by ``run_atom_native.py`` / ``chat_atom_native.py`` unchanged.

Per stream b and tick t the computation is the one of ``AtomNativeModel._advance``
+ ``transition_loss`` for that recipe (same ops, same order):

* compile packet features -> (r, φ, ω, E, κ)                 (AtomCompiler.mapping)
* α ← leak(α_{t-1}) + projection(atom)            (graph kept under --field-bptt)
* α ← RK4(α, drive = r) ; phase-sync Laplacian ; RMS cap (FieldAmplitudeController)
* consolidation: stability of the living-atom energy history -> persistence EMA
* hard readout: obl·α-MLP(â) + 0.3·frozen(â) + last-atom 2-gram readout (dropout)
* loss = CE(length) + CE(next byte) + W·(1 − cos(field_probe(feats), features))

Paths of ``_advance`` that cannot change the loss or gradients of this recipe are
skipped: MERGE (off), aggregation/abstraction (only feed the zeroed+frozen legacy
decoder under hard), field interaction energy (diagnostic), α-bank snapshots
(only used by contrast/ignorance losses, which must be 0).  Everything else that
is not supported raises ``BatchedUnsupported``.
"""

from __future__ import annotations

import math

import torch
import torch.nn.functional as F

from .atom_native import AtomNativeModel


class BatchedUnsupported(ValueError):
    """A model/trainer option that the batched byte-tick runtime does not implement."""


def check_batched_supported(model: AtomNativeModel, *, efference_every: int = 0,
                            free_run_aux_every: int = 0) -> None:
    """Raise ``BatchedUnsupported`` listing every unsupported option that is on."""
    bad: list[str] = []
    s = model.surface
    if int(model.max_payload_bytes) != 1 or int(model.atomizer.max_span_bytes) != 1:
        bad.append("max_span_bytes/max_payload_bytes != 1 (spans > 1 byte)")
    if not bool(model.field_obligatory_hard):
        bad.append("field_obligatory_hard off (only the hard readout is batched)")
    if bool(model.enable_merge):
        bad.append("enable_merge (MERGE)")
    if bool(getattr(s, "payload_enabled", True)):
        bad.append("payload copy (use --no-payload-copy)")
    if int(efference_every or 0) > 0 or int(getattr(model, "_efference_every", 0) or 0) > 0:
        bad.append("efference (--efference-every must be 0)")
    if int(free_run_aux_every or 0) > 0:
        bad.append("free-run aux")
    if getattr(model, "field_leak_spectrum", None) is not None:
        bad.append("field_leak_spectrum")
    if getattr(model, "slow_readout", None) is not None:
        bad.append("slow_mode_readout")
    if getattr(model, "field_leak_slow_block", None) is not None:
        bad.append("field_leak_slow_block")
    for name in ("field_contrast_weight", "field_ignorance_weight", "printable_aux_weight",
                 "field_next_packet_weight", "slow_memory_weight"):
        if float(getattr(model, name, 0.0) or 0.0) != 0.0:
            bad.append(f"{name} != 0")
    if int(model.slow_every) != 1:
        bad.append("slow_every != 1")
    if bool(getattr(s, "last_atom_scale_adaptive", False)):
        bad.append("last_atom_scale_adaptive")
    if float(getattr(s, "field_hard_couple_scale", 0.0) or 0.0) > 0.0:
        bad.append("field_hard_couple_scale > 0")
    if bad:
        raise BatchedUnsupported("batched byte-tick mode does not support: " + "; ".join(bad))


class BatchedByteTick:
    """Runtime state + per-tick forward for B parallel byte streams.

    ``step(cur_bytes, features, target_bytes)`` advances every stream by one byte
    and returns per-stream losses (shape [B]); ``loss.mean()`` is the batch loss.
    With ``model.field_bptt > 0`` (and train mode + grad) α keeps its autograd
    graph across ticks until ``detach_field_graph()``.
    """

    def __init__(self, model: AtomNativeModel, batch_size: int, *, atom_flush_every: int,
                 efference_every: int = 0, free_run_aux_every: int = 0) -> None:
        check_batched_supported(model, efference_every=efference_every,
                                free_run_aux_every=free_run_aux_every)
        if not 1 <= int(batch_size):
            raise ValueError("batch_size must be >= 1")
        cap = int(atom_flush_every or 0)
        if cap < 1:
            raise BatchedUnsupported("batched mode needs --atom-flush-every N >= 1 (bounded atom memory)")
        self.model = model
        self.B = int(batch_size)
        self.cap = cap
        st = model.core.state
        self.M = int(st.n_modes)
        self.D = int(st.d_model)
        self.reset_state()

    # ------------------------------------------------------------------ state
    @property
    def device(self) -> torch.device:
        return self.model.core.state.alpha.device

    @property
    def dtype(self) -> torch.dtype:
        return self.model.core.state.alpha.dtype

    def reset_state(self, mask: torch.Tensor | None = None) -> None:
        """Episode reset (``AtomNativeModel.reset_state``) of all streams, or of
        the streams where ``mask`` (bool [B]) is True."""
        B, M, D, dev, dt = self.B, self.M, self.D, self.device, self.dtype
        if mask is None or not hasattr(self, "alpha"):
            self.alpha = torch.zeros(B, M, D, device=dev, dtype=dt)
            self.persistence = torch.zeros(B, D, device=dev, dtype=dt)
            self.persistence_strength = torch.zeros(B, device=dev, dtype=dt)
            self.e_buf = torch.zeros(B, self.cap, D, device=dev, dtype=dt)
            self.h_buf = torch.zeros(B, self.cap, D, device=dev, dtype=dt)
            self.last_byte = torch.zeros(B, dtype=torch.long, device=dev)
            self.n_atoms = [0] * B
            self.ticks = [0] * B
            return
        m = mask.to(device=dev, dtype=torch.bool)
        keep = (~m).to(dt)
        self.alpha = self.alpha * keep.view(B, 1, 1)
        self.persistence = self.persistence * keep.view(B, 1)
        self.persistence_strength = self.persistence_strength * keep
        self.e_buf = self.e_buf * keep.view(B, 1, 1)
        self.h_buf = self.h_buf * keep.view(B, 1, 1)
        self.last_byte = torch.where(m, torch.zeros_like(self.last_byte), self.last_byte)
        for b, flag in enumerate(m.tolist()):
            if flag:
                self.n_atoms[b] = 0
                self.ticks[b] = 0

    def flush_atoms(self, mask: torch.Tensor | None = None) -> None:
        """``--atom-flush-every``: clear the living atoms + energy history (field kept)."""
        if mask is None:
            self.n_atoms = [0] * self.B
            return
        for b, flag in enumerate(mask.tolist()):
            if flag:
                self.n_atoms[b] = 0

    def detach_field_graph(self) -> None:
        """BPTT window boundary: keep α's value, drop its autograd history."""
        self.alpha = self.alpha.detach()

    def state_dict(self) -> dict:
        return {
            "B": self.B, "cap": self.cap,
            "alpha": self.alpha.detach().cpu().clone(),
            "persistence": self.persistence.detach().cpu().clone(),
            "persistence_strength": self.persistence_strength.detach().cpu().clone(),
            # only the live prefix (after a flush n_atoms == 0: nothing to store)
            "e_buf": self.e_buf[:, :max(self.n_atoms)].detach().cpu().clone(),
            "h_buf": self.h_buf[:, :max(self.n_atoms)].detach().cpu().clone(),
            "last_byte": self.last_byte.detach().cpu().clone(),
            "n_atoms": list(self.n_atoms), "ticks": list(self.ticks),
        }

    def load_state_dict(self, sd: dict) -> None:
        if int(sd["B"]) != self.B or int(sd["cap"]) != self.cap:
            raise ValueError(f"batched state B/cap {sd['B']}/{sd['cap']} != {self.B}/{self.cap}")
        dev, dt = self.device, self.dtype
        self.alpha = sd["alpha"].to(device=dev, dtype=dt).clone()
        self.persistence = sd["persistence"].to(device=dev, dtype=dt).clone()
        self.persistence_strength = sd["persistence_strength"].to(device=dev, dtype=dt).clone()
        for name in ("e_buf", "h_buf"):
            buf = torch.zeros(self.B, self.cap, self.D, device=dev, dtype=dt)
            src = sd[name].to(device=dev, dtype=dt)
            buf[:, :src.shape[1]] = src
            setattr(self, name, buf)
        self.last_byte = sd["last_byte"].to(device=dev).clone()
        self.n_atoms = [int(x) for x in sd["n_atoms"]]
        self.ticks = [int(x) for x in sd["ticks"]]

    def sync_to_model(self, stream: int = 0) -> None:
        """Copy stream ``stream``'s field/persistence into the single-stream model
        runtime buffers (what a single-stream checkpoint stores as live state).
        The structural atom list is left empty (it is episode memory)."""
        m = self.model
        with torch.no_grad():
            m.core.state.alpha.copy_(self.alpha[stream].detach())
            m.core.state.t.fill_(float(self.ticks[stream]) * float(m.core.dynamics.dt))
            m.core.consolidation.persistence.copy_(self.persistence[stream].detach())
            m.core.consolidation.persistence_strength.copy_(self.persistence_strength[stream].detach())
        m.core.atoms = type(m.core.atoms)()
        m.core.energy_history = []
        m._bptt_alpha = None

    # ---------------------------------------------------------------- pieces
    def _const(self, ref: torch.Tensor) -> dict:
        """Per-(device, dtype) constant tensors (avoids re-creating them every tick)."""
        key = (ref.device, ref.dtype)
        cache = self.__dict__.setdefault("_const_cache", {})
        if key not in cache:
            n, D = self.M, self.D
            st, dyn = self.model.core.state, self.model.core.dynamics.dynamics
            cache[key] = {
                "mode_phase": st.mode_phase.to(device=ref.device, dtype=ref.dtype).view(1, n, 1),
                "mode_index": torch.arange(n, device=ref.device, dtype=ref.dtype).view(1, n, 1),
                "dyn_phase": dyn.mode_phase.to(device=ref.device, dtype=ref.dtype).view(1, n, 1),
                "dim_phase": torch.arange(D, device=ref.device, dtype=ref.dtype).view(1, 1, D),
                "zeros_long": torch.zeros(self.B, dtype=torch.long, device=ref.device),
            }
            cache[key]["sin_phase"] = torch.sin(cache[key]["dyn_phase"])
        return cache[key]

    def _to_dev(self, t: torch.Tensor) -> torch.Tensor:
        """Host->device copy without a host sync on CUDA (pinned, non-blocking)."""
        if self.device.type == "cuda":
            return t.pin_memory().to(self.device, non_blocking=True)
        return t.to(self.device)

    def _bptt_active(self) -> bool:
        m = self.model
        return bool(int(getattr(m, "field_bptt", 0) or 0) > 0 and m.training and torch.is_grad_enabled())

    def _compile(self, features: torch.Tensor):
        c = self.model.compiler
        raw = c.mapping(features).view(features.shape[0], 8, c.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])
        return r, phi, omega, energy, kappa

    def _projection(self, r, phi, omega, E, kappa) -> torch.Tensor:
        """FractalSuperpositionState._projection per stream (no sum over atoms)."""
        n = self.M
        cst = self._const(r)
        mode_phase = cst["mode_phase"]
        phase = phi.mean(dim=-1, keepdim=True).unsqueeze(1)
        freq = omega.mean(dim=-1, keepdim=True).unsqueeze(1).abs()
        freq_angle = 2.0 * torch.pi * freq / (1.0 + freq)
        centre = (freq_angle / (2.0 * torch.pi)) * n
        mode_index = cst["mode_index"]
        delta = torch.remainder(mode_index - centre + n / 2, n) - n / 2
        width = 1.0 + 0.25 * torch.sigmoid(freq)
        kernel = torch.exp(-(delta / width) ** 2)
        modulation = 0.5 + 0.5 * torch.cos(mode_phase - phase)
        strength = (E * kappa).mean(dim=-1, keepdim=True).unsqueeze(1)
        vectors = r.unsqueeze(1) * strength
        return vectors * kernel * modulation

    def _dyn(self, alpha: torch.Tensor, tok: torch.Tensor) -> torch.Tensor:
        """ToroidalDynamics.forward with a per-stream drive token (B, D)."""
        dyn = self.model.core.dynamics.dynamics
        left = torch.roll(alpha, shifts=1, dims=-2)
        right = torch.roll(alpha, shifts=-1, dims=-2)
        laplacian = left + right - 2.0 * alpha
        coupling = dyn.phase_sync * laplacian
        cst = self._const(alpha)
        phase = cst["dyn_phase"]
        rotated = cst["sin_phase"] * alpha
        dim_phase = cst["dim_phase"]
        t = tok.reshape(tok.shape[0], 1, self.D)
        drive = t * torch.cos(phase + 0.5 * t * dim_phase)
        drive = dyn.coupling_scale * drive
        decay = (dyn.energy_decay - 1.0) * alpha
        return coupling + rotated + drive + decay

    def _evolve(self, alpha: torch.Tensor, tok: torch.Tensor) -> torch.Tensor:
        dt = self.model.core.dynamics.dt
        k1 = self._dyn(alpha, tok)
        k2 = self._dyn(alpha + dt * 0.5 * k1, tok)
        k3 = self._dyn(alpha + dt * 0.5 * k2, tok)
        k4 = self._dyn(alpha + dt * k3, tok)
        return alpha + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)

    def _cap(self, field: torch.Tensor):
        ctl = self.model.field_controller
        rms_before = torch.sqrt(field.pow(2).mean(dim=(1, 2)) + ctl.eps)
        if ctl.max_rms is None:
            return field, rms_before, torch.ones_like(rms_before)
        cst = self._const(field)
        if "max_rms" not in cst or cst["max_rms_val"] != ctl.max_rms:
            cst["max_rms"] = torch.full((), float(ctl.max_rms), device=field.device, dtype=field.dtype)
            cst["max_rms_val"] = ctl.max_rms
        scale = torch.minimum(torch.ones_like(rms_before), cst["max_rms"] / (rms_before + ctl.eps))
        return field * scale.view(-1, 1, 1), rms_before, scale

    def _last_atom(self, last: torch.Tensor, prev: torch.Tensor, phi: torch.Tensor) -> torch.Tensor:
        la = self.model.surface.last_atom
        pair = torch.cat([la.byte_embed(last), la.byte_embed(prev)], dim=-1)
        h = F.relu(la.expand(pair))
        keep = la.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 = la.contract(sparse)
        phase = torch.cat([torch.cos(phi), torch.sin(phi)], dim=-1)
        out = F.gelu(z + la.phase_proj(phase))
        return la.head(out)

    def _mean_upto(self, buf: torch.Tensor, lo: list[int], hi: list[int], std: bool = False):
        """Per-stream mean (and population std) of buf[b, lo_b:hi_b] (hi_b > lo_b)."""
        if all(x == lo[0] for x in lo) and all(x == hi[0] for x in hi):
            sl = buf[:, lo[0]:hi[0]]
            mean = sl.mean(dim=1)
            return (mean, sl.std(dim=1, unbiased=False)) if std else mean
        idx = torch.arange(buf.shape[1], device=buf.device).view(1, -1)
        lo_t = torch.tensor(lo, device=buf.device).view(-1, 1)
        hi_t = torch.tensor(hi, device=buf.device).view(-1, 1)
        w = ((idx >= lo_t) & (idx < hi_t)).to(buf.dtype).unsqueeze(-1)
        cnt = (hi_t - lo_t).clamp_min(1).to(buf.dtype).view(-1, 1)
        mean = (buf * w).sum(dim=1) / cnt
        if not std:
            return mean
        var = (((buf - mean.unsqueeze(1)) ** 2) * w).sum(dim=1) / cnt
        return mean, var.sqrt()

    # ------------------------------------------------------------------ tick
    def step(self, cur_bytes: torch.Tensor, features: torch.Tensor, target_bytes: torch.Tensor,
             *, dropout_mask: torch.Tensor | None = None) -> dict:
        """One byte tick for every stream.  ``cur_bytes``/``target_bytes``: long [B];
        ``features``: [B, 320] packet features of the current byte.
        ``dropout_mask`` (bool [B], tests only) overrides the last-atom dropout draw."""
        m = self.model
        B, dev = self.B, self.device
        cur_bytes = cur_bytes.to(dev)
        target_bytes = target_bytes.to(dev)
        features = features.to(device=dev, dtype=self.dtype)
        self.n_atoms = [n + 1 for n in self.n_atoms]
        self.ticks = [t + 1 for t in self.ticks]
        r, phi, omega, energy, kappa = self._compile(features)

        # ---- inject (leak, projection); graph kept across ticks under BPTT
        bptt = self._bptt_active()
        prev = self.alpha if bptt else self.alpha.detach()
        learn_lam = m.learnable_leak_rates() if getattr(m, "field_leak_logit", None) is not None else None
        if learn_lam is not None:
            prev = prev * (1.0 - learn_lam) if bptt else (prev * (1.0 - learn_lam)).detach()
        else:
            leak = float(getattr(m, "field_leak", 0.0) or 0.0)
            if leak > 0.0:
                prev = prev * (1.0 - leak)
        if bptt:
            alpha = prev + self._projection(r, phi, omega, energy, kappa)
        else:
            with torch.no_grad():
                alpha = prev + self._projection(r, phi, omega, energy, kappa)

        # ---- living-atom energy buffer (memory atoms are detached)
        n_list = self.n_atoms
        if max(n_list) > self.cap:
            raise RuntimeError(f"living atoms exceed atom_flush_every={self.cap}; call flush_atoms()")
        lockstep = all(n == n_list[0] for n in n_list)
        if lockstep:
            pos = n_list[0] - 1
            ar = slice(None)
        else:
            pos = self._to_dev(torch.tensor([n - 1 for n in n_list]))
            ar = torch.arange(B, device=dev)
        self.e_buf[ar, pos] = energy.detach()
        e_mean = self._mean_upto(self.e_buf, [0] * B, n_list)

        # ---- field dynamics
        alpha_new = self._evolve(alpha, r)
        ps = m.core.dynamics.dynamics.phase_sync
        left = torch.roll(alpha_new, 1, 1)
        right = torch.roll(alpha_new, -1, 1)
        alpha_new = alpha_new + 0.1 * ps * (left + right - 2.0 * alpha_new)
        alpha_new, rms_before, scale = self._cap(alpha_new)

        # ---- consolidation (slow path, every tick): streams with > 5 history entries
        hist_n = [n - 1 for n in n_list]
        act = [h > 5 for h in hist_n]
        if any(act):
            lo = [max(0, h - 10) if a else 0 for h, a in zip(hist_n, act)]
            hi = [h if a else 1 for h, a in zip(hist_n, act)]
            with torch.no_grad():
                h_mean, h_std = self._mean_upto(self.h_buf, lo, hi, std=True)
                cv = h_std / h_mean.abs().clamp_min(1e-6)
                stability = torch.exp(-cv).clamp(0, 1) * e_mean.clamp(0, 1)
                score = e_mean.clamp(0, 1) * stability.clamp(0, 1)
                sel = score > m.core.consolidation.consolidation_threshold  # [B, D]
                cnt = sel.sum(dim=1)
                upd = cnt > 0
                if not all(act):
                    upd = upd & self._to_dev(torch.tensor(act))
                a_det = alpha_new.detach()
                incoming = (a_det * sel.unsqueeze(1).to(a_det.dtype)).sum(dim=(1, 2)) / (
                    cnt.clamp_min(1).to(a_det.dtype) * self.M)
                strength = sel.to(a_det.dtype).mean(dim=1)
                pd = m.core.consolidation.persistence_decay
                new_p = self.persistence * pd + (1 - pd) * incoming.view(B, 1)
                new_s = self.persistence_strength * pd + (1 - pd) * strength
                self.persistence = torch.where(upd.view(B, 1), new_p, self.persistence)
                self.persistence_strength = torch.where(upd, new_s, self.persistence_strength)
        self.h_buf[ar, pos] = e_mean

        # ---- hard readout
        s = m.surface
        flat = alpha_new.reshape(B, -1)
        a_only = flat @ s.alpha_jl.t()
        a_rms = a_only.pow(2).mean(dim=-1, keepdim=True).sqrt().clamp_min(1e-8)
        a_hat = a_only / a_rms
        frozen_byte = a_hat @ s.alpha_byte_frozen.t()
        frozen_len = a_hat @ s.alpha_length_frozen.t()
        train_byte = s.alpha_byte_proj(a_hat)
        train_len = s.alpha_length_proj(a_hat)
        obl = s.obligatory_scale()
        byte_logits = obl * train_byte + 0.3 * frozen_byte
        la_on = getattr(s, "last_atom", None) is not None and not bool(getattr(s, "last_atom_force_off", False))
        drop = None
        if la_on:
            p_drop = float(getattr(s, "last_atom_dropout", 0.0) or 0.0)
            if dropout_mask is not None:
                drop = dropout_mask.detach().to("cpu", dtype=torch.bool).reshape(B)
            elif p_drop > 0.0 and m.training and torch.is_grad_enabled():
                # same draw as the single-stream path (one uniform per stream, CPU RNG)
                drop = torch.ones(B, dtype=torch.bool) if p_drop >= 1.0 else torch.rand(B) < p_drop
            if drop is not None and bool(drop.all()):
                la_on = False  # like zeros_like(la_byte): last-atom params get no grad at all
        if la_on:
            if lockstep:
                prev_b = self.last_byte if n_list[0] >= 2 else cur_bytes
            else:
                prev_b = torch.where(self._to_dev(torch.tensor([n >= 2 for n in n_list])), self.last_byte, cur_bytes)
            la_raw = self._last_atom(cur_bytes, prev_b, phi)
            la_scale = float(getattr(s, "last_atom_scale", 1.0))
            la_byte = la_raw if la_scale == 1.0 else la_raw * la_scale  # x*1.0 == x exactly
            if drop is not None and bool(drop.any()):
                la_byte = la_byte * self._to_dev((~drop).to(la_byte.dtype)).unsqueeze(-1)
            byte_logits = byte_logits + la_byte
        length_logits = obl * train_len + 0.3 * frozen_len

        # ---- losses (per stream)
        length_loss = F.cross_entropy(length_logits, self._const(length_logits)["zeros_long"], reduction="none")
        byte_loss = F.cross_entropy(byte_logits, target_bytes, reduction="none")
        loss = length_loss + byte_loss
        w = float(m.field_loss_weight)
        persist_loss = torch.zeros_like(byte_loss)
        if w > 0:
            mean = alpha_new.mean(dim=1)
            std = alpha_new.std(dim=1, unbiased=False)
            feats = s.field_feat_norm(torch.cat([mean, std, self.persistence, r], dim=-1))
            pred = m.field_probe(feats)
            persist_loss = 1.0 - F.cosine_similarity(pred, features, dim=-1)
            loss = loss + w * persist_loss

        self.alpha = alpha_new if bptt else alpha_new.detach()
        self.last_byte = cur_bytes
        return {
            "loss": loss, "byte_loss": byte_loss, "length_loss": length_loss,
            "persist_loss": persist_loss, "byte_logits": byte_logits,
            "field": alpha_new, "field_rms_before": rms_before, "field_scale": scale,
        }