Blame for python/spira/spiking.py

"""Spiking core: LIF reservoir + reward-modulated (three-factor) readout.

Paradigm properties — no backpropagation, no gradients, no autodiff:
  * Leaky integrate-and-fire (LIF) neurons with refractory periods and
    homeostatic adaptive thresholds;
  * Dale's law: each reservoir neuron is permanently excitatory or
    inhibitory (20% inhibitory);
  * Readout learning is a three-factor rule: eligibility (Hebbian pre/post
    coincidence = reservoir trace x outcome) gated by a dopaminergic reward
    prediction error d = r - baseline(r);
  * Recurrent synapses follow a light reward-modulated STDP (trace-based),
    sign-constrained to respect Dale's law;
  * One simulation tick = one token (event-time, not wall-clock ms).

All weights are constructed deterministically from a seeded mulberry32 PRNG
so the TypeScript port regenerates bit-identical connectivity.
"""
from __future__ import annotations

import numpy as np

from .hashing import Mulberry32


class SpikingCore:
    def __init__(
        self,
        d_in: int = 1024,
        k_in: int = 12,
        n_res: int = 512,
        n_out: int = 4097,          # 4096 vocab + 1 UNK column
        seed: int = 20261022,
        in_fanout: int = 8,
        rec_fanout: int = 24,
        inh_frac: float = 0.50,
        lam: float = 0.85,          # membrane leak
        base_thr: float = 1.0,
        thr_min: float = 0.4,
        thr_max: float = 8.0,
        refrac: int = 2,            # ticks
        tau_trace: float = 0.40,    # reservoir spike-trace decay
        w_in_scale: float = 3.0,
        w_rec_scale: float = 0.25,
        w_clip: float = 4.0,
        eta_delta: float = 0.5,     # local Rescorla-Wagner delta rate (LMS-normalized)
        eta_stdp: float = 0.0015,   # recurrent modulated STDP rate
        eta_homeo: float = 0.002,   # adaptive-threshold rate
        target_rate: float = 0.06,  # target per-neuron firing probability/tick
        baseline_momentum: float = 0.995,
        # ---- hierarchical slow area G (multi-scale cortex, v2) ----
        hier: bool = True,
        n_g: int = 256,              # slow-area neurons
        win: int = 8,                # fast ticks per slow step
        w_g_scale: float = 6.0,      # L -> G feedforward weight scale
        w_recg_scale: float = 0.15,  # G recurrent scale (fixed, no STDP)
        lam_g: float = 0.6,          # slow-area membrane leak
        tau_trace_g: float = 0.5,    # slow-area spike-trace decay
        target_rate_g: float = 0.10, # target firing per slow step
        eta_homeo_g: float = 0.01,
        eta_delta_g: float = 2.0,     # slow-readout delta rate (8x fewer updates)
        eta_beta: float = 0.02,      # top-down gain learning rate
        beta_init: float = 0.25,
        beta_max: float = 2.0,
        # ---- v3 (Muse): multi-timescale traces + direct input pathway ----
        trace_taus: tuple = (0.4,),   # reservoir trace decays; readout sees ALL banks
        direct_input: bool = False,   # readout also sees a decayed trace of the
                                      # exact input-neuron pattern (ESN-standard
                                      # [state; input] readout, omitted in v1/v2)
        direct_tau: float = 0.5,      # decay of the direct input trace
        direct_gain: float = 1.0,
    ) -> None:
        self.d_in, self.k_in, self.n_res, self.n_out = d_in, k_in, n_res, n_out
        self.seed = seed
        self.in_fanout, self.rec_fanout = in_fanout, rec_fanout
        self.inh_frac = inh_frac
        self.lam, self.base_thr = lam, base_thr
        self.thr_min, self.thr_max = thr_min, thr_max
        self.refrac = refrac
        self.tau_trace = tau_trace
        self.w_in_scale, self.w_rec_scale = w_in_scale, w_rec_scale
        self.w_clip = w_clip
        self.eta_delta, self.eta_stdp = eta_delta, eta_stdp
        self.eta_homeo, self.target_rate = eta_homeo, target_rate
        self.baseline_momentum = baseline_momentum
        self.hier = hier
        self.n_g, self.win = n_g, win
        self.w_g_scale, self.w_recg_scale = w_g_scale, w_recg_scale
        self.lam_g, self.tau_trace_g = lam_g, tau_trace_g
        self.target_rate_g, self.eta_homeo_g = target_rate_g, eta_homeo_g
        self.eta_delta_g = eta_delta_g
        self.eta_beta, self.beta0, self.beta_max = eta_beta, beta_init, beta_max
        # v3 structural params
        self.trace_taus = tuple(float(t) for t in trace_taus)
        self.direct_input = bool(direct_input)
        self.direct_tau, self.direct_gain = float(direct_tau), float(direct_gain)
        self.n_banks = len(self.trace_taus)
        self.n_state = self.n_res * self.n_banks + (self.d_in if self.direct_input else 0)

        # ------------------------------------------------ deterministic wiring
        rng = Mulberry32(seed)
        # Input -> reservoir: sparse fan-out per input neuron, dense storage.
        self.W_in = np.zeros((d_in, n_res), dtype=np.float32)
        for i in range(d_in):
            for _ in range(in_fanout):
                t = int(rng.next_float() * n_res)
                self.W_in[i, t] += rng.next_weight(w_in_scale)
        # Reservoir recurrent: sparse fan-out per source, Dale's law by source.
        self.W_rec = np.zeros((n_res, n_res), dtype=np.float32)
        self.inhibitory = np.zeros(n_res, dtype=np.bool_)
        w_rec_max = 3.0 * w_rec_scale
        for j in range(n_res):
            is_inh = rng.next_float() < inh_frac
            self.inhibitory[j] = is_inh
            sign = -1.0 if is_inh else 1.0
            for _ in range(rec_fanout):
                t = int(rng.next_float() * n_res)
                self.W_rec[j, t] += sign * abs(rng.next_weight(w_rec_scale))
        # Keep Dale's law under plasticity:
        #   excitatory rows clipped to [0, w_rec_max], inhibitory to [-w_rec_max, 0]
        self._rec_lo = np.where(self.inhibitory, -w_rec_max, 0.0).astype(np.float32)
        self._rec_hi = np.where(self.inhibitory, 0.0, w_rec_max).astype(np.float32)
        self.W_rec = np.clip(self.W_rec, self._rec_lo[:, None], self._rec_hi[:, None])

        # Readout: small seeded init, learned by the three-factor rule.
        rng2 = Mulberry32(seed + 1)
        self.W_out = np.zeros((self.n_state, n_out), dtype=np.float32)  # v3: n_state rows
        # sparse-ish init: each column connected from a random 25% of neurons
        for c in range(n_out):
            for _ in range(n_res // 4):
                r = int(rng2.next_float() * n_res)
                self.W_out[r, c] += rng2.next_weight(0.05)

        # ------------------------------------------ slow area G (hierarchy)
        # Wired from SEPARATE PRNG streams (seed+2 / seed+3) so the fast-area
        # sequences above are byte-identical to v1 — the TS port depends on
        # the exact call order.
        if hier:
            rng3 = Mulberry32(seed + 2)
            # L -> G: per fast-source fanout 2. Sign follows Dale's law of the
            # PRESYNAPTIC fast neuron: excitatory cells excite G, inhibitory
            # cells inhibit it (G sees the pooled fast spike volley).
            self.W_inG = np.zeros((n_res, n_g), dtype=np.float32)
            for i in range(n_res):
                sign = -1.0 if self.inhibitory[i] else 1.0
                for _ in range(2):
                    t = int(rng3.next_float() * n_g)
                    self.W_inG[i, t] += sign * abs(rng3.next_weight(w_g_scale))
            # G recurrent: sparse, Dale by source, FIXED (slower areas
            # plasticize slower — G plasticity lives in its readout only)
            self.W_recG = np.zeros((n_g, n_g), dtype=np.float32)
            g_rec_max = 3.0 * w_recg_scale
            self.inhibitory_g = np.zeros(n_g, dtype=np.bool_)
            for j in range(n_g):
                is_inh = rng3.next_float() < 0.5
                self.inhibitory_g[j] = is_inh
                sign = -1.0 if is_inh else 1.0
                for _ in range(8):
                    t = int(rng3.next_float() * n_g)
                    self.W_recG[j, t] += sign * abs(rng3.next_weight(w_recg_scale))
            self._recg_lo = np.where(self.inhibitory_g, -g_rec_max, 0.0).astype(np.float32)
            self._recg_hi = np.where(self.inhibitory_g, 0.0, g_rec_max).astype(np.float32)
            self.W_recG = np.clip(self.W_recG, self._recg_lo[:, None], self._recg_hi[:, None])
            # G readout: same seeded sparse init pattern as the fast readout
            rng4 = Mulberry32(seed + 3)
            self.W_outG = np.zeros((n_g, n_out), dtype=np.float32)
            for c in range(n_out):
                for _ in range(n_g // 4):
                    r = int(rng4.next_float() * n_g)
                    self.W_outG[r, c] += rng4.next_weight(0.05)

        self.reset()

    # ----------------------------------------------------------------- state
    def reset(self) -> None:
        self.v = np.zeros(self.n_res, dtype=np.float32)
        self.thr = np.full(self.n_res, self.base_thr, dtype=np.float32)
        self.rhat = np.full(self.n_res, self.target_rate, dtype=np.float32)
        self.refr = np.zeros(self.n_res, dtype=np.int8)
        self.x = np.zeros(self.n_state, dtype=np.float32)   # v3: augmented readout state
        self.x_prev = np.zeros(self.n_res, dtype=np.float32)  # reservoir trace t-1 (STDP)
        self.x_mean = np.zeros(self.n_state, dtype=np.float32)
        self.x_eff = np.zeros(self.n_state, dtype=np.float32)
        self._banks = np.zeros((self.n_banks, self.n_res), dtype=np.float32)
        self._direct = np.zeros(self.d_in if self.direct_input else 0, dtype=np.float32)
        self.s = np.zeros(self.n_res, dtype=np.float32)   # last spike vector
        self.baseline = 0.0
        self.tick = 0
        if self.hier:
            self.v_G = np.zeros(self.n_g, dtype=np.float32)
            self.thr_G = np.full(self.n_g, self.base_thr, dtype=np.float32)
            self.rhat_G = np.full(self.n_g, self.target_rate_g, dtype=np.float32)
            self.x_G = np.zeros(self.n_g, dtype=np.float32)
            self.x_mean_G = np.zeros(self.n_g, dtype=np.float32)
            self.s_G = np.zeros(self.n_g, dtype=np.float32)
            self.win_spikes = np.zeros(self.n_res, dtype=np.float32)
            self.u_slow = np.zeros(self.n_out, dtype=np.float32)
            self.beta = self.beta0

    # ------------------------------------------------------------- snapshots
    def snapshot(self) -> dict:
        """Full dynamical state (the homeostatic scalars are LEARNED state —
        thresholds and the trace mean co-adapt with the readout)."""
        d = {
            "v": self.v.copy(), "thr": self.thr.copy(), "rhat": self.rhat.copy(),
            "refr": self.refr.copy(), "x": self.x.copy(), "x_prev": self.x_prev.copy(),
            "x_mean": self.x_mean.copy(), "s": self.s.copy(), "baseline": self.baseline,
            "banks": self._banks.copy(), "direct": self._direct.copy(),
        }
        if self.hier:
            d.update({
                "v_G": self.v_G.copy(), "thr_G": self.thr_G.copy(),
                "rhat_G": self.rhat_G.copy(), "x_G": self.x_G.copy(),
                "x_mean_G": self.x_mean_G.copy(), "s_G": self.s_G.copy(),
                "win_spikes": self.win_spikes.copy(), "u_slow": self.u_slow.copy(),
                "beta": self.beta,
            })
        return d

    def restore(self, st: dict) -> None:
        self.v = st["v"].copy()
        self.thr = st["thr"].copy()
        self.rhat = st["rhat"].copy()
        self.refr = st["refr"].copy()
        self.x = st["x"].copy()
        self.x_prev = st["x_prev"].copy()
        self.x_mean = st["x_mean"].copy()
        self.s = st["s"].copy()
        self.baseline = st["baseline"]
        if "banks" in st:
            self._banks = st["banks"].copy()
            self._direct = st["direct"].copy()
        else:  # reconstruct banks from the composed state vector
            nb = self.n_banks * self.n_res
            self._banks = self.x[:nb].reshape(self.n_banks, self.n_res).copy()
            if self.direct_input:
                self._direct = self.x[nb:].copy()
        if self.hier and "v_G" in st:
            self.v_G = st["v_G"].copy()
            self.thr_G = st["thr_G"].copy()
            self.rhat_G = st["rhat_G"].copy()
            self.x_G = st["x_G"].copy()
            self.x_mean_G = st["x_mean_G"].copy()
            self.s_G = st["s_G"].copy()
            self.win_spikes = st["win_spikes"].copy()
            self.u_slow = st["u_slow"].copy()
            self.beta = st["beta"]

    def reset_dynamics(self) -> None:
        """Reset transient dynamics but KEEP learned homeostatic scalars
        (thresholds, rate estimates, trace mean, top-down gain) — used before
        generation."""
        thr, rhat, x_mean = self.thr, self.rhat, self.x_mean
        if self.hier:
            thr_G, rhat_G, x_mean_G, beta = self.thr_G, self.rhat_G, self.x_mean_G, self.beta
        self.reset()
        self.thr, self.rhat, self.x_mean = thr, rhat, x_mean
        if self.hier:
            self.thr_G, self.rhat_G, self.x_mean_G, self.beta = thr_G, rhat_G, x_mean_G, beta

    # ------------------------------------------------------------- inference
    def step(self, input_idx) -> np.ndarray:
        """Advance one tick with the given active input neurons.

        Returns the post-synaptic trace ``x`` (the state used for readout).
        """
        W_in, W_rec = self.W_in, self.W_rec
        I = W_in[input_idx].sum(axis=0, dtype=np.float32)
        I += self.s @ W_rec
        v = self.lam * self.v + I
        can = self.refr == 0
        v[~can] = 0.0
        spikes = (v >= self.thr) & can
        v[spikes] = 0.0
        np.subtract(self.refr, 1, out=self.refr, where=self.refr > 0)
        self.refr[spikes] = self.refrac
        self.v = v
        self.s = spikes.astype(np.float32)
        # homeostatic threshold via EMA firing-rate estimate
        self.rhat += 0.02 * (self.s - self.rhat)
        self.thr += self.eta_homeo * (self.rhat - self.target_rate)
        np.clip(self.thr, self.thr_min, self.thr_max, out=self.thr)
        # traces (v3): x_prev keeps the BANK-0 (fast) reservoir trace for STDP;
        # the readout state x is the concatenation of all trace banks plus the
        # direct input trace — context older than one tick survives in the
        # slow banks and in the direct trace's decayed identity history.
        self.x_prev = self._banks[0].copy()
        for i, tau in enumerate(self.trace_taus):
            if i == 0:
                self._banks[i] = tau * self._banks[i] + self.s
            else:
                # slow banks are NORMALIZED traces (decay-weighted averages):
                # a raw tau=0.95 sum would sit at ~20x the fast bank's scale
                # and swamp both the readout geometry and the LMS rate.
                self._banks[i] = tau * self._banks[i] + (1.0 - tau) * self.s
        if self.direct_input:
            self._direct *= np.float32(self.direct_tau)
            np.add.at(self._direct, input_idx, np.float32(self.direct_gain))
            self.x = np.concatenate([*self._banks, self._direct])
        elif self.n_banks > 1:
            self.x = np.concatenate([*self._banks])
        else:
            self.x = self._banks[0]
        # slow EMA of the trace: homeostatic "synaptic scaling" analogue used
        # to mean-center the readout input (kills the shared dense background
        # that otherwise dominates column norms)
        self.x_mean = 0.999 * self.x_mean + 0.001 * self.x
        self.x_eff = self.x - self.x_mean
        self.tick += 1
        if self.hier:
            self.win_spikes += self.s
        return self.x

    def scores(self, x: np.ndarray) -> np.ndarray:
        """Readout scores from a (raw) trace — centers with the EMA mean."""
        return (x - self.x_mean) @ self.W_out

    # ------------------------------------------------- slow area (hierarchy)
    def step_G(self) -> None:
        """Advance the slow area one slow step (one window of ``win`` fast
        ticks).

        Input: the pooled fast-area spike volley over the window (bottom-up,
        sign-preserving). Output: the cached slow readout ``u_slow``, used as
        a top-down prior fused into the fast prediction with gain ``beta``.
        Pure inference dynamics (homeostasis included, same regime as the
        fast area); learning happens in ``learn_G``.
        """
        c = self.win_spikes * np.float32(1.0 / self.win)
        I = c @ self.W_inG
        I += self.s_G @ self.W_recG
        v = self.lam_g * self.v_G + I
        spikes = v >= self.thr_G
        v[spikes] = 0.0
        self.v_G = v
        self.s_G = spikes.astype(np.float32)
        self.rhat_G += 0.05 * (self.s_G - self.rhat_G)
        self.thr_G += self.eta_homeo_g * (self.rhat_G - self.target_rate_g)
        np.clip(self.thr_G, self.thr_min, self.thr_max, out=self.thr_G)
        self.x_G = self.tau_trace_g * self.x_G + self.s_G
        self.x_mean_G = 0.999 * self.x_mean_G + 0.001 * self.x_G
        self.win_spikes[:] = 0.0
        self.u_slow = (self.x_G - self.x_mean_G) @ self.W_outG

    def learn_G(self, target: int, d: float) -> int:
        """Delta-rule update on the SLOW readout, gated by the same dopaminergic
        RPE that gates the fast one (a global neuromodulator broadcasts to both
        areas — each readout computes its own local error from its own
        activity). ``target`` is the majority vocabulary column of the NEXT
        window. Returns 1 if the slow readout's own top prediction matches it.
        """
        x_effG = self.x_G - self.x_mean_G
        u = self.u_slow
        V = self.n_out - 1
        own_pred = int(np.argmax(u[:V]))
        correct = 1 if own_pred == target else 0
        nx = 1.0 / (float(x_effG @ x_effG) + 1e-6)
        k = min(16, V)
        act = np.argpartition(u[:V], -k)[-k:]
        if target < V:
            cols = np.concatenate([act, np.array([target], dtype=np.int64)])
            desired = np.zeros(len(cols), dtype=np.float32)
            desired[-1] = 1.0
        else:
            cols = act
            desired = np.zeros(len(cols), dtype=np.float32)
        errs = u[cols] - desired
        gate = float(np.clip(1.0 + d, 0.25, 2.0))
        eta = np.float32(self.eta_delta_g * gate * nx)
        W = self.W_outG
        W[:, cols] -= eta * errs[None, :] * x_effG[:, None]
        W[:, cols] = np.clip(W[:, cols], -self.w_clip, self.w_clip)
        return correct

    # ------------------------------------------------------------- learning
    def learn(self, pred: int, target: int, x: np.ndarray,
              u: np.ndarray | None = None, active_k: int = 16) -> tuple[int, float, float]:
        """Local delta-rule readout update (Rescorla-Wagner) gated by a
        dopaminergic reward-prediction error, + modulated recurrent STDP.

        Each readout unit computes its OWN prediction error from its own
        activity (no gradients, no chain rule, no backprop): units in the
        active top-k set and the instructive target unit move their weights
        to reduce (score - desired). The learning rate is gated by dopamine
        d = reward - baseline (RPE), the classical Rescorla-Wagner / TD
        account of neuromodulation.

        Returns (correct, reward, dopamine).
        """
        correct = 1 if pred == target else 0
        reward = float(correct)
        self.baseline = self.baseline_momentum * self.baseline + (1 - self.baseline_momentum) * reward
        d = reward - self.baseline                       # dopaminergic RPE

        W_out = self.W_out
        x_eff = x - self.x_mean
        nx = 1.0 / (float(x_eff @ x_eff) + 1e-6)         # LMS normalization
        V = self.n_out - 1
        if u is None:
            u = x_eff @ W_out
        # active set: top-k scored columns (+ target if in vocabulary)
        k = min(active_k, V)
        act = np.argpartition(u[:V], -k)[-k:]
        if target < V:
            cols = np.concatenate([act, np.array([target], dtype=np.int64)])
        else:
            cols = act
        desired = np.zeros(len(cols), dtype=np.float32)
        if target < V:
            desired[-1] = 1.0                             # one-hot with margin 1
        errs = u[cols] - desired
        # dopamine gates plasticity amplitude (dips shrink, bursts amplify)
        gate = float(np.clip(1.0 + d, 0.25, 2.0))
        eta = np.float32(self.eta_delta * gate * nx)
        W_out[:, cols] -= eta * errs[None, :] * x_eff[:, None]
        # clip in place (fancy-indexed views need explicit writeback)
        W_out[:, cols] = np.clip(W_out[:, cols], -self.w_clip, self.w_clip)

        # light reward-modulated STDP on recurrent synapses of spiked neurons
        if self.s.any():
            spiked = np.nonzero(self.s)[0]
            W_rec = self.W_rec
            delta = np.float32(self.eta_stdp * d)
            # pre-trace (x_prev) paired with post spikes, gated by dopamine
            W_rec[:, spiked] += delta * self.x_prev[:, None]
            np.clip(W_rec[:, spiked], self._rec_lo[:, None], self._rec_hi[:, None],
                    out=W_rec[:, spiked])
        # top-down gain: TRUST dynamics, gated by the same RPE (a
        # noradrenergic gain-modulation analogue) with counterfactual credit:
        #   fused right + fast wrong  -> the slow prior SAVED the choice -> trust builds (additive)
        #   fused wrong + fast right  -> the slow prior CORRUPTED it   -> trust decays (multiplicative)
        # Trust building slowly and decaying fast (loss aversion) lets the
        # system keep a genuinely informative slow area while driving a
        # noise- or frequency-prior-dominated one to beta ~ 0.
        if self.hier and target < V:
            pred_fast = int(np.argmax(u[:V]))
            if pred_fast != target and correct:          # slow saved it
                self.beta = min(self.beta + self.eta_beta * gate, self.beta_max)
            elif pred_fast == target and not correct:    # slow corrupted it
                self.beta *= float(0.98) ** gate
                if self.beta < 1e-4:
                    self.beta = 0.0
            else:
                us_t = float(self.u_slow[target])
                us_p = float(self.u_slow[pred_fast])
                norm = abs(us_t) + abs(us_p) + 1e-6
                margin = float(np.tanh(3.0 * (us_t - us_p) / norm))
                self.beta = float(np.clip(
                    self.beta + 0.25 * self.eta_beta * gate * margin,
                    0.0, self.beta_max))
        return correct, reward, d

    def learn_reservoir(self, pred: int, target: int, x: np.ndarray,
                         u: np.ndarray | None = None) -> tuple[int, float, float]:
        """Training-regime dynamics WITHOUT readout plasticity: same reward,
        same dopaminergic RPE, same recurrent STDP — only the W_out delta
        update is skipped. Used at inference so the reservoir runs in its
        training operating regime (continual plasticity, frozen readout)."""
        correct = 1 if pred == target else 0
        reward = float(correct)
        self.baseline = self.baseline_momentum * self.baseline + (1 - self.baseline_momentum) * reward
        d = reward - self.baseline
        if self.s.any():
            spiked = np.nonzero(self.s)[0]
            W_rec = self.W_rec
            delta = np.float32(self.eta_stdp * d)
            W_rec[:, spiked] += delta * self.x_prev[:, None]
            W_rec[:, spiked] = np.clip(W_rec[:, spiked], self._rec_lo[:, None], self._rec_hi[:, None])
        return correct, reward, d

    # ------------------------------------------------------- training driver
    def train_pass(
        self,
        neurons: np.ndarray,      # (T, k_in) int32 active input neurons per tick
        cols: np.ndarray,         # (T,) int64 target column per tick (col of token t+1)
        log_every: int = 5000,
        log_fn=None,
        decay_every: int = 1024,
        decay: float = 1e-4,
        record_raster_from: int | None = None,
    ) -> dict:
        """One epoch over a token stream (teacher forcing on the next token).

        ``neurons[t]`` drives the reservoir at tick t; the prediction made
        from the resulting state is scored against ``cols[t]``.
        """
        self.reset()
        T = len(cols)
        acc_ema = None
        d_ema = 0.0
        n_correct = 0
        n_correct_fast = 0
        raster: list[tuple[int, np.ndarray]] = []
        V_vocab = self.n_out - 1          # UNK occupies the last column
        import time
        t0 = time.perf_counter()
        for t in range(T):
            x = self.step(neurons[t])
            u = (x - self.x_mean) @ self.W_out
            if self.hier:
                # behavioural prediction = fast readout + top-down slow prior
                fused = u + np.float32(self.beta) * self.u_slow
                pred = int(np.argmax(fused[:V_vocab]))
            else:
                pred = int(np.argmax(u[:V_vocab]))   # vocabulary-only prediction
            correct, reward, d = self.learn(pred, int(cols[t]), x, u=u)
            n_correct += correct
            if self.hier:
                pred_fast = int(np.argmax(u[:V_vocab]))
                if pred_fast == int(cols[t]):
                    n_correct_fast += 1
                if (t + 1) % self.win == 0:
                    self.step_G()
                    if t + 1 < T:
                        seg = cols[t + 1 : min(t + 1 + self.win, T)]
                        g_target = int(np.bincount(seg).argmax())
                        self.learn_G(g_target, d)
            acc_ema = reward if acc_ema is None else 0.999 * acc_ema + 0.001 * reward
            d_ema = 0.999 * d_ema + 0.001 * d
            if record_raster_from is not None and t >= record_raster_from:
                raster.append((t, np.nonzero(self.s)[0].astype(np.int16)))
            if (t + 1) % log_every == 0 and log_fn is not None:
                dt = time.perf_counter() - t0
                log_fn({
                    "tick": t + 1, "total": T,
                    "acc_ema": round(acc_ema, 4),
                    "acc_fast_ema": (round(0.999 * (n_correct_fast / max(t + 1, 1)), 4)
                                     if self.hier else None),
                    "beta": (round(self.beta, 3) if self.hier else None),
                    "dopamine_ema": round(d_ema, 4),
                    "ticks_per_sec": round((t + 1) / max(dt, 1e-9), 1),
                    "elapsed_s": round(dt, 1),
                })
            if (t + 1) % decay_every == 0:
                self.W_out *= np.float32(1.0 - decay)
                if self.hier:
                    self.W_outG *= np.float32(1.0 - decay)
        dt = time.perf_counter() - t0
        return {
            "ticks": T,
            "accuracy": n_correct / max(T, 1),
            "accuracy_fast_only": (n_correct_fast / max(T, 1)) if self.hier else None,
            "beta": float(self.beta) if self.hier else None,
            "acc_ema_final": acc_ema,
            "elapsed_s": dt,
            "ticks_per_sec": T / max(dt, 1e-9),
            "raster": raster,
        }

    # ------------------------------------------------------------ (de)hydrate
    def state_dict(self) -> dict:
        sd = {
            "config": {
                "d_in": self.d_in, "k_in": self.k_in, "n_res": self.n_res,
                "n_out": self.n_out, "seed": self.seed,
                "in_fanout": self.in_fanout, "rec_fanout": self.rec_fanout,
                "inh_frac": self.inh_frac, "lam": self.lam,
                "base_thr": self.base_thr, "refrac": self.refrac,
                "tau_trace": self.tau_trace,
                "w_in_scale": self.w_in_scale, "w_rec_scale": self.w_rec_scale,
                "w_clip": self.w_clip, "eta_delta": self.eta_delta, "eta_stdp": self.eta_stdp,
                "eta_homeo": self.eta_homeo,
                "target_rate": self.target_rate,
                "hier": self.hier, "n_g": self.n_g, "win": self.win,
                "w_g_scale": self.w_g_scale, "w_recg_scale": self.w_recg_scale,
                "lam_g": self.lam_g, "tau_trace_g": self.tau_trace_g,
                "target_rate_g": self.target_rate_g, "eta_homeo_g": self.eta_homeo_g,
                "eta_delta_g": self.eta_delta_g,
                "eta_beta": self.eta_beta, "beta_init": self.beta0, "beta_max": self.beta_max,
                "trace_taus": list(self.trace_taus), "direct_input": self.direct_input,
                "direct_tau": self.direct_tau, "direct_gain": self.direct_gain,
            },
            "W_rec": self.W_rec.astype(np.float32),
            "W_out": self.W_out.astype(np.float32),
            "inhibitory": self.inhibitory,
        }
        if self.hier:
            sd["W_out_G"] = self.W_outG.astype(np.float32)
        return sd

    def load_state_dict(self, sd: dict) -> None:
        cfg = sd["config"]
        self.__dict__.update(cfg)
        # v3 derived structure (defaults reproduce v2 for old checkpoints)
        self.trace_taus = tuple(float(t) for t in getattr(self, "trace_taus", (self.tau_trace,)))
        self.direct_input = bool(getattr(self, "direct_input", False))
        self.n_banks = len(self.trace_taus)
        self.n_state = int(sd["W_out"].shape[0])
        self.W_rec = np.asarray(sd["W_rec"], dtype=np.float32)
        self.W_out = np.asarray(sd["W_out"], dtype=np.float32)
        self.inhibitory = np.asarray(sd["inhibitory"], dtype=np.bool_)
        w_rec_max = 3.0 * self.w_rec_scale
        self._rec_lo = np.where(self.inhibitory, -w_rec_max, 0.0).astype(np.float32)
        self._rec_hi = np.where(self.inhibitory, 0.0, w_rec_max).astype(np.float32)
        if "W_out_G" in sd and self.hier:
            self.W_outG = np.asarray(sd["W_out_G"], dtype=np.float32)
        else:
            # v1 checkpoint without the slow area
            self.hier = False
        self.reset()