SPIRA v2: entropy-patch tokenization + hierarchical dopamine-gated spiking cortex + hippocampal one-shot memory with replay consolidation
032a961 verified 22 days ago
Raw History Blame Contribute Delete
29.6 kB
/**
 * SPIRA — Spiking Plastic Recurrent Architecture, TypeScript inference port.
 *
 * Loads the JSON artifact exported by `spira.export_deploy` and provides:
 *   - EPT tokenization (entropy patches)
 *   - LIF reservoir simulation (weights regenerated from seed, bit-identical)
 *   - delta-rule readout scoring (int8-quantized W_out, mean-centered trace)
 *   - hippocampal associative memory with hybrid addressing:
 *       primary full-key hash (pattern separation / capacity) +
 *       leave-one-out hashes (CA3-style completion under corruption)
 *   - closed-loop generation and ONE-SHOT episodic teaching (in browser)
 *
 * Everything runs on plain CPU in the browser — no GPU, no server.
 */

import { Mulberry32, fnv1a64, inputNeurons, laneKey32, splitmix64 } from "./hashing";
import { EPT, Patch } from "./ept";

const b64ToU8 = (b64: string): Uint8Array => {
  const bin = atob(b64);
  const u8 = new Uint8Array(bin.length);
  for (let i = 0; i < bin.length; i++) u8[i] = bin.charCodeAt(i);
  return u8;
};

export interface SpiraDeploy {
  format: string;
  config: Record<string, number>;
  ept: { budget_bits: number; min_len: number; max_len: number; alpha: number; log2p_b64: string };
  reservoir: { w_rec_rows: string; w_rec_cols: string; w_rec_vals_b64: string; w_rec_n: number };
  readout: {
    shape: [number, number];
    vocab_ids_hex: string[];
    w_out_q_b64: string;
    w_out_scale_b64: string;
  };
  hier?: {
    n_g: number; win: number; lam_g: number; tau_trace_g: number;
    target_rate_g: number; eta_homeo_g: number; beta: number; fusion_scale: number;
    w_ing_b64: string; w_recg_b64: string;
    w_outg_q_b64: string; w_outg_scale_b64: string;
  };
  am: { n_entries: number; packed_b64: string; entry_bytes: number };
  registry: { n: number; ids_b64: string; lens_b64: string; buffer_b64: string };
  core_state?: {
    x_mean_b64: string; thr_b64: string; rhat_b64: string;
    x_mean_g_b64?: string; thr_g_b64?: string; rhat_g_b64?: string;
  };
  fusion?: { scale: number; sigma_fast?: number; sigma_slow?: number };
  meta: Record<string, number>;
}

export interface TeachResult {
  tokens: number;
  newPatches: number;
  associations: number;
  ms: number;
}

const rotl = (x: number, r: number) => ((x << r) | (x >>> (32 - r))) >>> 0;

export class SpiraModel {
  ept: EPT;
  // dimensions
  dIn: number;
  kIn: number;
  nRes: number;
  nOut: number;
  vocabSize: number;
  // reservoir
  wIn!: Float32Array;
  wRec!: Float32Array;
  lam: number;
  baseThr: number;
  thrMin: number;
  thrMax: number;
  refrac: number;
  tauTrace: number;
  etaHomeo: number;
  targetRate: number;
  // state
  v!: Float32Array;
  thr!: Float32Array;
  rhat!: Float32Array;
  refr!: Int8Array;
  x!: Float32Array;
  xPrev!: Float32Array;
  xMean!: Float32Array;
  s!: Float32Array;
  // readout
  wOutQ!: Int8Array;
  wOutScale!: Float32Array;
  vocabIds: bigint[] = [];
  colOf = new Map<bigint, number>();
  // fusion
  betaAm: number;
  amOverride: number;
  temperature: number;
  topK: number;
  // hippocampus (hybrid addressing)
  amSize: number;
  amM: number;
  amMinLanes: number;
  amKeys!: Uint32Array;
  amVals!: BigUint64Array;
  amCounts!: Uint8Array;
  amUsed!: Uint8Array;
  // registry
  regIds: bigint[] = [];
  regLens: number[] = [];
  regOffsets: number[] = [];
  regBuffer!: Uint8Array;
  regIndexOf = new Map<bigint, number>();
  // hierarchical slow area (G) — v2
  hier: boolean;
  nG = 0;
  win = 0;
  lamG = 0;
  tauTraceG = 0;
  targetRateG = 0;
  etaHomeoG = 0;
  beta = 0;
  fusionScale = 1;
  wInG!: Float32Array;
  wRecG!: Float32Array;
  wOutGQ!: Int8Array;
  wOutGScale!: Float32Array;
  vG!: Float32Array;
  thrG!: Float32Array;
  rhatG!: Float32Array;
  xG!: Float32Array;
  xMeanG!: Float32Array;
  sG!: Float32Array;
  winSpikes!: Float32Array;
  uSlow!: Float32Array;
  winCount = 0;
  meta: Record<string, number>;

