"""Small LLM1: ~5M parameter Llama-style model trained from scratch on TinyStories."""
import math, os, time
import torch
from datasets import load_dataset
from tokenizers import Tokenizer, models, pre_tokenizers, decoders, trainers, processors
from transformers import LlamaConfig, LlamaForCausalLM, PreTrainedTokenizerFast

OUT = "small-llm1"
N_STORIES = 400_000
VOCAB = 4096
SEQ = 256
BATCH = 64
STEPS = 6000
LR = 2e-3
torch.manual_seed(0)

# --- data ---
print("loading data...")
ds = load_dataset("roneneldan/TinyStories", split="train", streaming=True)
texts = [ex["text"].strip() for _, ex in zip(range(N_STORIES), ds) if ex["text"].strip()]
print(f"{len(texts)} stories")

# --- tokenizer (byte-level BPE) ---
tok = Tokenizer(models.BPE())
tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
tok.decoder = decoders.ByteLevel()
trainer = trainers.BpeTrainer(vocab_size=VOCAB, special_tokens=["<|endoftext|>"],
                              initial_alphabet=pre_tokenizers.ByteLevel.alphabet())
tok.train_from_iterator(texts[:100_000], trainer)
tok.post_processor = processors.ByteLevel(trim_offsets=False)
hf_tok = PreTrainedTokenizerFast(tokenizer_object=tok, bos_token="<|endoftext|>",
                                 eos_token="<|endoftext|>", unk_token="<|endoftext|>",
                                 pad_token="<|endoftext|>")
eos = hf_tok.eos_token_id

ids = []
for i in range(0, len(texts), 10_000):
    for seq in hf_tok(texts[i:i + 10_000])["input_ids"]:
        ids.extend(seq + [eos])
data = torch.tensor(ids, dtype=torch.long)
n_val = SEQ * 2000
train_data, val_data = data[:-n_val], data[-n_val:]
print(f"{len(data)/1e6:.1f}M tokens")

# --- model ---
config = LlamaConfig(vocab_size=VOCAB, hidden_size=256, intermediate_size=512,
                     num_hidden_layers=6, num_attention_heads=8, num_key_value_heads=8,
                     max_position_embeddings=SEQ, tie_word_embeddings=True,
                     bos_token_id=eos, eos_token_id=eos, pad_token_id=eos)
dev = "cuda"
model = LlamaForCausalLM(config).to(dev)
n_params = sum(p.numel() for p in model.parameters())
print(f"params: {n_params/1e6:.2f}M")

opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=0.1, betas=(0.9, 0.95), fused=True)
warmup = 200
sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1, (s + 1) / warmup) *
                                          (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(s, STEPS) / STEPS))))

def batch(src):
    ix = torch.randint(len(src) - SEQ - 1, (BATCH,))
    x = torch.stack([src[i:i + SEQ] for i in ix]).to(dev, non_blocking=True)
    return x

@torch.no_grad()
def evaluate():
    model.eval()
    losses = []
    for i in range(0, len(val_data) - SEQ, SEQ * BATCH):
        x = torch.stack([val_data[j:j + SEQ] for j in range(i, min(i + SEQ * BATCH, len(val_data) - SEQ), SEQ)]).to(dev)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            losses.append(model(input_ids=x, labels=x).loss.item())
    model.train()
    return sum(losses) / len(losses)

t0 = time.time()
model.train()
for step in range(STEPS):
    x = batch(train_data)
    with torch.autocast("cuda", dtype=torch.bfloat16):
        loss = model(input_ids=x, labels=x).loss
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step(); sched.step(); opt.zero_grad(set_to_none=True)
    if step % 500 == 0 or step == STEPS - 1:
        print(f"step {step} loss {loss.item():.3f} lr {sched.get_last_lr()[0]:.2e} {time.time()-t0:.0f}s", flush=True)

val_loss = evaluate()
print(f"val loss {val_loss:.3f}  ppl {math.exp(val_loss):.2f}  time {time.time()-t0:.0f}s")

model.generation_config.pad_token_id = eos
model.save_pretrained(OUT)
hf_tok.save_pretrained(OUT)
with open(os.path.join(OUT, "train_stats.txt"), "w") as f:
    f.write(f"params={n_params}\ntokens={len(data)}\nsteps={STEPS}\nval_loss={val_loss:.4f}\n")
print("saved to", OUT)