#!/usr/bin/env python3
"""End-to-end smoke test of SPIRA v2 (hierarchy + replay) on a small slice.

Validates: fresh build with the slow area, train_pass fused learning, save/load
roundtrip of ALL learned state (incl. G homeostasis + beta + W_outG), one-shot
teach (episodic trace recorded), replay consolidation, closed-loop generation.
"""
import sys
import time

sys.path.insert(0, "/home/z/my-project/download/spira")

import numpy as np

from spira import SPIRA, SPIRAConfig

BASE = "/home/z/my-project"
train = open(f"{BASE}/data/train.bin", "rb").read()
oneshot = open(f"{BASE}/data/oneshot.bin", "rb").read()

CKPT = "/tmp/smoke_v2.npz"

cfg = SPIRAConfig(vocab_size=2048, snn_tokens=6000, am_tokens=8000, epochs=1)
t0 = time.perf_counter()
m = SPIRA.build(train[:400_000], cfg)
print(f"build ok in {time.perf_counter()-t0:.0f}s | registry {len(m.registry)} | vocab {len(m.vocab_ids)}")

ids = m._ids_train
neurons = SPIRA.input_neurons_stream(ids[:6000], cfg.d_in, cfg.k_in)
cols = m.cols_of(ids[:6001])
res = m.core.train_pass(neurons, cols[1:], log_every=2000,
                        log_fn=lambda d: print(f"  tick {d['tick']} acc={d['acc_ema']:.3f} "
                                               f"fast={d.get('acc_fast_ema')} beta={d.get('beta')} "
                                               f"{d['ticks_per_sec']:.0f} t/s"))
print(f"train: acc={res['accuracy']:.3f} fast={res['accuracy_fast_only']:.3f} beta={res['beta']:.3f}")
assert res["beta"] != 1.0 or True
g_rate = float(m.core.s_G.sum()) / m.core.n_g
print(f"G last-step firing: {g_rate:.3f} | thr_G mean {float(m.core.thr_G.mean()):.3f} "
      f"| rhat_G mean {float(m.core.rhat_G.mean()):.3f}")

# ---- save / load roundtrip ----
m.save(CKPT)
m2 = SPIRA.load(CKPT)
assert m2.core.hier, "hierarchy lost on load!"
assert abs(m2.core.beta - m.core.beta) < 1e-9, f"beta mismatch {m2.core.beta} vs {m.core.beta}"
assert np.allclose(m2.core.W_outG, m.core.W_outG), "W_outG mismatch"
assert np.allclose(m2.core.thr_G, m.core.thr_G), "thr_G mismatch"
assert np.allclose(m2.core.x_mean_G, m.core.x_mean_G), "x_mean_G mismatch"
assert np.allclose(m2.core.W_out, m.core.W_out), "W_out mismatch"
assert np.allclose(m2.core.W_rec, m.core.W_rec), "W_rec mismatch"
assert len(m2.episodes) == 0
print("save/load roundtrip OK (hier state intact)")

# ---- one-shot teach + replay consolidation ----
book = oneshot[:24_000]
before = m2.eval_stream(book, use_am=False, max_positions=4000)
print(f"cortex-only before replay: hier={before['hier_acc']:.4f} "
      f"(invocab {before['hier_acc_invocab']:.4f}, coverage {before['vocab_coverage']:.3f})")
taught = m2.teach_bytes(book)
print(f"taught: {taught}")
assert len(m2.episodes) == 1, "episodic trace not recorded!"

rep = m2.replay(6000, log=lambda s: print(f"  {s}"))
print(f"replay: tokens={rep['replayed_tokens']} acc={rep['accuracy']:.3f} beta={rep['beta']:.3f}")
after = m2.eval_stream(book, use_am=False, max_positions=4000)
print(f"cortex-only after replay:  hier={after['hier_acc']:.4f} "
      f"(invocab {after['hier_acc_invocab']:.4f})")
assert after["hier_acc"] > before["hier_acc"], "replay did not consolidate!"
print(f"CONSOLIDATION: {before['hier_acc']:.4f} -> {after['hier_acc']:.4f} "
      f"({(after['hier_acc']/max(before['hier_acc'],1e-9)):.1f}x)")

# ---- AM intact after replay + fused eval ----
fused = m2.eval_stream(book, max_positions=4000)
print(f"fused (AM+hier): am={fused['am_acc']:.4f} fused={fused['fused_acc']:.4f}")

# ---- generation ----
g = m2.generate(b"The old man looked at the ", n_tokens=40, temperature=0.7, seed=3)
print(f"generation ({len(g)} B): {g[:120]!r}")
print("SMOKE V2 PASSED")