  constructor(art: SpiraDeploy) {
    const c = art.config;
    this.dIn = c.d_in;
    this.kIn = c.k_in;
    this.nRes = c.n_res;
    this.vocabSize = c.vocab_size;
    this.nOut = this.vocabSize + 1;
    this.lam = c.lam;
    this.baseThr = c.base_thr;
    this.thrMin = c.thr_min;
    this.thrMax = c.thr_max;
    this.refrac = c.refrac;
    this.tauTrace = c.tau_trace;
    this.etaHomeo = c.eta_homeo;
    this.targetRate = c.target_rate;
    this.betaAm = c.beta_am;
    this.amOverride = c.am_override;
    this.temperature = c.temperature;
    this.topK = c.top_k;
    this.amSize = c.am_size;
    this.amM = c.am_m;
    this.amMinLanes = c.am_min_lanes;
    this.meta = art.meta;

    // ---- EPT ----
    const log2p = new Float32Array(b64ToU8(art.ept.log2p_b64).buffer);
    this.ept = new EPT(log2p, art.ept.budget_bits, art.ept.min_len, art.ept.max_len);

    // ---- reservoir: regenerate input weights from the seed ----
    this.buildW_in(c.seed, c.in_fanout, c.rec_fanout, c.inh_frac, c.w_in_scale, c.w_rec_scale);
    this.buildW_rec(art.reservoir);

    // ---- readout ----
    this.wOutQ = new Int8Array(b64ToU8(art.readout.w_out_q_b64).buffer);
    this.wOutScale = new Float32Array(b64ToU8(art.readout.w_out_scale_b64).buffer);
    this.vocabIds = art.readout.vocab_ids_hex.map((h) => BigInt(h));
    this.vocabIds.forEach((id, col) => this.colOf.set(id, col));

    // ---- hippocampus ----
    this.amKeys = new Uint32Array(this.amSize * this.amM);
    this.amVals = new BigUint64Array(this.amSize);
    this.amCounts = new Uint8Array(this.amSize);
    this.amUsed = new Uint8Array(this.amSize);
    this.loadAmEntries(art.am);

    // ---- registry ----
    this.loadRegistry(art.registry);

    // ---- hierarchical slow area (v2) ----
    this.hier = !!(art.hier && art.config.hierarchy);
    if (this.hier && art.hier) {
      const h = art.hier;
      this.nG = h.n_g;
      this.win = h.win;
      this.lamG = h.lam_g;
      this.tauTraceG = h.tau_trace_g;
      this.targetRateG = h.target_rate_g;
      this.etaHomeoG = h.eta_homeo_g;
      this.beta = h.beta;
      this.fusionScale = h.fusion_scale;
      this.wInG = new Float32Array(b64ToU8(h.w_ing_b64).buffer);
      this.wRecG = new Float32Array(b64ToU8(h.w_recg_b64).buffer);
      this.wOutGQ = new Int8Array(b64ToU8(h.w_outg_q_b64).buffer);
      this.wOutGScale = new Float32Array(b64ToU8(h.w_outg_scale_b64).buffer);
    }

    // ---- learned homeostatic state ----
    this.resetCore();
    if (art.core_state) {
      this.xMean = new Float32Array(b64ToU8(art.core_state.x_mean_b64).buffer);
      this.thr = new Float32Array(b64ToU8(art.core_state.thr_b64).buffer);
      this.rhat = new Float32Array(b64ToU8(art.core_state.rhat_b64).buffer);
      if (this.hier && art.core_state.x_mean_g_b64) {
        this.xMeanG = new Float32Array(b64ToU8(art.core_state.x_mean_g_b64).buffer);
        this.thrG = new Float32Array(b64ToU8(art.core_state.thr_g_b64).buffer);
        this.rhatG = new Float32Array(b64ToU8(art.core_state.rhat_g_b64).buffer);
      }
    }
  }

