#!/usr/bin/env python3
"""Quick smoke test of every SPIRA component on a tiny slice of the corpus.

Validates: hashing determinism, EPT segmentation, reservoir dynamics,
three-factor learning signal, hippocampal write/query/corruption, generation.
"""
import sys
import time

import numpy as np

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

from spira import EPT, Hippocampus, SpikingCore, SPIRA, SPIRAConfig
from spira.hashing import Mulberry32, fnv1a64, input_neurons, lane_key32, splitmix64

data = open("/home/z/my-project/data/train.bin", "rb").read()
print(f"corpus: {len(data)/1e6:.2f} MB")
small = data[:120_000]

# ---------------------------------------------------------------- hashing
h1 = fnv1a64(b"hello world")
h2 = fnv1a64(b"hello world")
assert h1 == h2 and h1 != fnv1a64(b"hello worle")
assert splitmix64(42) == splitmix64(42)
r = Mulberry32(7)
vals = [r.next_float() for _ in range(5)]
r2 = Mulberry32(7)
assert vals == [r2.next_float() for _ in range(5)]
neurons = input_neurons(h1, 1024, 12)
assert len(set(neurons)) >= 10 and all(0 <= n < 1024 for n in neurons)
print(f"[ok] hashing: fnv1a64('hello world') = {h1:#x}, 12 neurons, deterministic PRNG")

# -------------------------------------------------------------------- EPT
ept = EPT()
ept.fit(small[:80_000])
budget = ept.calibrate(small[80_000:100_000], target_len=4.5)
patches = ept.segment(small[:40_000])
st = ept.stats(patches)
print(f"[ok] EPT: budget={budget:.2f} bits -> avg patch len {st['avg_len']:.2f} B "
      f"(p50={st['p50_len']:.0f} p90={st['p90_len']:.0f} max={st['max_len']})")
assert 3.0 < st["avg_len"] < 7.0
# show a sample segmentation
sample = small[:220].decode("utf-8", "replace")
seg = ept.segment(small[:220])
lens = [len(p) for p in seg]
print(f"[ok] EPT sample: {' | '.join(repr(p.decode('utf-8','replace')) for p in seg[:24])}")
ids = [fnv1a64(p) for p in patches]
assert len(set(ids)) > 100

# ------------------------------------------------------------- hippocampus
am = Hippocampus(m=4, size=1 << 16)
ids_small = ids[:3000]
lane = am.lane_keys_stream(np.array(ids_small, dtype=np.uint64))
# scalar vs vectorized lane keys must agree
for i in (0, 17, 999):
    for l in range(4):
        assert int(lane[i, l]) == lane_key32(ids_small[i], l), (i, l)
print("[ok] hippocampus: scalar and vectorized lane keys agree")

ctx_lane_list = []
m = 4
n_ctx = len(ids_small) - m + 1
ctx = np.empty((n_ctx, m), dtype=np.uint32)
for li in range(m):
    ctx[:, li] = lane[li : li + n_ctx, li]
vals_idx = np.arange(m, m + n_ctx)   # continuation token index for each context
am.write_stream(ctx, vals_idx)
# context tokens 10..13 -> continuation = token 14
q, w = am.query(ctx[10])
assert q == 14 and w > 0, (q, w)
print(f"[ok] hippocampus: one-shot write + exact recall (query->next={q}, weight={w:.2f})")
# corruption: replace 1 of 4 context patches (slot 2 = token 12)
corrupt = ctx[10].copy()
corrupt[2] = lane_key32(ids_small[500], 2)
q2, w2 = am.query(corrupt)
print(f"[ok] hippocampus: corrupted-context query -> {q2} (true=14), weight={w2:.2f}")

# ------------------------------------------------------------- spiking core
core = SpikingCore(n_out=128, seed=1)
acc = []
T = 4000
rng = np.random.default_rng(0)
toy_ids = np.array(ids_small[:T], dtype=np.uint64)
neurons_stream = SPIRA.input_neurons_stream(toy_ids, 1024, 12)
# toy target: next id mod 128 -> learnable map? no: use the actual next patch (hard)
targets = (toy_ids % np.uint64(127)).astype(np.int64)[1:]
core.reset()
t0 = time.perf_counter()
res = core.train_pass(neurons_stream[: T - 1], targets, log_every=2000,
                      log_fn=lambda d: print(f"    {d}"))
dt = time.perf_counter() - t0
print(f"[ok] spiking core: {res['ticks']} ticks in {dt:.1f}s "
      f"({res['ticks_per_sec']:.0f} t/s), final acc_ema={res['acc_ema_final']:.3f}")
assert res["ticks_per_sec"] > 200

# ------------------------------------------------------- end-to-end (tiny)
cfg = SPIRAConfig(snn_tokens=6_000, am_tokens=6_000, epochs=1, vocab_size=1024)
model = SPIRA.build(small, cfg, log=lambda m: print("   ", m))
model.prepare()
model._patches_train = None
hist = model.train(log_fn=lambda m: print("   ", m))
print(f"[ok] end-to-end train: {hist['epochs'][0]['accuracy']:.3f} acc, am={hist['am']}")

# generation
gen = model.generate(b"The ", n_tokens=40, seed=1)
print(f"[ok] generation sample: {gen[:200]!r}")

# eval incl. one-shot teaching on held-out text
oneshot = open("/home/z/my-project/data/oneshot.bin", "rb").read()[:30_000]
ev0 = model.eval_stream(oneshot, max_positions=3000)
ev1 = model.eval_stream(oneshot, teach_am_first=True, max_positions=3000)
print(f"[ok] eval on held-out book: before-teach fused={ev0['fused_acc']:.3f}; "
      f"after ONE-SHOT teach fused={ev1['fused_acc']:.3f} (am={ev1['am_acc']:.3f})")
ev2 = model.eval_stream(oneshot, teach_am_first=True, corrupt_lanes=1, max_positions=3000)
print(f"[ok] eval with 1/4 corrupted context lanes: fused={ev2['fused_acc']:.3f}")

print("\nALL SMOKE TESTS PASSED")