#!/usr/bin/env python3
"""Train a small language model on TinyStories.

TinyStories (Eldan & Li, 2023) is a dataset of simple children's stories
generated by GPT-3.5/GPT-4, specifically designed to test what small
language models can learn. Even models with <10M parameters can produce
coherent stories after training on this dataset.

This script downloads TinyStories, tokenizes with GPT-2 BPE, and trains
a configurable model architecture. It demonstrates that our LMT library
can train on real, meaningful text -- not just Shakespeare character-level.

Usage:
    python examples/train_tinystories.py
    python examples/train_tinystories.py --epochs 5 --device mps
    python examples/train_tinystories.py --arch llama --embed-dim 256

Dataset:
    https://huggingface.co/datasets/roneneldan/TinyStories
    Paper: https://arxiv.org/abs/2305.07759
"""

from __future__ import annotations

import argparse
import math
import sys
import time
import urllib.request
from pathlib import Path

import torch
from torch.utils.data import DataLoader, Dataset

from lmt.generate import generate
from lmt.models.config import ModelConfig
from lmt.tokenizer.bpe import BPETokenizer

# TinyStories URLs (raw text files from HuggingFace)
TRAIN_URL = (
    'https://huggingface.co/datasets/roneneldan/TinyStories/'
    'resolve/main/TinyStoriesV2-GPT4-train.txt'
)
VALID_URL = (
    'https://huggingface.co/datasets/roneneldan/TinyStories/'
    'resolve/main/TinyStories-valid.txt'
)
DATA_DIR = Path('data')


def download_file(url: str, dest: Path) -> None:
    """Download a file if it doesn't already exist."""
    if dest.exists():
        size_mb = dest.stat().st_size / 1e6
        print(f'  Already cached: {dest} ({size_mb:.1f} MB)')
        return
    dest.parent.mkdir(parents=True, exist_ok=True)
    print(f'  Downloading {dest.name}...')
    urllib.request.urlretrieve(url, dest)
    size_mb = dest.stat().st_size / 1e6
    print(f'  Downloaded: {size_mb:.1f} MB')


class TokenDataset(Dataset):
    """Pre-tokenized dataset for next-token prediction.

    Stores all tokens in a single flat tensor and creates windows
    of (input, target) pairs with a sliding stride.

    Args:
        tokens: 1D tensor of token IDs.
        context_length: Length of each input/target window.
        stride: Step size between windows.
    """

    def __init__(
        self, tokens: torch.Tensor, context_length: int, stride: int
    ) -> None:
        """Initialize with a flat token tensor and sliding window params."""
        self.tokens = tokens
        self.context_length = context_length
        # Pre-compute valid start indices
        max_start = len(tokens) - context_length - 1
        self.starts = list(range(0, max_start, stride))

    def __len__(self) -> int:
        """Return number of samples."""
        return len(self.starts)

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        """Return (input, target) pair at index."""
        start = self.starts[idx]
        end = start + self.context_length
        return self.tokens[start:end], self.tokens[start + 1 : end + 1]


def tokenize_file(
    path: Path, tokenizer: BPETokenizer, max_tokens: int | None = None
) -> torch.Tensor:
    """Tokenize a text file and return a flat tensor of token IDs."""
    print(f'  Tokenizing {path.name}...')
    text = path.read_text()
    if max_tokens is not None:
        # Rough char estimate: ~4 chars per token for English
        text = text[: max_tokens * 4]
    tokens = tokenizer.encode(text)
    t = torch.tensor(tokens, dtype=torch.long)
    print(f'  {len(t):,} tokens')
    return t


def build_model(arch: str, config: ModelConfig) -> torch.nn.Module:
    """Build a model from architecture name and config."""
    if arch == 'gpt':
        from lmt.models.gpt import GPT

        return GPT(config)
    elif arch == 'llama':
        from lmt.models.llama import LLaMA

        return LLaMA(config)
    elif arch == 'mamba':
        from lmt.models.mamba import Mamba

        return Mamba(config)
    else:
        raise ValueError(f'Unknown architecture: {arch}')


