#!/usr/bin/env python3
"""Definitive mechanics test: can the three-factor readout overfit a
deterministic toy mapping? Separates mechanism bugs from data statistics."""
import sys

import numpy as np

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

rng = np.random.default_rng(7)
N_TOKENS = 64
d_in, k_in = 1024, 12

# deterministic toy: token i is always followed by token (i*7+3) % 64
patterns = [np.asarray(input_neurons(i * 0x9E3779B97F4A7C15 % (1 << 64), d_in, k_in), dtype=np.int32)
            for i in range(N_TOKENS)]
next_of = [(i * 7 + 3) % N_TOKENS for i in range(N_TOKENS)]

core = SpikingCore(n_out=N_TOKENS, seed=20261022,
                   w_in_scale=3.0, in_fanout=8, w_rec_scale=0.25, rec_fanout=24,
                   inh_frac=0.5, eta_pos=0.02, eta_r=0.0)

T = 20_000
seq = [0]
for _ in range(T):
    seq.append(next_of[seq[-1]])
correct_ema = None
for ep in range(1):
    core.reset()
    n_ok = 0
    for t in range(T):
        x = core.step(patterns[seq[t]])
        u = (x - core.x_mean) @ core.W_out
        pred = int(np.argmax(u))
        ok = pred == next_of[seq[t]]
        n_ok += ok
        r = 1.0 if ok else 0.0
        core.baseline = 0.995 * core.baseline + 0.005 * r
        d = r - core.baseline
        x_eff = x - core.x_mean
        W = core.W_out
        W[:, next_of[seq[t]]] += np.float32(core.eta_pos) * x_eff
        W[:, pred] += np.float32(core.eta_r * d) * x_eff
        correct_ema = r if correct_ema is None else 0.995 * correct_ema + 0.005 * r
        if (t + 1) % 5000 == 0:
            print(f"  tick {t+1}: acc={n_ok/(t+1):.3f} ema={correct_ema:.3f}")

# how stable is the state for the same token? (same-token x overlap)
core.reset()
xs = {}
for t in range(4000):
    x = core.step(patterns[seq[t]])
    xs.setdefault(seq[t], []).append(x.copy())
tok = max(xs, key=lambda k: len(xs[k]))
if len(xs[tok]) >= 2:
    a, b = xs[tok][0], xs[tok][-1]
    cos = float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-9))
    print(f"same-token cosine similarity: {cos:.3f} (n={len(xs[tok])})")
diff = [k for k in xs if k != tok][0]
c2 = float(np.dot(xs[tok][0], xs[diff][0]) / (np.linalg.norm(xs[tok][0]) * np.linalg.norm(xs[diff][0]) + 1e-9))
print(f"diff-token cosine similarity: {c2:.3f}")