  // ------------------------------------------------------------ build/load
  private buildW_in(seed: number, inFanout: number, recFanout: number, inhFrac: number,
                    wInScale: number, wRecScale: number) {
    const rng = new Mulberry32(seed);
    this.wIn = new Float32Array(this.dIn * this.nRes);
    for (let i = 0; i < this.dIn; i++) {
      for (let f = 0; f < inFanout; f++) {
        const t = Math.floor(rng.nextFloat() * this.nRes);
        this.wIn[i * this.nRes + t] += rng.nextWeight(wInScale);
      }
    }
    const wRec = new Float32Array(this.nRes * this.nRes);
    const wRecMax = 3 * wRecScale;
    const lo = new Float32Array(this.nRes);
    const hi = new Float32Array(this.nRes);
    for (let j = 0; j < this.nRes; j++) {
      const isInh = rng.nextFloat() < inhFrac;
      const sign = isInh ? -1 : 1;
      lo[j] = isInh ? -wRecMax : 0;
      hi[j] = isInh ? 0 : wRecMax;
      for (let f = 0; f < recFanout; f++) {
        const t = Math.floor(rng.nextFloat() * this.nRes);
        wRec[j * this.nRes + t] += sign * Math.abs(rng.nextWeight(wRecScale));
      }
    }
    for (let j = 0; j < this.nRes; j++) {
      for (let t = 0; t < this.nRes; t++) {
        const w = wRec[j * this.nRes + t];
        wRec[j * this.nRes + t] = Math.min(Math.max(w, lo[j]), hi[j]);
      }
    }
    this._wRecBase = wRec;
  }

  private _wRecBase!: Float32Array;

  private buildW_rec(res: SpiraDeploy["reservoir"]) {
    const wRec = this._wRecBase.slice();
    const rows = new Int32Array(b64ToU8(res.w_rec_rows).buffer);
    const cols = new Int32Array(b64ToU8(res.w_rec_cols).buffer);
    const vals = new Float32Array(b64ToU8(res.w_rec_vals_b64).buffer);
    for (let i = 0; i < res.w_rec_n; i++) {
      wRec[rows[i] * this.nRes + cols[i]] = vals[i];
    }
    this.wRec = wRec;
  }

  // ---- hippocampal addressing: primary + leave-one-out ----
  private primaryAddr(k: Uint32Array): number {
    const z =
      (k[0] ^ rotl(k[1], 8) ^ rotl(k[2], 16) ^ rotl(k[3], 24)) >>> 0;
    return Number(splitmix64(BigInt(z)) % BigInt(this.amSize));
  }

  private looAddr(k: Uint32Array, l: number): number {
    const a: number[] = [];
    for (let i = 0; i < this.amM; i++) if (i !== l) a.push(k[i]);
    const z = (a[0] ^ rotl(a[1], 13) ^ rotl(a[2], 26)) >>> 0;
    return Number(
      splitmix64(BigInt(z) ^ BigInt(0x9e3779b9 + l * 0x85ebca6b)) % BigInt(this.amSize)
    );
  }

  private loadAmEntries(am: SpiraDeploy["am"]) {
    const packed = b64ToU8(am.packed_b64);
    const m = this.amM;
    const eb = am.entry_bytes;
    const dv = new DataView(packed.buffer);
    for (let i = 0; i < am.n_entries; i++) {
      const off = i * eb;
      const key: number[] = [];
      for (let l = 0; l < m; l++) key.push(dv.getUint32(off + l * 4, true));
      const val = dv.getBigUint64(off + 16, true);
      const count = packed[off + 24];
      const k = Uint32Array.from(key);
      // primary slot
      const p = this.primaryAddr(k);
      for (let l2 = 0; l2 < m; l2++) this.amKeys[p * m + l2] = k[l2];
      this.amVals[p] = val;
      this.amCounts[p] = count;
      this.amUsed[p] = 1;
      // leave-one-out completion slots
      for (let l = 0; l < m; l++) {
        const a = this.looAddr(k, l);
        if (!this.amUsed[a]) {
          for (let l2 = 0; l2 < m; l2++) this.amKeys[a * m + l2] = k[l2];
          this.amVals[a] = val;
          this.amCounts[a] = count;
          this.amUsed[a] = 1;
        }
      }
    }
  }

