#!/usr/bin/env python3
"""Probe: sharp sparse input drive -> token-discriminative spike patterns."""
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=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, in_fanout, w_rec_scale, rec_fanout, inh_frac)
    (3.0, 8, 0.35, 48, 0.50),
    (2.5, 12, 0.35, 48, 0.50),
    (4.0, 8, 0.35, 48, 0.50),
    (3.0, 8, 0.25, 24, 0.50),
    (5.0, 6, 0.35, 48, 0.50),
]

def pattern_discrimination(core, neurons, n=400, stride=17):
    """Mean Jaccard overlap of spike patterns for tokens at distant positions
    with DIFFERENT patch ids; low overlap = high token information."""
    core.reset()
    pats = []
    id_seen = []
    for t in range(0, n * stride, stride):
        core.step(neurons[t])
        pats.append(set(np.nonzero(core.s)[0].tolist()))
        id_seen.append(int(neurons[t][0]))  # proxy: first neuron differs
    overlaps = []
    for i in range(len(pats) - 1):
        for j in range(i + 1, min(i + 5, len(pats))):
            a, b = pats[i], pats[j]
            if a or b:
                overlaps.append(len(a & b) / max(len(a | b), 1))
    return float(np.mean(overlaps))

for w_in, fan, w_rec, recf, inh in variants:
    core = SpikingCore(n_out=cfg.vocab_size + 1, seed=cfg.seed,
                       w_in_scale=w_in, in_fanout=fan,
                       w_rec_scale=w_rec, rec_fanout=recf, inh_frac=inh,
                       eta_r=0.002)
    core.reset()
    sp = 0
    for t in range(2000):
        core.step(neurons[t])
        sp += int(core.s.sum())
    rate = sp / 2000 / core.n_res
    jac = pattern_discrimination(core, neurons)
    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} fan={fan} w_rec={w_rec:.2f} recf={recf} | rate={rate:.3f} "
          f"jac={jac:.3f} overall={res['accuracy']:.4f} in-vocab~={res['accuracy']/cov:.4f}")