"""
Tokenizer Training — Built from Scratch using BPE
==================================================

What is a tokenizer and why do we need one?
- Neural networks can't read text directly — they need numbers
- A tokenizer converts text → sequence of integer IDs
- Each ID maps to a "token" (a word, subword, or character)

Why BPE (Byte Pair Encoding)?
- Used by GPT-2, GPT-3, GPT-4, Llama, Mistral, and most modern LLMs
- Starts with individual characters
- Repeatedly merges the most frequent pair of adjacent tokens
- Result: common words become single tokens, rare words get split into subwords
- Example: "unhappiness" → ["un", "happiness"] or ["un", "happ", "iness"]

Why train our OWN tokenizer?
- Pre-trained tokenizers are optimized for general English text
- A domain-specific tokenizer encodes your data more efficiently
- Fewer tokens per sentence = faster training = longer effective context
- Shows you understand the full pipeline end-to-end
"""

import os
import json
from pathlib import Path
from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders, processors
from datasets import load_dataset


def train_tokenizer(
    dataset_name: str = "roneneldan/TinyStories",
    vocab_size: int = 32000,
    output_dir: str = "data/tokenizer",
    split: str = "train",
    text_column: str = "text",
    num_samples: int = 100000,
):
    """
    Train a BPE tokenizer from scratch on your dataset.

    Steps:
    1. Load raw text data from the dataset
    2. Initialize a BPE tokenizer model
    3. Configure special tokens (padding, start, end, unknown)
    4. Train the tokenizer to learn merges
    5. Save the tokenizer for later use

    Args:
        dataset_name: HuggingFace dataset to train on
        vocab_size: number of unique tokens to learn
        output_dir: where to save the tokenizer
        split: which dataset split to use
        text_column: name of the text column in the dataset
        num_samples: how many samples to use for training (more = better but slower)
    """
    print(f"Training BPE tokenizer with vocab_size={vocab_size}")
    print(f"Dataset: {dataset_name}")

    # Load dataset
    print("Loading dataset...")
    dataset = load_dataset(dataset_name, split=split, streaming=True)

    # Collect text samples
    print(f"Collecting {num_samples} text samples...")
    texts = []
    for i, sample in enumerate(dataset):
        if i >= num_samples:
            break
        text = sample.get(text_column, "")
        if text:
            texts.append(text)

    print(f"Collected {len(texts)} samples")

    # Initialize BPE tokenizer
    tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))

    # Pre-tokenization: split text into words before BPE
    # ByteLevel pre-tokenizer handles all Unicode characters
    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)

    # Define special tokens
    special_tokens = ["<pad>", "<unk>", "<bos>", "<eos>"]

    # Configure the trainer
    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=special_tokens,
        show_progress=True,
        min_frequency=2,  # token pair must appear at least 2 times
    )

    # Train the tokenizer on our text corpus
    print("Training tokenizer (this may take a few minutes)...")
    tokenizer.train_from_iterator(texts, trainer=trainer)

    # Post-processing: add BOS/EOS tokens automatically
    bos_id = tokenizer.token_to_id("<bos>")
    eos_id = tokenizer.token_to_id("<eos>")
    tokenizer.post_processor = processors.TemplateProcessing(
        single=f"<bos>:0 $A:0 <eos>:0",
        special_tokens=[("<bos>", bos_id), ("<eos>", eos_id)],
    )

    # Decoder: convert tokens back to text
    tokenizer.decoder = decoders.ByteLevel()

    # Save
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    tokenizer.save(str(output_path / "tokenizer.json"))

    # Save config for reference
    config = {
        "vocab_size": tokenizer.get_vocab_size(),
        "dataset": dataset_name,
        "num_training_samples": len(texts),
        "special_tokens": {
            "pad": "<pad>",
            "unk": "<unk>",
            "bos": "<bos>",
            "eos": "<eos>",
        },
        "special_token_ids": {
            "pad": tokenizer.token_to_id("<pad>"),
            "unk": tokenizer.token_to_id("<unk>"),
            "bos": bos_id,
            "eos": eos_id,
        },
    }
    with open(output_path / "config.json", "w") as f:
        json.dump(config, f, indent=2)

    print(f"\nTokenizer saved to {output_path}")
    print(f"Vocabulary size: {tokenizer.get_vocab_size()}")

    # Demo
    test_text = "Once upon a time, there was a little cat."
    encoded = tokenizer.encode(test_text)
    print(f"\nDemo encoding:")
    print(f"  Text:   {test_text}")
    print(f"  Tokens: {encoded.tokens}")
    print(f"  IDs:    {encoded.ids}")
    print(f"  Decoded: {tokenizer.decode(encoded.ids)}")

    return tokenizer


def load_tokenizer(tokenizer_path: str = "data/tokenizer"):
    """Load a previously trained tokenizer."""
    path = Path(tokenizer_path) / "tokenizer.json"
    if not path.exists():
        raise FileNotFoundError(
            f"Tokenizer not found at {path}. Run tokenizer training first."
        )
    tokenizer = Tokenizer.from_file(str(path))
    return tokenizer


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Train a BPE tokenizer from scratch")
    parser.add_argument("--dataset", type=str, default="roneneldan/TinyStories")
    parser.add_argument("--vocab-size", type=int, default=32000)
    parser.add_argument("--output-dir", type=str, default="data/tokenizer")
    parser.add_argument("--num-samples", type=int, default=100000)
    args = parser.parse_args()

    train_tokenizer(
        dataset_name=args.dataset,
        vocab_size=args.vocab_size,
        output_dir=args.output_dir,
        num_samples=args.num_samples,
    )