"""
Quick smoke test: 3 steps of training to verify everything works.
"""
import sys; sys.path.insert(0, '.')
import torch
from pathlib import Path
from angstrom_nano import AngstromNanoConfig, AngstromNanoForCausalLM
from angstrom_nano.tokenizer import AngstromNanoTokenizer

text = Path("data/corpus.txt").read_text(encoding="utf-8")
tok = AngstromNanoTokenizer.train_bpe(["data/corpus.txt"], vocab_size=4096)
ids = torch.tensor(tok.encode(text, add_bos=True, add_eos=True), dtype=torch.long)
print(f"Vocab: {len(tok)}, Tokens: {len(ids):,}")

cfg = AngstromNanoConfig(
    vocab_size=len(tok), hidden_size=192, intermediate_size=512,
    num_hidden_layers=6, num_attention_heads=6, num_key_value_heads=3,
    head_dim=32, num_local_experts=4, num_experts_per_tok=2,
    max_position_embeddings=256, sliding_window=64,
    scoring_func="sigmoid", use_qk_norm=True, use_routing_bias=True,
    tie_word_embeddings=True,
)
model = AngstromNanoForCausalLM(cfg)
p = sum(p.numel() for p in model.parameters())
print(f"Model: {p:,} params")

opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
for step in range(1, 4):
    i = torch.randint(0, len(ids) - 64 - 1, (1,)).item()
    x = ids[i:i+64].unsqueeze(0)
    y = ids[i+1:i+65].unsqueeze(0)
    out = model(x, labels=y, output_router_logits=True)
    out["loss"].backward()
    opt.step()
    opt.zero_grad()
    loss = out["loss"].item()
    aux = out["aux_loss"].item()
    print(f"step {step}: loss={loss:.4f} aux={aux:.6f}")

print("SMOKE TEST PASSED")