Blame for python/spira/model.py

"""SPIRA — Spiking Plastic Recurrent Architecture.

Full system pipeline:

    raw bytes ──EPT──▶ content-addressed patches ──hashing trick──▶ input spikes
                                   │                                     │
                                   │                              LIF reservoir
                                   │                              (STDP, Dale)
                                   │                                     │
                                   │                              three-factor
                                   │                              readout (R-STDP)
                                   ▼                                     ▼
                        hippocampal associative memory ──fused──▶ prediction

    generation runs the same loop closed: predicted patch is fed back as
    input spikes, one patch (multi-byte) at a time.
"""
from __future__ import annotations

import json
import time
from dataclasses import dataclass, field, asdict

import numpy as np

from .ept import EPT
from .hashing import fnv1a64, splitmix64
from .hippocampus import Hippocampus
from .spiking import SpikingCore

UNK_COL = -1  # sentinel before adding the +1 offset


@dataclass
class SPIRAConfig:
    # EPT
    target_avg_len: float = 4.5
    min_len: int = 2
    max_len: int = 32
    alpha: float = 0.5
    # spiking core
    d_in: int = 1024
    k_in: int = 12
    n_res: int = 512
    vocab_size: int = 16384
    seed: int = 20261022
    # hierarchical multi-scale cortex (v2)
    hierarchy: bool = True
    n_g: int = 256
    win: int = 8
    eta_pos: float = 0.01
    eta_r: float = 0.01
    eta_stdp: float = 0.0015
    # hippocampus
    am_size: int = 1 << 23
    am_m: int = 4
    am_min_lanes: int = 3
    # training
    snn_tokens: int = 600_000
    am_tokens: int = 1_200_000
    epochs: int = 2
    # hippocampal replay -> cortical consolidation
    replay_tokens: int = 180_000
    replay_slice: int = 512
    # v3 cortex fixes (Muse)
    direct_input: bool = True
    trace_taus: tuple = (0.4, 0.8, 0.95)
    direct_tau: float = 0.5
    structured_inputs: bool = True
    k_tri: int = 6
    # generation / fusion
    beta_am: float = 0.5
    am_override: float = 6.0
    temperature: float = 0.8
    top_k: int = 8


class PatchRegistry:
    """Compact content-addressed patch storage.

    Sorted uint64 id array + offsets into a single concatenated byte buffer.
    Lookup is O(log n) via searchsorted (vectorizable). Also stores counts.
    """

    def __init__(self) -> None:
        self.ids = np.zeros(0, dtype=np.uint64)      # sorted unique
        self.counts = np.zeros(0, dtype=np.uint32)
        self.offsets = np.zeros(0, dtype=np.uint64)  # len = n+1
        self.buffer = b""

    @classmethod
    def build(cls, ids: np.ndarray, patches: list[bytes], keep: set[int] | None = None) -> "PatchRegistry":
        reg = cls()
        id2bytes: dict[int, bytes] = {}
        for i, b in zip(ids.tolist(), patches):
            if keep is not None and i not in keep:
                continue
            id2bytes[i] = b          # last occurrence wins; identical anyway
        if not id2bytes:
            return reg
        reg.ids = np.fromiter(id2bytes.keys(), dtype=np.uint64, count=len(id2bytes))
        reg.ids.sort()
        buf = bytearray()
        offsets = np.zeros(len(reg.ids) + 1, dtype=np.uint64)
        for j, pid in enumerate(reg.ids.tolist()):
            b = id2bytes[pid]
            offsets[j] = len(buf)
            buf.extend(b)
        offsets[-1] = len(buf)
        reg.offsets = offsets
        reg.buffer = bytes(buf)
        # counts from the stream
        uniq, cnt = np.unique(ids, return_counts=True)
        pos = np.searchsorted(reg.ids, uniq)
        pos_ok = (pos < len(reg.ids)) & (reg.ids[np.minimum(pos, len(reg.ids) - 1)] == uniq)
        reg.counts = np.zeros(len(reg.ids), dtype=np.uint32)
        reg.counts[pos[pos_ok]] = cnt[pos_ok].astype(np.uint32)
        return reg

    def __len__(self) -> int:
        return len(self.ids)

    def has(self, ids: np.ndarray | int) -> np.ndarray | bool:
        if len(self.ids) == 0:
            return np.zeros(len(ids), dtype=bool) if isinstance(ids, np.ndarray) else False
        if isinstance(ids, (int, np.integer)):
            pos = np.searchsorted(self.ids, np.uint64(ids))
            return bool(pos < len(self.ids) and self.ids[min(pos, len(self.ids) - 1)] == ids)
        pos = np.searchsorted(self.ids, ids)
        pos_c = np.minimum(pos, len(self.ids) - 1)
        return (pos < len(self.ids)) & (self.ids[pos_c] == ids)

    def bytes_of(self, pid: int) -> bytes | None:
        pos = int(np.searchsorted(self.ids, np.uint64(pid)))
        if pos >= len(self.ids) or self.ids[pos] != np.uint64(pid):
            return None
        a, b = int(self.offsets[pos]), int(self.offsets[pos + 1])
        return self.buffer[a:b]

    def index_of(self, pid: int) -> int:
        """Registry index of a patch id, or -1."""
        pos = int(np.searchsorted(self.ids, np.uint64(pid)))
        if pos >= len(self.ids) or self.ids[pos] != np.uint64(pid):
            return -1
        return pos

    def extend(self, id2bytes: dict[int, bytes]) -> int:
        """Append new patches (one-shot episodic learning).

        Returns the number of patches added. Keeps arrays sorted.
        """
        if not id2bytes:
            return 0
        cand = np.array(sorted(id2bytes.keys()), dtype=np.uint64)
        cand = cand[~self.has(cand)]
        if len(cand) == 0:
            return 0
        n_old = len(self.ids)
        merged = np.concatenate([self.ids, cand])
        order = np.argsort(merged, kind="stable")
        merged = merged[order]
        old_lens = np.diff(self.offsets).astype(np.uint64)
        new_lens = np.array([len(id2bytes[int(i)]) for i in cand], dtype=np.uint64)
        old_off = self.offsets[:-1]
        out = bytearray()
        offsets = np.zeros(len(merged) + 1, dtype=np.uint64)
        for k, idx in enumerate(order.tolist()):
            if idx < n_old:
                a, b = int(old_off[idx]), int(old_off[idx] + old_lens[idx])
                out.extend(self.buffer[a:b])
            else:
                out.extend(id2bytes[int(cand[idx - n_old])])
            offsets[k] = len(out)
        offsets[-1] = len(out)
        self.ids = merged
        self.offsets = offsets
        self.buffer = bytes(out)
        self.counts = np.concatenate([self.counts, np.zeros(len(cand), dtype=np.uint32)])[order]
        return len(cand)


