#!/usr/bin/env python3
"""Post-training sanity: hierarchy ablation + generation quality."""
import sys
import time

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

from spira import SPIRA

BASE = "/home/z/my-project"
m = SPIRA.load(f"{BASE}/download/spira/artifacts/spira_model.npz")
print(f"vocab={m.cfg.vocab_size} hier={m.core.hier} beta={m.core.beta:.3f} "
      f"thr_G={float(m.core.thr_G.mean()):.3f} rhat_G={float(m.core.rhat_G.mean()):.3f}")

m._ids_train = None
train = open(f"{BASE}/data/train.bin", "rb").read()
test = open(f"{BASE}/data/test.bin", "rb").read()

import numpy as np
from spira.hashing import fnv1a64
ids = np.fromiter((fnv1a64(p) for p in m.ept.segment(train)), dtype=np.uint64, count=len(m.ept.segment(train)))
m._ids_train = ids

for name, data, pos in [("train-region", train[10_000_000:10_200_000], 5000),
                        ("test", test[:200_000], 8000)]:
    ev = m.eval_stream(data, max_positions=pos)
    print(f"{name}: cov={ev['vocab_coverage']:.3f} fast={ev['snn_acc']:.4f} "
          f"hier={ev['hier_acc']:.4f} am={ev['am_acc']:.4f} fused={ev['fused_acc']:.4f} "
          f"| invocab: fast={ev['snn_acc_invocab']:.4f} hier={ev['hier_acc_invocab']:.4f} "
          f"am={ev['am_acc_invocab']:.4f}")

t0 = time.perf_counter()
g = m.generate(b"The old man looked at the ", n_tokens=120, temperature=0.7, seed=42)
dt = time.perf_counter() - t0
print(f"\ngeneration ({len(g)} B, {len(g)/max(dt,1e-9):.0f} B/s):")
print(g.decode("utf-8", "replace")[:400])
g2 = m.generate(b"def process(data):\n    ", n_tokens=100, temperature=0.7, seed=7)
print("\ncode:", g2.decode("utf-8", "replace")[:300])