"""
Train angstrom_nano with BPE tokenizer on the full merged corpus.
Laptop-friendly medium config: ~8.4M params, fits in <500 MB RAM.
Trains on CPU in ~1-2 hours.
"""
import sys; sys.path.insert(0, '.')
import math, time, torch
from pathlib import Path
from angstrom_nano import AngstromNanoConfig, AngstromNanoForCausalLM
from angstrom_nano.tokenizer import AngstromNanoTokenizer

torch.manual_seed(42)

# ------------------------------------------------------------------
# 1. Train BPE tokenizer on the merged corpus
# ------------------------------------------------------------------
corpus_path = Path("data/corpus.txt")
if not corpus_path.exists():
    print("corpus.txt not found — run prepare_data.py first")
    sys.exit(1)

text = corpus_path.read_text(encoding="utf-8")
print(f"Data: {len(text):,} chars ({len(text)/1e6:.1f} MB)")

tok = AngstromNanoTokenizer.train_bpe(
    [str(corpus_path)],
    vocab_size=4096,
    out_path="checkpoints/tokenizer.json",
)
vocab_size = len(tok)
print(f"BPE vocab: {vocab_size} tokens")

ids = torch.tensor(tok.encode(text, add_bos=True, add_eos=True), dtype=torch.long)
print(f"Tokenized: {len(ids):,} tokens ({len(ids)/1e6:.1f}M)")

# ------------------------------------------------------------------
# 2. Medium config — fits laptop CPU comfortably
# ------------------------------------------------------------------
cfg = AngstromNanoConfig(
    vocab_size=vocab_size,
    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())
m = sum(p.numel() * p.element_size() for p in model.parameters()) / 1e6
print(f"Model: {p:,} params ({m:.2f} MB in {'FP32' if m > 0 else 'N/A'})")

# ------------------------------------------------------------------
# 3. Training
# ------------------------------------------------------------------
seq_len = 64
lr = 3e-3
steps = 5000
print_every = 500
opt = torch.optim.AdamW(model.parameters(), lr=lr)
t0 = time.time()

for step in range(1, steps + 1):
    i = torch.randint(0, len(ids) - seq_len - 1, (1,)).item()
    x = ids[i : i + seq_len].unsqueeze(0)
    y = ids[i + 1 : i + seq_len + 1].unsqueeze(0)

    out = model(x, labels=y, output_router_logits=True)
    loss = out["loss"]
    opt.zero_grad()
    loss.backward()
    opt.step()

    if step % print_every == 0 or step == 1:
        print(f"  step {step:>4d}  loss={loss.item():.4f}  ppl={math.exp(loss.item()):.2f}  "
              f"aux={out['aux_loss'].item():.6f}  {time.time()-t0:.0f}s")

# ------------------------------------------------------------------
# 4. Save .safetensors
# ------------------------------------------------------------------
out_dir = Path("checkpoints")
out_dir.mkdir(exist_ok=True)
path = out_dir / "medium_model.safetensors"

import json
from safetensors.torch import save_file

cfg_path = out_dir / "medium_config.json"
cfg_path.write_text(json.dumps(
    {"vocab_size": vocab_size, "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_type": "angstrom_nano"}, indent=2))

sd = model.state_dict()
if sd["lm_head.weight"].data_ptr() == sd["model.embed_tokens.weight"].data_ptr():
    sd.pop("lm_head.weight")
save_file({k: v.contiguous() for k, v in sd.items()}, str(path))
print(f"\nSaved: {path}  ({path.stat().st_size / 1e6:.2f} MB)")
print(f"Saved: {cfg_path}")

# ------------------------------------------------------------------
# 5. Verify: reload and generate
# ------------------------------------------------------------------
from safetensors.torch import load_file as safe_load
tok = AngstromNanoTokenizer.from_bpe_file("checkpoints/tokenizer.json")
d = json.loads(cfg_path.read_text())
cfg2 = AngstromNanoConfig(**d)
model2 = AngstromNanoForCausalLM(cfg2)
sd = safe_load(str(path))
if "lm_head.weight" not in sd:
    model2.lm_head.weight = model2.model.embed_tokens.weight
model2.load_state_dict(sd, strict=False)
model2.eval()

prompt = "def fibonacci"
ids_p = torch.tensor([tok.encode(prompt, add_bos=True, add_eos=False)])
out_ids = model2.generate(ids_p, max_new_tokens=20, temperature=0.8)
gen = tok.decode(out_ids[0].tolist(), skip_special_tokens=True)
clean = ''.join(c if 32 <= ord(c) <= 126 or c in '\n\t' else ' ' for c in gen)
print(f"\nPrompt: {prompt}")
print(f"Gen   : {clean}")
print("TRAINING COMPLETE")