  private loadRegistry(reg: SpiraDeploy["registry"]) {
    const ids = new BigUint64Array(b64ToU8(reg.ids_b64).buffer);
    const lens = new Uint16Array(b64ToU8(reg.lens_b64).buffer);
    this.regBuffer = b64ToU8(reg.buffer_b64);
    let off = 0;
    for (let i = 0; i < reg.n; i++) {
      const id = ids[i];
      this.regIds.push(id);
      this.regLens.push(lens[i]);
      this.regOffsets.push(off);
      this.regIndexOf.set(id, i);
      off += lens[i];
    }
  }

  // ------------------------------------------------------------- reservoir
  resetCore() {
    this.v = new Float32Array(this.nRes);
    this.thr = new Float32Array(this.nRes).fill(this.baseThr);
    this.rhat = new Float32Array(this.nRes).fill(this.targetRate);
    this.refr = new Int8Array(this.nRes);
    this.x = new Float32Array(this.nRes);
    this.xPrev = new Float32Array(this.nRes);
    this.xMean = new Float32Array(this.nRes);
    this.s = new Float32Array(this.nRes);
    if (this.hier) {
      this.vG = new Float32Array(this.nG);
      this.thrG = new Float32Array(this.nG).fill(this.baseThr);
      this.rhatG = new Float32Array(this.nG).fill(this.targetRateG);
      this.xG = new Float32Array(this.nG);
      this.xMeanG = new Float32Array(this.nG);
      this.sG = new Float32Array(this.nG);
      this.winSpikes = new Float32Array(this.nRes);
      this.uSlow = new Float32Array(this.nOut);
      this.winCount = 0;
    }
  }

  /** Reset transients, KEEP learned homeostatic scalars (before generation). */
  resetDynamics() {
    const { thr, rhat, xMean } = this;
    let thrG: Float32Array | undefined, rhatG: Float32Array | undefined,
        xMeanG: Float32Array | undefined;
    if (this.hier) { thrG = this.thrG; rhatG = this.rhatG; xMeanG = this.xMeanG; }
    this.resetCore();
    this.thr = thr;
    this.rhat = rhat;
    this.xMean = xMean;
    if (this.hier && thrG && rhatG && xMeanG) {
      this.thrG = thrG;
      this.rhatG = rhatG;
      this.xMeanG = xMeanG;
    }
  }

  /** Advance one tick with the given active input neurons; returns trace x. */
  step(inputIdx: number[]): Float32Array {
    const n = this.nRes;
    const I = new Float32Array(n);
    for (const idx of inputIdx) {
      const base = idx * n;
      for (let j = 0; j < n; j++) I[j] += this.wIn[base + j];
    }
    for (let j = 0; j < n; j++) {
      const sj = this.s[j];
      if (sj !== 0) {
        const base = j * n;
        for (let t = 0; t < n; t++) I[t] += this.wRec[base + t] * sj;
      }
    }
    const v = this.v,
      thr = this.thr,
      refr = this.refr,
      rhat = this.rhat;
    const spikes = new Float32Array(n);
    for (let j = 0; j < n; j++) {
      let vj = this.lam * v[j] + I[j];
      if (refr[j] !== 0) vj = 0;
      if (vj >= thr[j] && refr[j] === 0) {
        spikes[j] = 1;
        v[j] = 0;
      } else {
        v[j] = vj;
      }
      if (refr[j] > 0) refr[j]--;
      if (spikes[j]) refr[j] = this.refrac;
      rhat[j] += 0.02 * (spikes[j] - rhat[j]);
      thr[j] += this.etaHomeo * (rhat[j] - this.targetRate);
      if (thr[j] < this.thrMin) thr[j] = this.thrMin;
      else if (thr[j] > this.thrMax) thr[j] = this.thrMax;
    }
    this.s = spikes;
    if (this.hier) {
      const ws = this.winSpikes;
      for (let j = 0; j < n; j++) ws[j] += spikes[j];
    }
    this.xPrev = this.x;
    const x = this.x;
    for (let j = 0; j < n; j++) x[j] = this.tauTrace * x[j] + spikes[j];
    // slow EMA of the trace (homeostatic scaling analogue; readout centering)
    const xm = this.xMean;
    for (let j = 0; j < n; j++) xm[j] += 0.001 * (x[j] - xm[j]);
    return x;
  }

