#!/usr/bin/env python3
"""Quick hyperparameter probe for the three-factor rule (40k ticks each)."""
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.hashing import fnv1a64
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])
inv_mask = (cols[1:] != cfg.vocab_size)[: N - 1]
print(f"coverage: {inv_mask.mean():.3f}")

variants = [
    ("eta_r=0.010 d-full", 0.010, -1.0),
    ("eta_r=0.001 d-full", 0.001, -1.0),
    ("eta_r=0.010 d>=-0.25", 0.010, -0.25),
    ("eta_r=0.000 pure-inst", 0.000, 0.0),
]
for name, eta_r, d_floor in variants:
    core = SpikingCore(n_out=cfg.vocab_size + 1, seed=cfg.seed, eta_r=eta_r)
    # monkeypatch learn with d floor
    orig_learn = core.learn
    def learn(pred, target, x, core=core, floor=d_floor, orig=orig_learn):
        # replicate learn() but with d floor on the chosen-column update
        correct = 1 if pred == target else 0
        reward = float(correct)
        core.baseline = core.baseline_momentum * core.baseline + (1 - core.baseline_momentum) * reward
        d = reward - core.baseline
        W_out = core.W_out
        if target != core.n_out - 1:
            W_out[:, target] += np.float32(core.eta_pos) * x
        d_eff = d if d > 0 else max(d, floor)
        W_out[:, pred] += np.float32(core.eta_r * d_eff) * x
        np.clip(W_out[:, target], -core.w_clip, core.w_clip, out=W_out[:, target])
        np.clip(W_out[:, pred], -core.w_clip, core.w_clip, out=W_out[:, pred])
        if core.s.any():
            spiked = np.nonzero(core.s)[0]
            W_rec = core.W_rec
            delta = np.float32(core.eta_stdp * d)
            W_rec[:, spiked] += delta * core.x_prev[:, None]
            np.clip(W_rec[:, spiked], core._rec_lo[:, None], core._rec_hi[:, None],
                    out=W_rec[:, spiked])
        return correct, reward, d
    core.learn = learn
    t0 = time.perf_counter()
    res = core.train_pass(neurons[: N - 1], cols[1:N], log_every=10_000, log_fn=lambda d: None)
    # in-vocab accuracy measured inline: recompute via final pass? use overall + coverage
    overall = res["accuracy"]
    print(f"{name:24s} overall={overall:.4f} in-vocab~={overall/max(inv_mask.mean(),1e-9):.4f} "
          f"({res['ticks_per_sec']:.0f} t/s)")