"""Train the BPE tokenizer on a sample of the corpus.

Pure-Python BPE training is superlinear, so we train on a sample rather than
the full 2 GB. TinyStories' restricted vocabulary means a ~50 MB sample yields
essentially the same merges as the whole corpus.

Usage:
    python scripts/train_tokenizer.py                 # defaults: 50 MB, 4096 vocab
    python scripts/train_tokenizer.py --sample-mb 25 --vocab-size 4096

Writes the tokenizer to  artifacts/tokenizer.json  and prints the compression
ratio (bytes per token) measured on held-out text.
"""

import argparse
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "src"))

from tokenizer import BPETokenizer  # noqa: E402

EOT = "<|endoftext|>"
TRAIN_TXT = ROOT / "data" / "raw" / "TinyStoriesV2-GPT4-train.txt"
VALID_TXT = ROOT / "data" / "raw" / "TinyStoriesV2-GPT4-valid.txt"
OUT = ROOT / "artifacts" / "tokenizer.json"


def read_prefix(path: Path, mb: float) -> str:
    """Read the first `mb` megabytes of a file as UTF-8 (ignoring a split char)."""
    with open(path, "rb") as f:
        data = f.read(int(mb * 1_000_000))
    return data.decode("utf-8", errors="ignore")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--sample-mb", type=float, default=50.0)
    ap.add_argument("--vocab-size", type=int, default=4096)
    args = ap.parse_args()

    if not TRAIN_TXT.exists():
        sys.exit(f"missing {TRAIN_TXT} — run scripts/download_data.py first")

    print(f"reading {args.sample_mb:g} MB sample from {TRAIN_TXT.name}")
    sample = read_prefix(TRAIN_TXT, args.sample_mb)

    print(f"training BPE: target vocab {args.vocab_size}, special token {EOT!r}")
    t0 = time.perf_counter()
    tok = BPETokenizer.train(sample, vocab_size=args.vocab_size, special_tokens=[EOT])
    dt = time.perf_counter() - t0
    print(f"trained in {dt:.1f}s — {len(tok.merges)} merges, "
          f"vocab_size {tok.vocab_size}")

    OUT.parent.mkdir(parents=True, exist_ok=True)
    tok.save(OUT)
    print(f"saved tokenizer to {OUT}")

    # Compression ratio on held-out validation text (not used for training).
    held = read_prefix(VALID_TXT, 5.0) if VALID_TXT.exists() else sample[: 5_000_000]
    ids = tok.encode(held)
    n_bytes = len(held.encode("utf-8"))
    print(f"\ncompression on held-out text: {n_bytes:,} bytes -> {len(ids):,} tokens"
          f"  ({n_bytes / len(ids):.2f} bytes/token)")


if __name__ == "__main__":
    main()