#!/usr/bin/env python3
"""
StoryLM-10M: ~10.5M param LLaMA-style language model trained on TinyStories.

Architecture:
  d_model=256, n_heads=8, n_kv_heads=8 (MHA), n_layers=8
  SwiGLU FFN (4x), RoPE, RMSNorm, vocab 8192, tied embed/head, ctx 512
  Total: ~10.5M learnable parameters

Training:
  ~2B tokens from TinyStories, AdamW 3e-4, cosine + warmup
  batch 64 (effective), seq 512, ~6,100 steps

Usage:
  python3 train_story10m.py --stage prepare   # download + tokenize
  python3 train_story10m.py --stage train     # train
  python3 train_story10m.py --stage eval      # evaluate
  python3 train_story10m.py --stage all
"""
import os, sys, math, json, time, argparse, glob
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

# ============================================================================
# Config
# ============================================================================
VOCAB = 8192
D_MODEL = 256
N_HEADS = 8
N_KV_HEADS = 8
N_LAYERS = 8
HEAD_DIM = D_MODEL // N_HEADS  # 32
KV_DIM = N_KV_HEADS * HEAD_DIM  # 256
FFN_DIM = D_MODEL * 4  # 1024
SEQ_LEN = 512
BATCH = 4
LR = 3e-4
WARMUP_STEPS = 500
TOTAL_STEPS = 24400
CHECKPOINT_EVERY = 500
OUT_DIR = "/work/story10m"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
TOKENIZER_PATH = "/work/tokenizer.json"

os.makedirs(OUT_DIR, exist_ok=True)

# ============================================================================
# Architecture
# ============================================================================
class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
    def forward(self, x):
        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x * norm).type_as(x) * self.weight

def precompute_rope(dim, max_seq, base=10000.0):
    freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
    t = torch.arange(max_seq).float()
    angles = torch.outer(t, freqs)
    return torch.stack([angles.cos(), angles.sin()], dim=-1)

def apply_rope(x, cos, sin):
    # x: (B, H, T, D)
    B, H, T, D = x.shape
    x1 = x[..., 0::2]
    x2 = x[..., 1::2]
    c = cos[:T].unsqueeze(0).unsqueeze(0)  # (1, 1, T, D/2)
    s = sin[:T].unsqueeze(0).unsqueeze(0)
    out = torch.empty_like(x)
    out[..., 0::2] = x1 * c - x2 * s
    out[..., 1::2] = x1 * s + x2 * c
    return out

class Attention(nn.Module):
    def __init__(self, d_model, n_heads, n_kv_heads, head_dim):
        super().__init__()
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.head_dim = head_dim
        self.qkv = nn.Linear(d_model, (n_heads + 2 * n_kv_heads) * head_dim, bias=False)
        self.o = nn.Linear(n_heads * head_dim, d_model, bias=False)

    def forward(self, x, cos, sin):
        B, T, _ = x.shape
        qkv = self.qkv(x)
        q, k, v = qkv.split([self.n_heads * self.head_dim,
                             self.n_kv_heads * self.head_dim,
                             self.n_kv_heads * self.head_dim], dim=-1)
        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        if self.n_kv_heads != self.n_heads:
            rep = self.n_heads // self.n_kv_heads
            k = k.repeat_interleave(rep, dim=1)
            v = v.repeat_interleave(rep, dim=1)
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)
        attn = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        attn = attn.transpose(1, 2).contiguous().view(B, T, -1)
        return self.o(attn)

class SwiGLU(nn.Module):
    def __init__(self, d_model, ffn_dim):
        super().__init__()
        self.gate = nn.Linear(d_model, ffn_dim, bias=False)
        self.up = nn.Linear(d_model, ffn_dim, bias=False)
        self.down = nn.Linear(ffn_dim, d_model, bias=False)
    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))