  /**
   * Advance the slow area one slow step (one window of `win` fast ticks).
   * Input: the pooled fast spike volley, sign-preserved through Dale's law
   * of the presynaptic fast neuron. Output: the cached slow readout uSlow,
   * used as a variance-matched top-down prior with gain beta.
   */
  stepG() {
    const nG = this.nG;
    const I = new Float32Array(nG);
    const inv = 1 / this.win;
    for (let j = 0; j < this.nRes; j++) {
      const c = this.winSpikes[j] * inv;
      if (c === 0) continue;
      const base = j * nG;
      for (let g = 0; g < nG; g++) I[g] += c * this.wInG[base + g];
    }
    for (let j = 0; j < nG; j++) {
      const sj = this.sG[j];
      if (sj !== 0) {
        const base = j * nG;
        for (let g = 0; g < nG; g++) I[g] += this.wRecG[base + g] * sj;
      }
    }
    const v = this.vG, thr = this.thrG, rhat = this.rhatG, x = this.xG,
          xm = this.xMeanG;
    const spikes = new Float32Array(nG);
    for (let j = 0; j < nG; j++) {
      const vj = this.lamG * v[j] + I[j];
      if (vj >= thr[j]) {
        spikes[j] = 1;
        v[j] = 0;
      } else {
        v[j] = vj;
      }
      rhat[j] += 0.05 * (spikes[j] - rhat[j]);
      thr[j] += this.etaHomeoG * (rhat[j] - this.targetRateG);
      if (thr[j] < this.thrMin) thr[j] = this.thrMin;
      else if (thr[j] > this.thrMax) thr[j] = this.thrMax;
      x[j] = this.tauTraceG * x[j] + spikes[j];
      xm[j] += 0.001 * (x[j] - xm[j]);
    }
    this.sG = spikes;
    this.winSpikes.fill(0);
    // cached slow readout (int8-quantized, per-column scale)
    const nOut = this.nOut;
    const raw = new Float32Array(nOut);
    for (let j = 0; j < nG; j++) {
      const xj = x[j] - xm[j];
      if (xj === 0) continue;
      const base = j * nOut;
      for (let c = 0; c < nOut; c++) raw[c] += xj * this.wOutGQ[base + c];
    }
    const u = this.uSlow;
    for (let c = 0; c < nOut; c++) u[c] = raw[c] * this.wOutGScale[c];
  }

  /** Feed one token: fast step (+ slow step at window boundaries). */
  feed(pid: bigint): Float32Array {
    const x = this.step(inputNeurons(pid, this.dIn, this.kIn));
    if (this.hier) {
      this.winCount++;
      if (this.winCount % this.win === 0) this.stepG();
    }
    return x;
  }

  /** Behavioural scores: fast readout + variance-matched top-down prior. */
  scores(x: Float32Array): Float32Array {
    const nRes = this.nRes,
      nOut = this.nOut;
    const raw = new Float32Array(nOut);
    for (let j = 0; j < nRes; j++) {
      const xj = x[j] - this.xMean[j];
      if (xj === 0) continue;
      const base = j * nOut;
      for (let c = 0; c < nOut; c++) raw[c] += xj * this.wOutQ[base + c];
    }
    const u = new Float32Array(nOut);
    if (this.hier) {
      const bs = this.beta * this.fusionScale;
      for (let c = 0; c < nOut; c++) u[c] = raw[c] * this.wOutScale[c] + bs * this.uSlow[c];
    } else {
      for (let c = 0; c < nOut; c++) u[c] = raw[c] * this.wOutScale[c];
    }
    return u;
  }