class SPIRA:
    EPI_CAP_TOKENS = 2_000_000   # hippocampal episodic-trace ring cap

    def __init__(self, cfg: SPIRAConfig) -> None:
        self.cfg = cfg
        self.ept = EPT(budget_bits=16.0, min_len=cfg.min_len, max_len=cfg.max_len, alpha=cfg.alpha)
        self.core = SpikingCore(
            d_in=cfg.d_in, k_in=cfg.k_in, n_res=cfg.n_res,
            n_out=cfg.vocab_size + 1, seed=cfg.seed,
            hier=cfg.hierarchy, n_g=cfg.n_g, win=cfg.win,
            trace_taus=tuple(cfg.trace_taus), direct_input=cfg.direct_input,
            direct_tau=cfg.direct_tau,
        )
        self._row_cache: dict[int, np.ndarray] = {}
        self.am = Hippocampus(m=cfg.am_m, size=cfg.am_size, min_lanes=cfg.am_min_lanes)
        self.registry = PatchRegistry()
        self.vocab_ids: np.ndarray = np.zeros(0, dtype=np.uint64)   # col -> id
        self.n_out = cfg.vocab_size + 1
        self.meta: dict = {}
        # inference-time fusion calibration (variance matching): the raw slow
        # readout operates at a different score scale than the fast one, so the
        # top-down prior is rescaled by sigma_fast/sigma_slow — measured once
        # on a reference stream (calibrate_fusion), never trained.
        self.fusion: dict = {"scale": 1.0, "sigma_fast": None, "sigma_slow": None}
        # episodic traces retained for hippocampal replay (CA3-style sequence
        # store; the AM itself only keeps hashed associations)
        self.episodes: list[np.ndarray] = []

    # ================================================================ build
    @classmethod
    def build(cls, train_bytes: bytes, cfg: SPIRAConfig, log=print) -> "SPIRA":
        m = cls(cfg)
        fit_slice = train_bytes[: 2_000_000]
        calib_slice = train_bytes[2_000_000: 2_400_000]
        log("[spira] fitting EPT byte-level Markov model ...")
        m.ept.fit(fit_slice)
        if calib_slice:
            budget = m.ept.calibrate(calib_slice, target_len=cfg.target_avg_len)
            log(f"[spira] calibrated surprisal budget = {budget:.2f} bits "
                f"(target avg patch len {cfg.target_avg_len})")
        log("[spira] segmenting full train stream ...")
        t0 = time.perf_counter()
        patches = m.ept.segment(train_bytes)
        ids = np.fromiter((fnv1a64(p) for p in patches), dtype=np.uint64, count=len(patches))
        seg_s = time.perf_counter() - t0
        st = m.ept.stats(patches)
        log(f"[spira] {st['n_patches']} patches, avg len {st['avg_len']:.2f} B "
            f"({st['bytes']/1e6:.1f} MB segmented in {seg_s:.1f}s, "
            f"{st['bytes']/1e6/max(seg_s,1e-9):.2f} MB/s)")
        uniq, cnt = np.unique(ids, return_counts=True)
        log(f"[spira] distinct patches in train: {len(uniq)}")
        # vocabulary: top-V by count (deterministic tie-break: larger id first)
        order = np.lexsort((uniq, -cnt))[: cfg.vocab_size]
        m.vocab_ids = np.sort(uniq[order]).astype(np.uint64)
        # registry: vocab + anything with count >= 2 (repeatable content)
        keep = set(m.vocab_ids.tolist()) | set(uniq[cnt >= 2].tolist())
        log(f"[spira] building patch registry ({len(keep)} ids) ...")
        m.registry = PatchRegistry.build(ids, patches, keep=keep)
        m.meta = {
            "train_bytes": len(train_bytes),
            "n_patches": st["n_patches"],
            "avg_patch_len": st["avg_len"],
            "distinct_patches": int(len(uniq)),
            "registry_size": len(m.registry),
            "budget_bits": m.ept.budget_bits,
            "segmentation_MB_per_s": st["bytes"] / 1e6 / max(seg_s, 1e-9),
        }
        m._ids_train = ids
        m._patches_train = None   # registry carries everything; free ~100MB
        return m

    # =========================================================== array prep
    @staticmethod
    def input_neurons_stream(ids: np.ndarray, d_in: int, k_in: int) -> np.ndarray:
        """Vectorized input-neuron projection: (T, k_in) int32."""
        T = len(ids)
        out = np.empty((T, k_in), dtype=np.int32)
        u64 = ids.astype(np.uint64)
        for j in range(k_in):
            salt = (0x51ED270B733D4C31 + j * 0x9E3779B97F4A7C15) & 0xFFFFFFFFFFFFFFFF
            z = u64 ^ np.uint64(salt)
            z = z + np.uint64(0x9E3779B97F4A7C15)
            z = (z ^ (z >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9)
            z = (z ^ (z >> np.uint64(27))) * np.uint64(0x94D049BB133111EB)
            z = z ^ (z >> np.uint64(31))
            out[:, j] = ((z >> np.uint64(32)) % np.uint64(d_in)).astype(np.int32)
        return out

    def cols_of(self, ids: np.ndarray) -> np.ndarray:
        """Vocabulary column per patch id (UNK -> last column)."""
        V = self.cfg.vocab_size
        pos = np.searchsorted(self.vocab_ids, ids)
        pos_c = np.minimum(pos, V - 1)
        ok = (pos < V) & (self.vocab_ids[pos_c] == ids)
        cols = np.where(ok, pos_c, V).astype(np.int64)  # V = UNK column
        return cols

    # ------------------------------------------------------- v3 input rows
    def _structured_row(self, pid: int, pbytes: bytes | None) -> np.ndarray:
        """Input-neuron row for one patch: (k_in - k_tri) whole-patch hash
        neurons + k_tri byte-trigram hash neurons. Patches sharing trigrams
        share input neurons, so similar contexts produce similar reservoir
        states (v1/v2 rows were pure whole-patch hashes: fully orthogonal).
        Falls back to the v2 row when the patch bytes are unknown."""
        from .hashing import fnv1a64, input_neurons, splitmix64
        cfg = self.cfg
        if not cfg.structured_inputs or pbytes is None:
            return np.asarray(input_neurons(pid, cfg.d_in, cfg.k_in), dtype=np.int32)
        k_whole = cfg.k_in - cfg.k_tri
        row = list(input_neurons(pid, cfg.d_in, k_whole))
        grams: set[bytes] = set()
        if len(pbytes) >= 3:
            grams = {pbytes[i:i + 3] for i in range(len(pbytes) - 2)}
        elif len(pbytes) == 2:
            grams = {pbytes}
        cand = sorted({splitmix64(fnv1a64(g)) % cfg.d_in for g in grams})
        tri = cand[: cfg.k_tri]
        while len(tri) < cfg.k_tri:  # short patches: fill from whole-hash salts
            tri.append(int(input_neurons(pid, cfg.d_in, cfg.k_tri + len(tri))[-1]))
        return np.asarray(row + tri, dtype=np.int32)

    def neuron_rows(self, ids: np.ndarray, patches: list[bytes] | None = None) -> np.ndarray:
        """(T, k_in) input rows; structured when enabled, v2-identical else."""
        cfg = self.cfg
        if not cfg.structured_inputs:
            return self.input_neurons_stream(ids, cfg.d_in, cfg.k_in)
        by_id: dict[int, bytes] = {}
        if patches is not None:
            for pid, pb in zip(ids.tolist(), patches):
                by_id.setdefault(pid, pb)
        rows = np.empty((len(ids), cfg.k_in), dtype=np.int32)
        cache = self._row_cache
        for t, pid in enumerate(ids.tolist()):
            r = cache.get(pid)
            if r is None:
                pb = by_id.get(pid)
                if pb is None:
                    pb = self.registry.bytes_of(pid)
                r = self._structured_row(pid, pb)
                cache[pid] = r
            rows[t] = r
        return rows

    def prepare(self) -> None:
        """Precompute training arrays for the SNN and AM subsets only."""
        cfg = self.cfg
        ids = self._ids_train
        n_snn = min(cfg.snn_tokens, len(ids) - 1)
        n_am = min(cfg.am_tokens, len(ids) - 1)
        log = print
        log(f"[spira] precomputing projections for {n_snn} SNN + {n_am} AM tokens ...")
        t0 = time.perf_counter()
        self._neurons = self.neuron_rows(ids[:n_snn])
        self._cols = self.cols_of(ids[: n_snn + 1])
        # AM: val = registry index of the NEXT token (skip unregistered)
        am_ids = ids[:n_am]
        reg_idx = np.full(n_am, -1, dtype=np.int64)
        pos = np.searchsorted(self.registry.ids, am_ids)
        pos_c = np.minimum(pos, len(self.registry.ids) - 1)
        ok = (pos < len(self.registry.ids)) & (self.registry.ids[pos_c] == am_ids)
        reg_idx[ok] = pos_c[ok]
        self._reg_idx = reg_idx
        ls = self.am.lane_keys_stream(am_ids)
        # AM contexts (diagonal keys): ctx[r, l] = lane key of token r+l at
        # slot l; the context ending at t = r+m-1 continues with token r+m.
        m = self.am.m
        n_ctx = len(ls) - m
        ctx = np.empty((n_ctx, m), dtype=np.uint32)
        for lane in range(m):
            ctx[:, lane] = ls[lane : lane + n_ctx, lane]
        self._am_ctx = ctx
        log(f"    done in {time.perf_counter()-t0:.1f}s")

    # ================================================================ train
    def train(self, log_fn=print, ckpt_path: str | None = None) -> dict:
        cfg = self.cfg
        n_snn = len(self._neurons)
        n_ctx_am = len(self._am_ctx)
        history = {"epochs": []}
        for ep in range(cfg.epochs):
            log_fn(f"[spira] === SNN epoch {ep+1}/{cfg.epochs} ({n_snn} ticks) ===")
            res = self.core.train_pass(
                self._neurons[:n_snn], self._cols[1 : n_snn + 1],
                log_every=10_000,
                log_fn=lambda d, ep=ep: log_fn(
                    f"    tick {d['tick']}/{d['total']} acc={d['acc_ema']:.3f} "
                    f"dop={d['dopamine_ema']:+.3f} {d['ticks_per_sec']:.0f} t/s"),
                record_raster_from=(n_snn - 1500 if ep == cfg.epochs - 1 else None),
            )
            raster = res.pop("raster", None)
            if raster and ckpt_path:
                ticks = np.concatenate([np.full(len(sp), t, dtype=np.int32) for t, sp in raster])
                neurons_r = np.concatenate([sp for _, sp in raster]).astype(np.int32)
                np.savez_compressed(ckpt_path + ".raster.npz", ticks=ticks, neurons=neurons_r)
            history["epochs"].append({k: v for k, v in res.items()})
            if ckpt_path:
                self.save(ckpt_path)
        # hippocampal one-shot pass over the (larger) AM stream
        log_fn(f"[spira] === hippocampus one-shot write pass ({n_ctx_am} contexts) ===")
        t0 = time.perf_counter()
        m = self.am.m
        ctx = self._am_ctx[:n_ctx_am]
        vals = self._ids_train[m : m + n_ctx_am].astype(np.uint64)
        self.am.write_stream(ctx, vals)
        log_fn(f"    wrote {n_ctx_am} associations in "
               f"{time.perf_counter()-t0:.1f}s; {self.am.stats()}")
        history["am"] = self.am.stats()
        if ckpt_path:
            self.save(ckpt_path)
        return history

    # =============================================================== predict
    def _scores(self, x: np.ndarray) -> np.ndarray:
        """Behavioural scores: fast readout + variance-matched top-down slow
        prior (if the hierarchical cortex is active). No hippocampal vote."""
        u = (x - self.core.x_mean) @ self.core.W_out
        if self.core.hier:
            u = u + np.float32(self.core.beta * self.fusion["scale"]) * self.core.u_slow
        return u

    def calibrate_fusion(self, data: bytes, n_tokens: int = 3000) -> dict:
        """Measure fast/slow readout score scales on a reference stream and set
        the fusion normalization sigma_fast/sigma_slow.

        The slow readout's raw scores live at a different (much larger) scale
        than the fast one; fusing them raw lets even a tiny beta swamp the fast
        argmax. Variance matching bounds the top-down prior's contribution to
        beta * sigma_fast — the same units the fast readout itself uses. Pure
        measurement, no plasticity (same spirit as the EPT budget calibration).
        """
        patches = self.ept.segment(data)
        ids = np.fromiter((fnv1a64(p) for p in patches), dtype=np.uint64,
                          count=len(patches))[: n_tokens + 1]
        neurons = self.neuron_rows(ids[:-1], patches[:-1])
        snap = self.core.snapshot()
        w_rec_snap = self.core.W_rec.copy()
        self.core.reset_dynamics()
        self._win_count = 0
        sig_f: list[float] = []
        sig_s: list[float] = []
        for t in range(len(ids) - 1):
            x = self.core.step(neurons[t])
            if self.core.hier:
                self._win_count += 1
                if self._win_count % self.core.win == 0:
                    self.core.step_G()
                    if len(sig_s) < 400 or len(sig_s) % 8 == 0:
                        sig_s.append(float(np.std(self.core.u_slow)))
            u_fast = (x - self.core.x_mean) @ self.core.W_out
            if len(sig_f) < 400 or len(sig_f) % 8 == 0:
                sig_f.append(float(np.std(u_fast)))
        self.core.restore(snap)
        self.core.W_rec = w_rec_snap
        sf = float(np.mean(sig_f)) if sig_f else 1.0
        ss = float(np.mean(sig_s)) if sig_s else 1.0
        self.fusion = {"scale": sf / max(ss, 1e-9), "sigma_fast": sf, "sigma_slow": ss}
        return dict(self.fusion)

    def _feed(self, pid: int) -> np.ndarray:
        """Feed one token: fast step (+ slow step at window boundaries)."""
        x = self.core.step(self._neuron_row(pid))
        if self.core.hier:
            self._win_count += 1
            if self._win_count % self.core.win == 0:
                self.core.step_G()
        return x

    def _col_of_id(self, pid: int) -> int | None:
        """Vocab column of a patch id, or None if out of vocabulary."""
        V = self.cfg.vocab_size
        col = int(np.searchsorted(self.vocab_ids, np.uint64(pid)))
        if col < V and self.vocab_ids[col] == np.uint64(pid):
            return col
        return None

    def fuse_scores(self, x: np.ndarray, last_ids: list[int]):
        """Combine behavioural scores (fast + top-down slow prior) with a
        hippocampal vote.

        Returns (fused_scores, am_val, am_weight, bypass_id).
        """
        u = self._scores(x)
        am_val, am_w = None, 0.0
        bypass_id = None
        if len(last_ids) >= self.am.m:
            lanes = self.am.lane_keys(last_ids[-self.am.m :])
            am_val, am_w = self.am.query(lanes)
        if am_val is None:
            return u, am_val, am_w, None
        pid = int(am_val)
        col = self._col_of_id(pid)
        if col is not None:
            fused = u.copy()
            fused[col] += self.cfg.beta_am * am_w
            return fused, am_val, am_w, None
        if am_w >= self.cfg.am_override and self.registry.bytes_of(pid) is not None:
            bypass_id = pid   # strong episodic recall outside vocabulary
        return u, am_val, am_w, bypass_id

    def generate(self, prompt: bytes, n_tokens: int = 120, temperature: float | None = None,
                 top_k: int | None = None, use_am: bool = True, seed: int | None = None,
                 on_token=None) -> bytes:
        """Closed-loop generation: predicted patches are fed back as spikes."""
        cfg = self.cfg
        temperature = cfg.temperature if temperature is None else temperature
        top_k = cfg.top_k if top_k is None else top_k
        rng = np.random.default_rng(seed if seed is not None else int(time.time_ns() & 0x7FFFFFFF))
        self.core.reset_dynamics()
        self._win_count = 0
        patches = self.ept.segment(prompt)
        if not patches:
            return prompt
        ids = [fnv1a64(p) for p in patches]
        out = bytearray(prompt)
        V = cfg.vocab_size
        UNK = V  # UNK column index
        # warm-up: feed the entire prompt (state now reflects the last patch)
        x = np.zeros(self.core.n_res, dtype=np.float32)
        for pid in ids:
            x = self._feed(pid)
        for step in range(n_tokens):
            if use_am:
                fused, am_val, am_w, bypass_id = self.fuse_scores(x, ids)
            else:
                fused, bypass_id = self._scores(x), None
            if bypass_id is not None:
                chosen_id = bypass_id
            else:
                scores = np.asarray(fused, dtype=np.float64).copy()
                scores[UNK] = -np.inf          # cannot emit UNK bytes
                k = max(1, min(top_k, V))
                top = np.argpartition(scores, -k)[-k:]
                z = (scores[top] - scores[top].max()) / max(temperature, 1e-6)
                w = np.exp(z)
                w /= w.sum()
                pick = int(top[rng.choice(len(top), p=w)])
                chosen_id = int(self.vocab_ids[pick])
            b = self.registry.bytes_of(chosen_id)
            if b is None:
                break  # unknown emission: stop rather than corrupt
            out.extend(b)
            ids.append(chosen_id)
            x = self._feed(chosen_id)
            if on_token is not None:
                on_token(step, b)
        return bytes(out)

    def _neuron_row(self, pid: int) -> np.ndarray:
        r = self._row_cache.get(pid)
        if r is None:
            r = self._structured_row(pid, self.registry.bytes_of(pid))
            self._row_cache[pid] = r
        return r

    # ============================================================= one-shot
    def teach_bytes(self, data: bytes) -> dict:
        """ONE-SHOT episodic teaching: segment, register unseen patches, and
        write every context -> continuation into the hippocampus — a single
        pass, no gradients, no epochs, no repetition."""
        patches = self.ept.segment(data)
        ids = np.fromiter((fnv1a64(p) for p in patches), dtype=np.uint64, count=len(patches))
        return self.teach_ids(ids, patches)

    def teach_ids(self, ids: np.ndarray, patches: list[bytes] | None = None,
                  episodic: bool = True) -> dict:
        n = len(ids)
        # 0) retain the episodic trace for hippocampal replay (the sequence
        #    the cortex will re-run during "sleep") — only for genuinely
        #    episodic teaching, not for the corpus baseline fill
        if episodic:
            self._remember_episode(ids)
        # 1) register unseen patches so continuations are always storable
        known = self.registry.has(ids)
        if patches is not None:
            new_map: dict[int, bytes] = {}
            for i, b in zip(ids[~known].tolist(), (p for p, k in zip(patches, known) if not k)):
                if i not in new_map:
                    new_map[i] = b
            n_new = self.registry.extend(new_map)
        else:
            n_new = 0
        # 2) write associations (vals are stable patch ids)
        lane = self.am.lane_keys_stream(ids)
        m = self.am.m
        n_ctx = n - m
        ctx = np.empty((n_ctx, m), dtype=np.uint32)
        for li in range(m):
            ctx[:, li] = lane[li : li + n_ctx, li]
        vals = ids[m : m + n_ctx].astype(np.uint64)
        self.am.write_stream(ctx, vals)
        return {
            "tokens": int(n),
            "new_patches": int(n_new),
            "associations": int(n_ctx),
            "episode_tokens": int(n),
        }

    # ================================================================ replay
    def _remember_episode(self, ids: np.ndarray) -> None:
        """Append to the hippocampal episodic trace (ring-capped)."""
        if len(ids) < 2:
            return
        self.episodes.append(np.asarray(ids, dtype=np.uint64))
        total = sum(len(e) for e in self.episodes)
        while total > self.EPI_CAP_TOKENS and len(self.episodes) > 1:
            total -= len(self.episodes[0])
            self.episodes.pop(0)

    def replay(self, n_tokens: int | None = None, slice_len: int | None = None,
               seed: int = 0, log=print) -> dict:
        """Hippocampal replay -> cortical consolidation (closing the loop).

        Random slices of the episodic traces retained by the hippocampus are
        re-run through the cortex in full training mode: the same local rules
        (three-factor readout delta, modulated recurrent STDP, dopaminergic
        RPE, slow-area learning) now act on material the cortex never saw
        externally. One-shot memories are gradually consolidated into slow
        cortical weights — the complementary-systems loop, closed.

        Biologically this is Sharpewave-ripple-like offline replay: the cortex
        does not re-run isolated (context -> continuation) pairs but short
        SEQUENCES, so reservoir state statistics match the awake stream.
        """
        cfg = self.cfg
        n_tokens = cfg.replay_tokens if n_tokens is None else n_tokens
        slice_len = cfg.replay_slice if slice_len is None else slice_len
        if not self.episodes:
            return {"replayed_tokens": 0, "note": "no episodic traces to replay"}
        rng = np.random.default_rng(seed)
        weights = np.array([len(e) for e in self.episodes], dtype=np.float64)
        weights /= weights.sum()
        chunks: list[np.ndarray] = []
        got = 0
        while got < n_tokens:
            e = self.episodes[int(rng.choice(len(self.episodes), p=weights))]
            start = int(rng.integers(0, max(1, len(e) - slice_len)))
            take = min(slice_len, len(e) - start)
            chunks.append(e[start : start + take])
            got += take
        ids = np.concatenate(chunks)
        neurons = self.neuron_rows(ids[:-1])
        cols = self.cols_of(ids)[1:]
        log(f"[replay] re-running {len(ids)} cortical tokens sampled from "
            f"{len(self.episodes)} episodic trace(s) (slice {slice_len})")
        res = self.core.train_pass(
            neurons, cols, log_every=25_000,
            log_fn=lambda d: log(f"    tick {d['tick']}/{d['total']} acc={d['acc_ema']:.3f} "
                                 f"beta={d.get('beta')} {d['ticks_per_sec']:.0f} t/s"),
        )
        res.pop("raster", None)
        return {"replayed_tokens": int(len(ids)), **res}

    # ================================================================ eval
    def eval_stream(self, data: bytes, use_am: bool = True, teach_am_first: bool = False,
                    corrupt_lanes: int = 0, seed: int = 0, max_positions: int | None = None) -> dict:
        """Evaluate next-patch prediction on a byte stream.

        Metrics for the SNN readout alone, the hippocampus alone, and the
        fused system — both over all positions and over in-vocab targets only
        (targets outside the readout vocabulary cannot be predicted by the
        readout, but the hippocampus can still retrieve them).
        """
        from .hashing import lane_key32
        patches = self.ept.segment(data)
        ids = np.fromiter((fnv1a64(p) for p in patches), dtype=np.uint64, count=len(patches))
        n = len(ids)
        cols = self.cols_of(ids)
        V = self.cfg.vocab_size
        if teach_am_first:
            # ONE-SHOT teaching pass (registers new patches too)
            taught = self.teach_ids(ids, patches)
        # SNN inference pass in the TRAINING regime: the reservoir keeps its
        # modulated-STDP plasticity (the operating point is part of the
        # learned system) while the readout stays frozen for measurement.
        snap = self.core.snapshot()
        w_rec_snap = self.core.W_rec.copy()
        neurons = self.neuron_rows(ids, patches)
        lane = self.am.lane_keys_stream(ids)
        # diagonal context-key stream: ctxq[r] = key of context tokens r..r+m-1
        m_ctx = self.am.m
        n_ctxq = n - m_ctx + 1
        ctxq = np.empty((n_ctxq, m_ctx), dtype=np.uint32)
        for li in range(m_ctx):
            ctxq[:, li] = lane[li : li + n_ctxq, li]
        rng = np.random.default_rng(seed)
        train_ids = self._ids_train if hasattr(self, "_ids_train") else ids
        tot = inv = 0
        snn_ok = snn_ok_inv = am_ok = am_ok_inv = fused_ok = fused_ok_inv = 0
        hier_ok = hier_ok_inv = 0
        stop = n - 1 if max_positions is None else min(n - 1, max_positions)
        x = np.zeros(self.core.n_res, dtype=np.float32)
        for t in range(stop):
            x = self.core.step(neurons[t])
            target_col = int(cols[t + 1])
            target_id = int(ids[t + 1])
            u_fast = (x - self.core.x_mean) @ self.core.W_out
            if self.core.hier:
                u = u_fast + np.float32(self.core.beta * self.fusion["scale"]) * self.core.u_slow
                hier_pred = int(np.argmax(u[:V]))   # cortex (fast + top-down)
            else:
                u = u_fast
                hier_pred = None
            pred = int(np.argmax(u_fast[:V]))       # fast-area-only ablation
            # keep training-regime reservoir dynamics (readout frozen)
            self.core.learn_reservoir(hier_pred if hier_pred is not None else pred,
                                      target_col, x, u=u_fast)
            if self.core.hier and (t + 1) % self.core.win == 0:
                self.core.step_G()
            # hippocampal query with optional context corruption
            am_val, am_w = None, 0.0
            if use_am and t >= m_ctx - 1:
                ctx_lane = ctxq[t - m_ctx + 1].copy()
                if corrupt_lanes > 0:
                    pidxs = rng.choice(m_ctx, size=min(corrupt_lanes, m_ctx), replace=False)
                    for pidx in pidxs:
                        rand_id = int(train_ids[rng.integers(0, len(train_ids))])
                        ctx_lane[pidx] = lane_key32(rand_id, int(pidx))
                am_val, am_w = self.am.query(ctx_lane)
            am_id = int(am_val) if am_val is not None else None
            am_hit = am_id == target_id
            # fused prediction (hierarchical cortex + hippocampal vote)
            fused_pred, bypass = (hier_pred if hier_pred is not None else pred), False
            if am_id is not None:
                col = self._col_of_id(am_id)
                if col is not None:
                    fused = u.copy()
                    fused[col] += self.cfg.beta_am * am_w
                    fused_pred = int(np.argmax(fused))
                elif am_w >= self.cfg.am_override:
                    bypass = True
            fused_hit = am_hit if bypass else (fused_pred == target_col)
            tot += 1
            if target_col != V:
                inv += 1
                if pred == target_col:
                    snn_ok_inv += 1
                if hier_pred is not None and hier_pred == target_col:
                    hier_ok_inv += 1
                if am_hit:
                    am_ok_inv += 1
                if fused_hit:
                    fused_ok_inv += 1
            if pred == target_col:
                snn_ok += 1
            if hier_pred is not None and hier_pred == target_col:
                hier_ok += 1
            if am_hit:
                am_ok += 1
            if fused_hit:
                fused_ok += 1
        self.core.restore(snap)
        self.core.W_rec = w_rec_snap
        return {
            "n_positions": tot,
            "vocab_coverage": inv / max(tot, 1),
            "snn_acc": snn_ok / max(tot, 1),
            "hier_acc": hier_ok / max(tot, 1),
            "am_acc": am_ok / max(tot, 1),
            "fused_acc": fused_ok / max(tot, 1),
            "snn_acc_invocab": snn_ok_inv / max(inv, 1),
            "hier_acc_invocab": hier_ok_inv / max(inv, 1),
            "am_acc_invocab": am_ok_inv / max(inv, 1),
            "fused_acc_invocab": fused_ok_inv / max(inv, 1),
            "corrupt_lanes": corrupt_lanes,
        }

    # =========================================================== save/load
    def save(self, path: str, am_cap: int = 1_500_000) -> None:
        """Checkpoint everything. ALL hippocampal entries are packed (a
        one-shot episodic memory must persist single-occurrence traces!) in a
        single memory-frugal pass.

        AM packing dedups EXACTLY: every slot an entry occupies is derivable
        from its key (primary full-key hash + leave-one-out hashes), so each
        logical entry is packed once with its surviving copies marked.
        """
        from array import array
        cfg = self.cfg
        am = self.am
        visited = np.zeros(am.size, dtype=np.bool_)
        keys_buf = array("I")
        vals_buf = array("Q")
        counts_buf = array("B")
        for addr in np.nonzero(am.used)[0]:
            if visited[addr]:
                continue
            k = am.keys[addr]
            c = int(am.counts[addr])
            for a in am._slots_of(k):
                if am.used[a] and np.array_equal(am.keys[a], k):
                    visited[a] = True
            keys_buf.extend(k.tolist())
            vals_buf.append(int(am.vals[addr]))
            counts_buf.append(min(c, 255))
            if len(vals_buf) >= am_cap:
                break
        n_e = len(vals_buf)
        if n_e == 0:
            am_keys = np.zeros((0, am.m), dtype=np.uint32)
            am_vals = np.zeros(0, dtype=np.uint64)
            am_counts = np.zeros(0, dtype=np.uint8)
        else:
            am_keys = np.frombuffer(keys_buf, dtype=np.uint32).reshape(n_e, am.m).copy()
            am_vals = np.frombuffer(vals_buf, dtype=np.uint64).copy()
            am_counts = np.frombuffer(counts_buf, dtype=np.uint8).copy()
        np.savez_compressed(
            path,
            config=json.dumps(asdict(cfg)),
            meta=json.dumps(self.meta),
            fusion=json.dumps(self.fusion),
            log2p=self.ept.log2p,
            ept_params=np.array([self.ept.budget_bits, self.ept.min_len, self.ept.max_len, self.ept.alpha]),
            W_rec=self.core.W_rec, W_out=self.core.W_out,
            W_out_G=(self.core.W_outG if self.core.hier
                     else np.zeros((0, 0), dtype=np.float32)),
            inhibitory=self.core.inhibitory,
            vocab_ids=self.vocab_ids,
            reg_ids=self.registry.ids, reg_offsets=self.registry.offsets,
            reg_counts=self.registry.counts,
            reg_buffer=np.frombuffer(self.registry.buffer, dtype=np.uint8),
            am_keys=am_keys, am_vals=am_vals, am_counts=am_counts,
            core_state=json.dumps({
                k: (v if isinstance(v, float) else v.tolist())
                for k, v in self.core.snapshot().items()
            }),
        )

    @classmethod
    def load(cls, path: str) -> "SPIRA":
        z = np.load(path, allow_pickle=False)
        cfg = SPIRAConfig(**json.loads(str(z["config"])))
        m = cls(cfg)
        if "fusion" in z:
            m.fusion = json.loads(str(z["fusion"]))
        m.ept.log2p = z["log2p"].astype(np.float32)
        p = z["ept_params"]
        m.ept.budget_bits, m.ept.min_len, m.ept.max_len, m.ept.alpha = (float(p[0]), int(p[1]), int(p[2]), float(p[3]))
        m.core.W_rec = z["W_rec"].astype(np.float32)
        m.core.W_out = z["W_out"].astype(np.float32)
        m.core.inhibitory = z["inhibitory"].astype(np.bool_)
        w_rec_max = 3.0 * m.core.w_rec_scale
        m.core._rec_lo = np.where(m.core.inhibitory, -w_rec_max, 0.0).astype(np.float32)
        m.core._rec_hi = np.where(m.core.inhibitory, 0.0, w_rec_max).astype(np.float32)
        if "W_out_G" in z and cfg.hierarchy and z["W_out_G"].size > 0:
            m.core.W_outG = z["W_out_G"].astype(np.float32)
        else:
            # v1 checkpoint: hierarchy absent — disable the slow area
            cfg.hierarchy = False
            m.core.hier = False
            m.core.W_outG = None
        m.core.reset()
        if "core_state" in z:
            st = json.loads(str(z["core_state"]))
            restored = {}
            for k, v in st.items():
                if k in ("baseline", "beta"):
                    restored[k] = float(v)
                else:
                    restored[k] = np.asarray(v, dtype=np.float32 if k not in ("refr",) else np.int8)
            m.core.restore(restored)
        m.vocab_ids = z["vocab_ids"].astype(np.uint64)
        m.registry.ids = z["reg_ids"].astype(np.uint64)
        m.registry.offsets = z["reg_offsets"].astype(np.uint64)
        m.registry.counts = z["reg_counts"].astype(np.uint32)
        m.registry.buffer = z["reg_buffer"].tobytes()
        m.meta = json.loads(str(z["meta"]))
        # rebuild hippocampus from packed entries: each entry returns to its
        # derivable slots (primary full-key + leave-one-out), with the same
        # best-effort collision policy as a live write — entries whose slots
        # are taken by a different key are simply not redundantly copied.
        keys = z["am_keys"]; vals = z["am_vals"]; counts = z["am_counts"]
        am = m.am
        for i in range(len(vals)):
            k, v, c = keys[i], int(vals[i]), int(counts[i])
            paddr = am.primary_addr(k)
            if not am.used[paddr] or np.array_equal(am.keys[paddr], k):
                am.keys[paddr] = k
                am.vals[paddr] = v
                am.counts[paddr] = c
                am.used[paddr] = True
            for lane in range(am.m):
                addr = am.loo_addr(k, lane)
                if not am.used[addr]:
                    am.keys[addr] = k
                    am.vals[addr] = v
                    am.counts[addr] = c
                    am.used[addr] = True
        return m