"""CPU smoke test for jlens_v3 (faithful J-lens machinery). Random tiny model."""
import sys, os
sys.path.insert(0, '.')
sys.path.insert(0, 'src')
import torch, numpy as np
from model import GPT, GPTConfig
import jlens_v3

torch.manual_seed(0)

cfg = GPTConfig(n_layer=3, n_head=4, n_embd=16, block_size=32, bias=False,
                vocab_size=65, dropout=0.0)
model = GPT(cfg).eval()

data = np.random.randint(0, 65, 2000).astype(np.uint16)
batches = jlens_v3.make_batches(data, 32, 4, 3, 'cpu')

J = jlens_v3.compute_faithful_jlens(model, 1, batches, 'cpu', chunk=8)
print("faithful vecs shape:", J.shape, "finite:", torch.isfinite(J).all().item())
print("row norms (first 5):", [round(v, 4) for v in J.norm(dim=1)[:5].tolist()])

pn = jlens_v3.compute_proxy_norms(model, 1, batches, 'cpu', 65)
print("proxy norms finite:", all(np.isfinite(v) for v in pn.values()))

fn = {k: J[k].norm().item() for k in range(65)}
freq = np.bincount(data, minlength=65).astype(float)
freq = freq / freq.sum() * 100
r_p = np.corrcoef(np.array([pn[k] for k in range(65)]), freq)[0, 1]
r_f = np.corrcoef(np.array([fn[k] for k in range(65)]), freq)[0, 1]
print(f"random-init sanity: proxy r={r_p:+.3f} faithful r={r_f:+.3f} (should be ~0, model is random)")
assert J.shape == (65, 16) and torch.isfinite(J).all()
print("SMOKE TEST PASSED")