  // ------------------------------------------------------------ hippocampus
  private writeAm(laneKeys: Uint32Array, val: bigint) {
    const m = this.amM;
    // primary slot: exact-context capacity
    const p = this.primaryAddr(laneKeys);
    if (!this.amUsed[p]) {
      for (let l = 0; l < m; l++) this.amKeys[p * m + l] = laneKeys[l];
      this.amVals[p] = val;
      this.amCounts[p] = 1;
      this.amUsed[p] = 1;
    } else {
      let same = true;
      for (let l = 0; l < m; l++)
        if (this.amKeys[p * m + l] !== laneKeys[l]) {
          same = false;
          break;
        }
      if (same) {
        if (this.amVals[p] === val) {
          if (this.amCounts[p] < 255) this.amCounts[p]++;
        } else {
          if (this.amCounts[p] > 1) this.amCounts[p]--;
          else {
            this.amVals[p] = val;
            this.amCounts[p] = 1;
          }
        }
      } else if (this.amCounts[p] <= 1) {
        for (let l = 0; l < m; l++) this.amKeys[p * m + l] = laneKeys[l];
        this.amVals[p] = val;
        this.amCounts[p] = 1;
      }
    }
    // leave-one-out completion slots
    for (let l = 0; l < m; l++) {
      const a = this.looAddr(laneKeys, l);
      if (!this.amUsed[a]) {
        for (let l2 = 0; l2 < m; l2++) this.amKeys[a * m + l2] = laneKeys[l2];
        this.amVals[a] = val;
        this.amCounts[a] = 1;
        this.amUsed[a] = 1;
      }
    }
  }

  /** Query the hippocampus: (patchId, weight) or null. Two-tier voting. */
  queryAm(laneKeys: Uint32Array): { val: bigint; weight: number } | null {
    const m = this.amM;
    const votes = new Map<bigint, number>();
    const maxMatch = new Map<bigint, number>();
    // primary probe: exact full-key retrieval
    const p = this.primaryAddr(laneKeys);
    if (this.amUsed[p]) {
      let same = true;
      for (let l = 0; l < m; l++)
        if (this.amKeys[p * m + l] !== laneKeys[l]) {
          same = false;
          break;
        }
      if (same) {
        const v = this.amVals[p];
        const w = m * (1 + Math.log1p(this.amCounts[p]));
        votes.set(v, (votes.get(v) ?? 0) + w);
        maxMatch.set(v, m);
      }
    }
    // leave-one-out probes: corrupted-cue completion
    for (let l = 0; l < m; l++) {
      const a = this.looAddr(laneKeys, l);
      if (!this.amUsed[a]) continue;
      let match = 0;
      for (let l2 = 0; l2 < m; l2++)
        if (this.amKeys[a * m + l2] === laneKeys[l2]) match++;
      if (match >= this.amMinLanes) {
        const v = this.amVals[a];
        const w = match * (1 + Math.log1p(this.amCounts[a]));
        votes.set(v, (votes.get(v) ?? 0) + w);
        if (match > (maxMatch.get(v) ?? 0)) maxMatch.set(v, match);
      }
    }
    if (!votes.size) return null;
    let top = 0;
    for (const mm of maxMatch.values()) if (mm > top) top = mm;
    let bestV: bigint | null = null;
    let bestW = 0;
    for (const [v, w] of votes) {
      const eff = top === m ? (maxMatch.get(v) === m ? w : -1) : w;
      if (eff >= 0 && eff > bestW) {
        bestW = eff;
        bestV = v;
      }
    }
    if (bestV === null) return null;
    return { val: bestV, weight: bestW };
  }

  private laneKeysOf(ids: bigint[]): Uint32Array {
    const m = this.amM;
    const out = new Uint32Array(m);
    for (let l = 0; l < m; l++) out[l] = laneKey32(ids[ids.length - m + l], l);
    return out;
  }

  // -------------------------------------------------------------- registry
  bytesOfId(id: bigint): Uint8Array | null {
    const idx = this.regIndexOf.get(id);
    if (idx === undefined) return null;
    const off = this.regOffsets[idx];
    return this.regBuffer.slice(off, off + this.regLens[idx]);
  }

  private registerPatches(patches: Patch[]): number {
    let added = 0;
    for (const p of patches) {
      if (!this.regIndexOf.has(p.id)) {
        const idx = this.regIds.length;
        this.regIds.push(p.id);
        const off = this.regOffsets.length
          ? this.regOffsets[this.regOffsets.length - 1] +
            this.regLens[this.regLens.length - 1]
          : 0;
        this.regOffsets.push(off);
        this.regLens.push(p.bytes.length);
        this.regIndexOf.set(p.id, idx);
        const newBuf = new Uint8Array(this.regBuffer.length + p.bytes.length);
        newBuf.set(this.regBuffer);
        newBuf.set(p.bytes, this.regBuffer.length);
        this.regBuffer = newBuf;
        added++;
      }
    }
    return added;
  }

