from __future__ import annotations


import math
from pathlib import Path


import torch
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer


ROOT = Path(__file__).resolve().parent
EXPECTED_PARAMS = 96_130_040


print("1/6 tokenizer load + UTF-8 round trip")
tok = AutoTokenizer.from_pretrained(
    ROOT,
    trust_remote_code=True,
    use_fast=False,
)
sample = "TransWave — raw bytes ✓"
ids = tok(sample, return_tensors="pt")["input_ids"]
decoded = tok.decode(ids[0].tolist())
assert decoded == sample, (decoded, sample)
assert tok.vocab_size == 256


print("2/6 config")
cfg = AutoConfig.from_pretrained(ROOT, trust_remote_code=True)
assert cfg.architecture_id == "causal_transwave_blocklifting_meyer_byte_lm_g4_v2"
assert cfg.trained_context_length == 20480
assert cfg.vocab_size == 256


print("3/6 trained checkpoint load")
model = AutoModelForCausalLM.from_pretrained(
    ROOT,
    trust_remote_code=True,
).eval()
n = sum(p.numel() for p in model.transwave.parameters())
assert n == EXPECTED_PARAMS, n


print("4/6 exact finite NLL")
torch.manual_seed(20261012)
x = torch.randint(0, 256, (1, 32), dtype=torch.long)
with torch.no_grad():
    direct = model.transwave.nll(x, checkpoint_layers=False)
    wrapped = model(input_ids=x, labels=x).loss
assert torch.isfinite(direct)
assert torch.isfinite(wrapped)
assert abs(float(direct - wrapped)) < 2e-5
print("NLL:", float(wrapped), "PPL:", math.exp(float(wrapped)))


print("5/6 future-perturbation causality")
a = torch.randint(0, 256, (1, 32), dtype=torch.long)
b = a.clone()
b[:, 16:] = torch.randint(0, 256, b[:, 16:].shape, dtype=torch.long)
with torch.no_grad():
    ca, _ = model.transwave.contexts(a, checkpoint_layers=False)
    cb, _ = model.transwave.contexts(b, checkpoint_layers=False)
diff = float((ca[:, :3, :] - cb[:, :3, :]).abs().max())
assert diff <= 2e-4, diff
print("future->past max diff:", diff)


print("6/6 reference generation")
prompt = tok("abcdefgh", return_tensors="pt")["input_ids"]
with torch.no_grad():
    y = model.generate_bytes(prompt, max_new_tokens=2, do_sample=False)
assert y.shape == (1, 10)
assert int(y.min()) >= 0 and int(y.max()) <= 255


print("SMOKE TEST: PASS")