#!/usr/bin/env python3
"""A/B test: ORIGINAL reservoir params (rich recurrent states) with honest
vocab-masked metrics, with/without trace centering, various eta_r."""
import sys

import numpy as np

sys.path.insert(0, "/home/z/my-project/download/spira")
from spira import SPIRA, SPIRAConfig
from spira.spiking import SpikingCore

data = open("/home/z/my-project/data/train.bin", "rb").read()
cfg = SPIRAConfig(snn_tokens=60_000, am_tokens=1000, epochs=1)
model = SPIRA.build(data, cfg, log=lambda m: None)
ids = model._ids_train
N = 60_000
neurons = SPIRA.input_neurons_stream(ids[:N], cfg.d_in, cfg.k_in)
cols = model.cols_of(ids[: N + 1])
cov = (cols[1:] != cfg.vocab_size)[: N - 1].mean()
print(f"coverage: {cov:.3f}")

# ORIGINAL params that reached EMA 0.52-0.82 (with the UNK shortcut metric)
variants = [
    ("sharp+delta e=.70", dict(w_in_scale=3.0, in_fanout=8, w_rec_scale=0.25, rec_fanout=24, inh_frac=0.5, eta_delta=0.70), True),
    ("sharp+delta e=1.0", dict(w_in_scale=3.0, in_fanout=8, w_rec_scale=0.25, rec_fanout=24, inh_frac=0.5, eta_delta=1.0), True),
    ("sharp+delta e=1.5", dict(w_in_scale=3.0, in_fanout=8, w_rec_scale=0.25, rec_fanout=24, inh_frac=0.5, eta_delta=1.5), True),
    ("sharp+delta e=1.0 tau.3", dict(w_in_scale=3.0, in_fanout=8, w_rec_scale=0.25, rec_fanout=24, inh_frac=0.5, eta_delta=1.0, tau_trace=0.35), True),
]
for name, kw, center in variants:
    core = SpikingCore(n_out=cfg.vocab_size + 1, seed=cfg.seed, **kw)
    if not center:
        # disable centering by pinning x_mean to zero
        class NoCenter:
            pass
        core._center = False
        orig_step = core.step
        import types
        # monkeypatch: after step, reset x_mean to 0
        def step_noc(self, input_idx, _orig=orig_step):
            x = _orig(input_idx)
            self.x_mean = np.zeros_like(self.x_mean)
            return x
        core.step = types.MethodType(step_noc, core)
    res = core.train_pass(neurons[: N - 1], cols[1:N], log_every=30_000, log_fn=lambda d: None)
    print(f"{name:28s} overall={res['accuracy']:.4f} in-vocab~={res['accuracy']/cov:.4f} "
          f"ema={res['acc_ema_final']:.4f}")