  // ------------------------------------------------------------- one-shot
  /** ONE-SHOT episodic teaching: one pass, no gradients, no repetition. */
  teach(text: string | Uint8Array): TeachResult {
    const t0 = performance.now();
    const data = typeof text === "string" ? new TextEncoder().encode(text) : text;
    const patches = this.ept.segment(data);
    const ids = patches.map((p) => p.id);
    const newPatches = this.registerPatches(patches);
    const m = this.amM;
    let associations = 0;
    for (let r = 0; r + m < ids.length; r++) {
      const laneKeys = new Uint32Array(m);
      for (let l = 0; l < m; l++) laneKeys[l] = laneKey32(ids[r + l], l);
      this.writeAm(laneKeys, ids[r + m]);
      associations++;
    }
    return { tokens: ids.length, newPatches, associations, ms: performance.now() - t0 };
  }

  /**
   * Retrieve the memorized continuation of a context (multi-patch replay),
   * with optional corruption: `corruptSlots` entries of the 4-slot context
   * are replaced by random registered patches before querying (pattern
   * completion test). Repeats the completion `steps` times, sliding the
   * context window.
   */
  recall(context: string, corruptSlots = 0, rngSeed = 11, steps = 8): {
    continuation: string;
    weight: number;
  } | null {
    const data = new TextEncoder().encode(context);
    const patches = this.ept.segment(data);
    if (patches.length < this.amM) return null;
    const ids = patches.map((p) => p.id);
    const enc = new TextEncoder();
    // initial (optionally corrupted) lane keys
    const laneKeys = this.laneKeysOf(ids);
    if (corruptSlots > 0 && this.regIds.length > 8) {
      let seed = rngSeed >>> 0;
      const rand = () => {
        seed ^= seed << 13; seed >>>= 0;
        seed ^= seed >>> 17;
        seed ^= seed << 5; seed >>>= 0;
        return seed / 4294967296;
      };
      const slots = new Set<number>();
      let guard = 0;
      while (slots.size < Math.min(corruptSlots, this.amM) && guard++ < 32) {
        slots.add(Math.floor(rand() * this.amM) % this.amM);
      }
      for (const slot of slots) {
        const randIdx = Math.floor(rand() * this.regIds.length) % this.regIds.length;
        laneKeys[slot] = laneKey32(this.regIds[randIdx], slot);
      }
    }
    const hit = this.queryAm(laneKeys);
    if (!hit) return null;
    const first = this.bytesOfId(hit.val);
    if (!first) return null;
    // greedy multi-patch replay through the memory
    let out: Uint8Array = first;
    let weight = hit.weight;
    const ctx = ids.slice();
    ctx.push(hit.val);
    for (let s = 1; s < steps; s++) {
      const lk = this.laneKeysOf(ctx);
      const h = this.queryAm(lk);
      if (!h) break;
      const b = this.bytesOfId(h.val);
      if (!b) break;
      const merged = new Uint8Array(out.length + b.length);
      merged.set(out);
      merged.set(b, out.length);
      out = merged;
      weight = Math.min(weight, h.weight);
      ctx.push(h.val);
    }
    return { continuation: new TextDecoder().decode(out), weight };
  }