def main() -> int:
    """Train on TinyStories."""
    parser = argparse.ArgumentParser(
        description='Train a small LM on TinyStories'
    )
    parser.add_argument(
        '--arch',
        choices=['gpt', 'llama', 'mamba'],
        default='llama',
        help='Model architecture (default: llama)',
    )
    parser.add_argument(
        '--epochs', type=int, default=3, help='Training epochs'
    )
    parser.add_argument(
        '--embed-dim', type=int, default=128, help='Embedding dimension'
    )
    parser.add_argument(
        '--num-layers', type=int, default=4, help='Number of layers'
    )
    parser.add_argument(
        '--num-heads', type=int, default=4, help='Attention heads'
    )
    parser.add_argument(
        '--context-len', type=int, default=256, help='Context length'
    )
    parser.add_argument(
        '--batch-size', type=int, default=32, help='Batch size'
    )
    parser.add_argument('--lr', type=float, default=3e-4, help='Learning rate')
    parser.add_argument('--device', type=str, default='cpu', help='Device')
    parser.add_argument(
        '--max-train-tokens',
        type=int,
        default=2_000_000,
        help='Max training tokens (default: 2M, use 0 for full dataset)',
    )
    parser.add_argument(
        '--use-valid-only',
        action='store_true',
        help='Train on validation set only (faster, ~19MB)',
    )
    parser.add_argument('--seed', type=int, default=42, help='Random seed')
    args = parser.parse_args()

    torch.manual_seed(args.seed)
    device = torch.device(args.device)

    # Tokenizer
    tokenizer = BPETokenizer(
        encoding_name='gpt2', allowed_special={'<|endoftext|>'}
    )
    vocab_size = 50257  # GPT-2 vocab size
    random_loss = math.log(vocab_size)

    # Download data
    print('Data:')
    valid_path = DATA_DIR / 'tinystories_valid.txt'
    download_file(VALID_URL, valid_path)

    if args.use_valid_only:
        train_path = valid_path
        print('  Using validation set for training (--use-valid-only)')
    else:
        train_path = DATA_DIR / 'tinystories_train.txt'
        download_file(TRAIN_URL, train_path)

    # Tokenize
    print('\nTokenization:')
    max_tok = args.max_train_tokens if args.max_train_tokens > 0 else None
    train_tokens = tokenize_file(train_path, tokenizer, max_tokens=max_tok)
    val_tokens = tokenize_file(valid_path, tokenizer, max_tokens=500_000)

    # Create datasets
    stride = args.context_len  # Non-overlapping windows
    train_ds = TokenDataset(train_tokens, args.context_len, stride)
    val_ds = TokenDataset(val_tokens, args.context_len, stride)
    train_loader = DataLoader(
        train_ds,
        batch_size=args.batch_size,
        shuffle=True,
        drop_last=True,
    )
    val_loader = DataLoader(
        val_ds,
        batch_size=args.batch_size,
        shuffle=False,
        drop_last=True,
    )

    # Model
    config = ModelConfig(
        vocab_size=vocab_size,
        embed_dim=args.embed_dim,
        num_heads=args.num_heads,
        num_kv_heads=args.num_heads,
        num_layers=args.num_layers,
        context_length=args.context_len,
        dropout=0.1,
    )
    model = build_model(args.arch, config)
    model.to(device)
    params = sum(p.numel() for p in model.parameters() if p.requires_grad)

    print(f'\nModel: {args.arch.upper()}')
    print(f'  Parameters: {params:,}')
    print(f'  Embed dim: {args.embed_dim}, Layers: {args.num_layers}')
    print(f'  Context: {args.context_len}, Heads: {args.num_heads}')
    print(f'  Device: {device}')
    print('\nData:')
    print(f'  Train: {len(train_ds):,} samples')
    print(f'  Val: {len(val_ds):,} samples')
    print(f'  Random loss: {random_loss:.3f} (ln({vocab_size}))')

    # Train
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=args.lr, weight_decay=0.01
    )
    print(f'\nTraining for {args.epochs} epochs...')
    start = time.time()

    for epoch in range(args.epochs):
        # Train
        model.train()
        total_loss, n = 0.0, 0
        for inp, tgt in train_loader:
            inp, tgt = inp.to(device), tgt.to(device)
            optimizer.zero_grad()
            logits = model(inp)
            loss = torch.nn.functional.cross_entropy(
                logits.flatten(0, 1), tgt.flatten()
            )
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            total_loss += loss.item()
            n += 1
        train_loss = total_loss / n

        # Eval
        model.eval()
        vl, vn = 0.0, 0
        with torch.no_grad():
            for inp, tgt in val_loader:
                inp, tgt = inp.to(device), tgt.to(device)
                logits = model(inp)
                vl += torch.nn.functional.cross_entropy(
                    logits.flatten(0, 1), tgt.flatten()
                ).item()
                vn += 1
        val_loss = vl / vn if vn else 0

        improvement = (1 - val_loss / random_loss) * 100
        elapsed = time.time() - start
        print(
            f'  Epoch {epoch + 1}/{args.epochs} | '
            f'train: {train_loss:.3f} | val: {val_loss:.3f} | '
            f'{improvement:.1f}% below random | {elapsed:.0f}s'
        )

    total_time = time.time() - start

    # Generate samples
    model.eval()
    prompts = [
        'Once upon a time',
        'The little girl',
        'One day, a big',
    ]
    print('\nGeneration samples (temperature=0.8):')
    for prompt in prompts:
        input_ids = torch.tensor(
            [tokenizer.encode(prompt)], dtype=torch.long, device=device
        )
        with torch.no_grad():
            out = generate(
                model,
                input_ids,
                max_new_tokens=100,
                context_size=args.context_len,
                temperature=0.8,
                top_k=40,
            )
        text = tokenizer.decode(out[0].tolist())
        # Truncate at first endoftext or newline for clean display
        for sep in ['<|endoftext|>', '\n\n']:
            if sep in text[len(prompt) :]:
                text = text[: text.index(sep, len(prompt))]
                break
        print(f'  "{text[:200]}"')

    print(f'\nDone in {total_time:.1f}s ({total_time / 60:.1f} min)')
    return 0


if __name__ == '__main__':
    sys.exit(main())