class Block(nn.Module):
    def __init__(self, d_model, n_heads, n_kv_heads, head_dim, ffn_dim):
        super().__init__()
        self.norm1 = RMSNorm(d_model)
        self.attn = Attention(d_model, n_heads, n_kv_heads, head_dim)
        self.norm2 = RMSNorm(d_model)
        self.ffn = SwiGLU(d_model, ffn_dim)
    def forward(self, x, cos, sin):
        x = x + self.attn(self.norm1(x), cos, sin)
        x = x + self.ffn(self.norm2(x))
        return x

class StoryLM(nn.Module):
    def __init__(self):
        super().__init__()
        self.embed = nn.Embedding(VOCAB, D_MODEL)
        self.layers = nn.ModuleList([
            Block(D_MODEL, N_HEADS, N_KV_HEADS, HEAD_DIM, FFN_DIM)
            for _ in range(N_LAYERS)
        ])
        self.norm = RMSNorm(D_MODEL)
        self.lm_head = nn.Linear(D_MODEL, VOCAB, bias=False)
        self.lm_head.weight = self.embed.weight  # tied
        rope = precompute_rope(HEAD_DIM, SEQ_LEN)
        self.register_buffer('rope_cos', rope[..., 0])
        self.register_buffer('rope_sin', rope[..., 1])

    def forward(self, x, targets=None):
        h = self.embed(x)
        for block in self.layers:
            h = block(h, self.rope_cos, self.rope_sin)
        h = self.norm(h)
        logits = self.lm_head(h)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, VOCAB), targets.view(-1))
        return logits, loss

# ============================================================================
# Data
# ============================================================================
def prepare_data():
    """Download TinyStories and tokenize."""
    print("Loading tokenizer...")
    from tokenizers import Tokenizer
    tok = Tokenizer.from_file(TOKENIZER_PATH)
    
    print("Downloading TinyStories...")
    import urllib.request
    stories_path = os.path.join(OUT_DIR, "tinystories.txt")
    if not os.path.exists(stories_path):
        url = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesv2_2M.txt"
        urllib.request.urlretrieve(url, stories_path)
    
    with open(stories_path, 'r') as f:
        stories = f.readlines()
    print(f"  {len(stories)} stories")
    
    print("Tokenizing...")
    all_tokens = []
    for i, story in enumerate(stories):
        story = story.strip()
        if len(story) < 20:
            continue
        ids = tok.encode(story, add_special_tokens=False).ids
        all_tokens.extend(ids)
        if (i+1) % 100000 == 0:
            print(f"  {i+1}/{len(stories)} stories, {len(all_tokens)/1e6:.1f}M tokens so far")
    
    tokens = np.array(all_tokens, dtype=np.uint16)
    np.save(os.path.join(OUT_DIR, "tokens.npy"), tokens)
    print(f"  Total: {len(tokens)/1e6:.1f}M tokens")
    print(f"  Saved to {OUT_DIR}/tokens.npy")

def load_data():
    tokens = np.load(os.path.join(OUT_DIR, "tokens.npy"))
    print(f"Loaded {len(tokens)/1e6:.1f}M tokens")
    # 99% train, 1% val
    split = int(len(tokens) * 0.99)
    train = tokens[:split]
    val = tokens[split:]
    print(f"Train: {len(train)/1e6:.1f}M, Val: {len(val)/1e6:.1f}M")
    return train, val

def get_batch(data, batch_size, seq_len):
    ix = torch.randint(len(data) - seq_len - 1, (batch_size,))
    x = torch.stack([torch.from_numpy((data[i:i+seq_len]).astype(np.int64)) for i in ix])
    y = torch.stack([torch.from_numpy((data[i+1:i+1+seq_len]).astype(np.int64)) for i in ix])
    return x.to(DEVICE), y.to(DEVICE)

