#!/usr/bin/env python3
"""Probe reservoir scaling (E/I balance) + resulting in-vocab accuracy."""
import sys
import time

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=40_000, am_tokens=1000, epochs=1)
model = SPIRA.build(data, cfg, log=lambda m: None)
ids = model._ids_train
N = 40_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}")

variants = [
    # (w_in_scale, w_rec_scale, inh_frac, lam)
    (1.5, 0.35, 0.50, 0.90),
    (1.2, 0.30, 0.50, 0.90),
    (1.5, 0.25, 0.50, 0.90),
    (1.0, 0.35, 0.50, 0.90),
    (1.5, 0.35, 0.50, 0.70),
    (2.0, 0.40, 0.50, 0.90),
]
for w_in, w_rec, inh, lam in variants:
    core = SpikingCore(n_out=cfg.vocab_size + 1, seed=cfg.seed,
                       w_in_scale=w_in, w_rec_scale=w_rec, inh_frac=inh, lam=lam,
                       eta_r=0.002)
    # firing rate over first 2000 ticks
    core.reset()
    sp = 0
    for t in range(2000):
        core.step(neurons[t])
        sp += int(core.s.sum())
    rate = sp / 2000 / core.n_res
    res = core.train_pass(neurons[: N - 1], cols[1:N], log_every=20_000, log_fn=lambda d: None)
    print(f"w_in={w_in:.1f} w_rec={w_rec:.2f} inh={inh:.2f} lam={lam:.2f} | "
          f"rate={rate:.3f} overall={res['accuracy']:.4f} in-vocab~={res['accuracy']/cov:.4f}")