#!/usr/bin/env python3
"""Decisive roundtrip test: fresh 60k-tick train -> measure -> save -> load -> measure."""
import sys

import numpy as np

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

data = open("/home/z/my-project/data/train.bin", "rb").read()
cfg = SPIRAConfig(snn_tokens=60_000, am_tokens=2000, 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()

res = model.core.train_pass(neurons[: N - 1], cols[1:N], log_every=30_000, log_fn=lambda d: None)
print(f"IN-PROCESS training: overall={res['accuracy']:.4f} in-vocab~={res['accuracy']/cov:.4f}")

# in-process eval on the SAME slice (no learning)
ev = model.eval_stream(data[:260_000], max_positions=12_000)
print(f"IN-PROCESS eval (same slice): snn_invocab={ev['snn_acc_invocab']:.4f} "
      f"cov={ev['vocab_coverage']:.3f}")

model.save("/tmp/roundtrip.npz")
model2 = SPIRA.load("/tmp/roundtrip.npz")
ev2 = model2.eval_stream(data[:260_000], max_positions=12_000)
print(f"POST-LOAD eval (same slice):  snn_invocab={ev2['snn_acc_invocab']:.4f}")

# compare weight matrices directly
same = np.allclose(model.core.W_out, model2.core.W_out, atol=1e-6)
print(f"W_out identical after roundtrip: {same}")
print(f"W_out col norm stats: max={np.linalg.norm(model.core.W_out, axis=0).max():.4f} "
      f"mean={np.linalg.norm(model.core.W_out, axis=0).mean():.4f}")
print(f"threshold range: [{model2.core.thr.min():.3f}, {model2.core.thr.max():.3f}] "
      f"mean={model2.core.thr.mean():.3f}")
print(f"x_mean norm: {np.linalg.norm(model2.core.x_mean):.3f} vs in-proc {np.linalg.norm(model.core.x_mean):.3f}")