# ============================================================================
# Training
# ============================================================================
def train():
    train_data, val_data = load_data()
    model = StoryLM().to(DEVICE)
    total_params = sum(p.numel() for p in model.parameters())
    print(f"Parameters: {total_params:,}")
    
    optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=0.0, betas=(0.9, 0.95))
    
    best_val = float('inf')
    t0 = time.time()
    
    for step in range(1, TOTAL_STEPS + 1):
        # LR schedule: warmup + cosine
        if step < WARMUP_STEPS:
            lr = LR * step / WARMUP_STEPS
        else:
            progress = (step - WARMUP_STEPS) / (TOTAL_STEPS - WARMUP_STEPS)
            lr = LR * 0.5 * (1 + math.cos(math.pi * progress))
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr
        
        x, y = get_batch(train_data, BATCH, SEQ_LEN)
        _, loss = model(x, y)
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        
        if step % 100 == 0:
            elapsed = time.time() - t0
            tok_per_s = (step * BATCH * SEQ_LEN) / elapsed
            print(f"step {step}/{TOTAL_STEPS} | loss {loss.item():.4f} | lr {lr:.2e} | {tok_per_s/1000:.1f}k tok/s | {elapsed:.0f}s", flush=True)
        
        if step % CHECKPOINT_EVERY == 0:
            # Eval
            model.eval()
            val_losses = []
            for _ in range(20):
                vx, vy = get_batch(val_data, BATCH, SEQ_LEN)
                _, vloss = model(vx, vy)
                val_losses.append(vloss.item())
            val_loss = sum(val_losses) / len(val_losses)
            model.train()
            print(f"  val_loss: {val_loss:.4f}")
            
            ckpt_path = os.path.join(OUT_DIR, f"ckpt_{step}.pt")
            torch.save(model.state_dict(), ckpt_path)
            print(f"  checkpoint: {ckpt_path}")
            
            if val_loss < best_val:
                best_val = val_loss
                torch.save(model.state_dict(), os.path.join(OUT_DIR, "best.pt"))
                print(f"  new best: {val_loss:.4f}")
    
    # Save final
    torch.save(model.state_dict(), os.path.join(OUT_DIR, "final.pt"))
    print(f"\nTraining done. Best val_loss: {best_val:.4f}")
    print(f"Final model: {OUT_DIR}/final.pt")

# ============================================================================
# Evaluation
# ============================================================================
def evaluate():
    model = StoryLM().to(DEVICE)
    ckpt = torch.load(os.path.join(OUT_DIR, "best.pt"), map_location=DEVICE)
    model.load_state_dict(ckpt)
    model.eval()
    
    _, val_data = load_data()
    
    # Perplexity
    val_losses = []
    for _ in range(50):
        vx, vy = get_batch(val_data, BATCH, SEQ_LEN)
        _, vloss = model(vx, vy)
        val_losses.append(vloss.item())
    avg_loss = sum(val_losses) / len(val_losses)
    ppl = math.exp(avg_loss)
    print(f"Val perplexity: {ppl:.2f} (loss {avg_loss:.4f})")
    
    # Generation samples
    from tokenizers import Tokenizer
    tok = Tokenizer.from_file(TOKENIZER_PATH)
    prompts = [
        "Once upon a time, there was a little",
        "The cat sat on the",
        "In the beginning, the world was",
        "A small robot named",
        "Every morning, the sun",
    ]
    print("\n=== Generation Samples ===")
    for prompt in prompts:
        ids = tok.encode(prompt, add_special_tokens=False).ids
        x = torch.tensor([ids], dtype=torch.long, device=DEVICE)
        with torch.no_grad():
            for _ in range(30):
                out, _ = model(x)
                next_id = out[0, -1].argmax().item()
                x = torch.cat([x, torch.tensor([[next_id]], device=DEVICE)], dim=1)
        full_ids = ids + [next_id]
        text = tok.decode(full_ids, skip_special_tokens=True)
        print(f"Prompt: {prompt}")
        print(f"Output: {text}")
        print()

# ============================================================================
# Main
# ============================================================================
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--stage", choices=["prepare", "train", "eval", "all"], default="all")
    args = parser.parse_args()
    
    if args.stage in ("prepare", "all"):
        prepare_data()
    if args.stage in ("train", "all"):
        train()
    if args.stage in ("eval", "all"):
        evaluate()

if __name__ == "__main__":
    main()