  // ------------------------------------------------------------ generation
  /** Closed-loop generation: predicted patches are fed back as spikes. */
  generate(prompt: string, opts?: {
    nTokens?: number; temperature?: number; topK?: number; useAm?: boolean;
    seed?: number; onToken?: (step: number, bytes: Uint8Array) => void;
  }): string {
    const nTokens = opts?.nTokens ?? 120;
    const temperature = opts?.temperature ?? this.temperature;
    const topK = opts?.topK ?? this.topK;
    const useAm = opts?.useAm ?? true;
    let seed = opts?.seed ?? 12345;
    const rand = () => {
      seed ^= seed << 13; seed >>>= 0;
      seed ^= seed >>> 17;
      seed ^= seed << 5; seed >>>= 0;
      return seed / 4294967296;
    };
    this.resetDynamics();
    const data = new TextEncoder().encode(prompt);
    const patches = this.ept.segment(data);
    if (!patches.length) return prompt;
    const ids: bigint[] = patches.map((p) => p.id);
    const out: number[] = Array.from(data);
    for (const pid of ids) {
      this.feed(pid);
    }
    for (let step = 0; step < nTokens; step++) {
      const x = this.x;
      let u = this.scores(x);
      let bypassId: bigint | null = null;
      if (useAm && ids.length >= this.amM) {
        const hit = this.queryAm(this.laneKeysOf(ids));
        if (hit) {
          const col = this.colOf.get(hit.val);
          if (col !== undefined) {
            u = u.slice();
            u[col] += this.betaAm * hit.weight;
          } else if (hit.weight >= this.amOverride && this.bytesOfId(hit.val)) {
            bypassId = hit.val;
          }
        }
      }
      let chosenId: bigint;
      if (bypassId !== null) {
        chosenId = bypassId;
      } else {
        // top-k selection by a single partial scan (k is small)
        const k = Math.max(1, Math.min(topK, this.vocabSize));
        const top: number[] = [];
        const topScores: number[] = [];
        for (let c = 0; c < this.vocabSize; c++) {
          const sc = u[c];
          if (topScores.length < k) {
            top.push(c); topScores.push(sc);
            // keep the small list sorted descending
            for (let i = topScores.length - 1; i > 0 && topScores[i] > topScores[i - 1]; i--) {
              const ts = topScores[i]; topScores[i] = topScores[i - 1]; topScores[i - 1] = ts;
              const ti = top[i]; top[i] = top[i - 1]; top[i - 1] = ti;
            }
          } else if (sc > topScores[k - 1]) {
            top[k - 1] = c; topScores[k - 1] = sc;
            for (let i = k - 1; i > 0 && topScores[i] > topScores[i - 1]; i--) {
              const ts = topScores[i]; topScores[i] = topScores[i - 1]; topScores[i - 1] = ts;
              const ti = top[i]; top[i] = top[i - 1]; top[i - 1] = ti;
            }
          }
        }
        const maxS = topScores[0];
        const z = top.map((c) => (u[c] - maxS) / Math.max(temperature, 1e-6));
        const exps = z.map(Math.exp);
        const tot = exps.reduce((a, b) => a + b, 0);
        let r = rand() * tot;
        let pick = top[0];
        for (let i = 0; i < top.length; i++) {
          r -= exps[i];
          if (r <= 0) {
            pick = top[i];
            break;
          }
        }
        chosenId = this.vocabIds[pick];
      }
      const b = this.bytesOfId(chosenId);
      if (!b) break;
      for (const byte of b) out.push(byte);
      ids.push(chosenId);
      this.feed(chosenId);
      opts?.onToken?.(step, b);
    }
    return new TextDecoder().decode(new Uint8Array(out));
  }

  // ------------------------------------------------------------- utilities
  surpriseProfile(text: string): { surprise: Float32Array; patches: Patch[] } {
    const data = new TextEncoder().encode(text);
    return { surprise: this.ept.surprise(data), patches: this.ept.segment(data) };
  }

  stats() {
    const slowAmortized = this.hier
      ? (this.nG * (this.vocabSize + 1) + this.nRes * this.nG + this.nG * this.nG) / this.win
      : 0;
    const macs =
      this.kIn * this.nRes + this.nRes * this.nRes + this.nRes * (this.vocabSize + 1) +
      slowAmortized;
    return {
      neurons: this.nRes,
      inputNeurons: this.dIn,
      readout: this.vocabSize + 1,
      hierarchy: this.hier
        ? { neurons: this.nG, window: this.win, beta: this.beta }
        : null,
      synapses: {
        input: this.dIn * 8,
        recurrent: this.nRes * this.nRes,
        readout: this.nRes * (this.vocabSize + 1),
        slowReadout: this.hier ? this.nG * (this.vocabSize + 1) : 0,
      },
      macsPerToken: macs,
      registry: this.regIds.length,
      amEntries: this.amUsed.reduce((a, b) => a + b, 0),
    };
  }
}

let _model: SpiraModel | null = null;

/** Lazy singleton loader (the artifact is ~29 MB — fetched once). */
export async function loadSpira(): Promise<SpiraModel> {
  if (_model) return _model;
  const res = await fetch("/spira_deploy.json");
  if (!res.ok) throw new Error(`artifact fetch failed: ${res.status}`);
  const art = (await res.json()) as SpiraDeploy;
  _model = new SpiraModel(art);
  return _model;
}