"""
Train a small GPT model on TinyStories using PyTorch (CPU).
~10M parameters, character-level, trains in ~30-60 min on CPU.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import os
import pickle
import requests
from tqdm import tqdm
import math

# ============================================================
# Config
# ============================================================
CONFIG = {
    "block_size": 128,
    "batch_size": 64,
    "n_embd": 256,
    "n_head": 8,
    "n_layer": 6,
    "dropout": 0.1,
    "learning_rate": 3e-4,
    "max_steps": 3000,
    "eval_interval": 300,
    "eval_iters": 100,
    "warmup_steps": 200,
    "weight_decay": 0.01,
}

DATA_DIR = "/home/kongyaolang/tinystories/data"
MODEL_DIR = "/home/kongyaolang/tinystories/model"
os.makedirs(DATA_DIR, exist_ok=True)
os.makedirs(MODEL_DIR, exist_ok=True)

# ============================================================
# Data
# ============================================================
def download_tinystories():
    """Download TinyStories dataset."""
    data_path = os.path.join(DATA_DIR, "TinyStories.txt")
    if os.path.exists(data_path) and os.path.getsize(data_path) > 100_000:
        print(f"Data exists: {data_path} ({os.path.getsize(data_path):,} bytes)")
        return data_path

    print("Downloading TinyStories...")
    # Try mirror first, then official
    urls = [
        "https://hf-mirror.com/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-train.txt",
        "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-train.txt",
    ]
    resp = None
    for url in urls:
        try:
            print(f"  Trying: {url}")
            resp = requests.get(url, timeout=60, stream=True)
            if resp.status_code == 200:
                break
        except Exception as e:
            print(f"  Failed: {e}")
            continue
    if resp is None or resp.status_code != 200:
        raise RuntimeError("Failed to download TinyStories from all mirrors")
    total = int(resp.headers.get("content-length", 0))
    
    with open(data_path, "wb") as f:
        with tqdm(total=total, unit="B", unit_scale=True, desc="Downloading") as pbar:
            for chunk in resp.iter_content(chunk_size=8192):
                f.write(chunk)
                pbar.update(len(chunk))
    
    print(f"Downloaded to {data_path}")
    return data_path


def prepare_data(data_path, max_chars=10_000_000):
    """Load text, build character-level vocab, create train/val splits."""
    with open(data_path, "r", encoding="utf-8") as f:
        text = f.read(max_chars)
    
    # Build vocab: printable ASCII + common chars
    chars = sorted(set(text))
    # Keep chars that appear at least a few times
    char_counts = {}
    for c in text:
        char_counts[c] = char_counts.get(c, 0) + 1
    
    # Filter to chars appearing > 5 times, keep ASCII first
    ascii_chars = [c for c in chars if ord(c) < 128 and char_counts[c] > 5]
    non_ascii = [c for c in chars if ord(c) >= 128 and char_counts[c] > 50]
    vocab_chars = ascii_chars + non_ascii
    
    stoi = {ch: i for i, ch in enumerate(vocab_chars)}
    itos = {i: ch for i, ch in enumerate(vocab_chars)}
    vocab_size = len(vocab_chars)
    
    # Encode
    valid_chars = set(vocab_chars)
    filtered = "".join(c for c in text if c in valid_chars)
    data = torch.tensor([stoi[c] for c in filtered], dtype=torch.long)
    
    # Split
    n = int(0.9 * len(data))
    train_data = data[:n]
    val_data = data[n:]
    
    print(f"Vocab size: {vocab_size}, Train tokens: {len(train_data):,}, Val tokens: {len(val_data):,}")
    return train_data, val_data, stoi, itos, vocab_size


def get_batch(data, block_size, batch_size):
    """Get a random batch."""
    ix = torch.randint(0, len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i+block_size] for i in ix])
    y = torch.stack([data[i+1:i+block_size+1] for i in ix])
    return x, y


# ============================================================
# Model: Mini-GPT
# ============================================================
class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        
        self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)
        self.proj = nn.Linear(n_embd, n_embd, bias=False)
        self.dropout = nn.Dropout(dropout)
        
        # Causal mask
        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size))
                             .view(1, 1, block_size, block_size))
    
    def forward(self, x):
        B, T, C = x.shape
        qkv = self.qkv(x)  # (B, T, 3*C)
        q, k, v = qkv.chunk(3, dim=-1)
        
        # Reshape to multi-head
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)  # (B, nh, T, hs)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        
        # Attention
        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
        att = F.softmax(att, dim=-1)
        att = self.dropout(att)
        
        y = att @ v  # (B, nh, T, hs)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.proj(y)
        y = self.dropout(y)
        return y


class MLP(nn.Module):
    def __init__(self, n_embd, dropout):
        super().__init__()
        self.fc1 = nn.Linear(n_embd, 4 * n_embd)
        self.gelu = nn.GELU()
        self.fc2 = nn.Linear(4 * n_embd, n_embd)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.gelu(x)
        x = self.fc2(x)
        x = self.dropout(x)
        return x


class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout)
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = MLP(n_embd, dropout)
    
    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x


class MiniGPT(nn.Module):
    def __init__(self, vocab_size, n_embd, n_head, n_layer, block_size, dropout):
        super().__init__()
        self.block_size = block_size
        
        self.tok_emb = nn.Embedding(vocab_size, n_embd)
        self.pos_emb = nn.Embedding(block_size, n_embd)
        self.drop = nn.Dropout(dropout)
        
        self.blocks = nn.Sequential(*[
            Block(n_embd, n_head, block_size, dropout) for _ in range(n_layer)
        ])
        
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
        
        # Weight tying
        self.tok_emb.weight = self.lm_head.weight
        
        # Init
        self.apply(self._init_weights)
    
    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
    
    def forward(self, idx):
        B, T = idx.shape
        assert T <= self.block_size
        
        tok = self.tok_emb(idx)  # (B, T, C)
        pos = torch.arange(T, device=idx.device)
        pos_emb = self.pos_emb(pos)  # (T, C)
        x = self.drop(tok + pos_emb)
        
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)  # (B, T, vocab_size)
        return logits
    
    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=0.8):
        self.eval()
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.block_size:]
            logits = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, next_token], dim=1)
        self.train()
        return idx


# ============================================================
# Training
# ============================================================
@torch.no_grad()
def estimate_loss(model, data, block_size, batch_size, eval_iters):
    model.eval()
    losses = []
    for _ in range(eval_iters):
        x, y = get_batch(data, block_size, batch_size)
        logits = model(x)
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1))
        losses.append(loss.item())
    model.train()
    return sum(losses) / len(losses)


def train(model, train_data, val_data, config, stoi, itos):
    print(f"\n{'='*60}")
    n_params = sum(p.numel() for p in model.parameters())
    print(f"MiniGPT: {n_params/1e6:.1f}M parameters")
    print(f"  Layers: {config['n_layer']}, Dim: {config['n_embd']}, Heads: {config['n_head']}")
    print(f"  Steps: {config['max_steps']}, Batch: {config['batch_size']}, LR: {config['learning_rate']}")
    print(f"{'='*60}\n")
    
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=config["learning_rate"],
        weight_decay=config["weight_decay"],
        betas=(0.9, 0.95),
    )
    
    losses = []
    best_val_loss = float("inf")
    
    pbar = tqdm(range(1, config["max_steps"] + 1))
    for step in pbar:
        # Learning rate warmup + cosine decay
        if step <= config["warmup_steps"]:
            lr = config["learning_rate"] * step / config["warmup_steps"]
        else:
            progress = (step - config["warmup_steps"]) / (config["max_steps"] - config["warmup_steps"])
            lr = config["learning_rate"] * 0.5 * (1 + math.cos(math.pi * progress))
        
        for param_group in optimizer.param_groups:
            param_group["lr"] = lr
        
        # Get batch
        x, y = get_batch(train_data, config["block_size"], config["batch_size"])
        
        # Forward
        logits = model(x)
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1))
        
        # Backward
        optimizer.zero_grad()
        loss.backward()
        
        # Gradient clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        
        optimizer.step()
        
        losses.append(loss.item())
        
        # Eval
        if step % config["eval_interval"] == 0 or step == config["max_steps"]:
            train_loss = sum(losses[-100:]) / min(100, len(losses))
            val_loss = estimate_loss(model, val_data, config["block_size"], 
                                     config["batch_size"], config["eval_iters"])
            pbar.set_description(f"train: {train_loss:.4f}, val: {val_loss:.4f}")
            
            # Save best
            if val_loss < best_val_loss:
                best_val_loss = val_loss
                torch.save({
                    "model": model.state_dict(),
                    "config": config,
                    "stoi": stoi,
                    "itos": itos,
                    "step": step,
                    "val_loss": val_loss,
                }, os.path.join(MODEL_DIR, "best_model.pt"))
        
        # Generate sample
        if step % 500 == 0:
            model.eval()
            context = torch.zeros((1, 1), dtype=torch.long)
            gen = model.generate(context, max_new_tokens=150, temperature=0.8)
            text = "".join(itos.get(t.item(), "?") for t in gen[0])
            model.train()
            print(f"\n--- Sample at step {step} (lr={lr:.2e}) ---")
            print(text[:250])
            print("---\n")
    
    return losses


# ============================================================
# Main
# ============================================================
if __name__ == "__main__":
    print("=" * 60)
    print("MiniGPT: Training on TinyStories (PyTorch CPU)")
    print("=" * 60)
    
    # Download data
    data_path = download_tinystories()
    
    # Prepare data
    train_data, val_data, stoi, itos, vocab_size = prepare_data(data_path, max_chars=8_000_000)
    CONFIG["vocab_size"] = vocab_size
    
    # Create model
    model = MiniGPT(
        vocab_size=vocab_size,
        n_embd=CONFIG["n_embd"],
        n_head=CONFIG["n_head"],
        n_layer=CONFIG["n_layer"],
        block_size=CONFIG["block_size"],
        dropout=CONFIG["dropout"],
    )
    
    n_params = sum(p.numel() for p in model.parameters())
    print(f"Model: {n_params/1e6:.1f}M parameters")
    
    # Train
    losses = train(model, train_data, val_data, CONFIG, stoi, itos)
    
    # Save final model
    torch.save({
        "model": model.state_dict(),
        "config": CONFIG,
        "stoi": stoi,
        "itos": itos,
    }, os.path.join(MODEL_DIR, "final_model.pt"))
    print(f"\nModel saved to {MODEL_DIR}")
    
    # Final generation
    print("\n" + "=" * 60)
    print("Final Generation Samples")
    print("=" * 60)
    
    model.eval()
    for temp in [0.5, 0.7, 0.9]:
        print(f"\n--- Temperature: {temp} ---")
        context = torch.zeros((1, 1), dtype=torch.long)
        gen = model.generate(context, max_new_tokens=200, temperature=temp)
        text = "".join(itos.get(t.item(), "?") for t in gen[0])
        print(text[:300])
        print()