# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""Download TinyStories, train (or reuse) a forgebpe tokenizer, tokenize,
and write train.bin / val.bin as uint16 token streams with the llm.c-style
header {magic 20240520, version 1, num_tokens} (RESEARCH.md §7).

Encoder equivalence: src/tokenizer/bpe.cpp repeatedly applies the *lowest-id*
merge present. Because a merge (a,b)->c always has c > a and c > b, applying
merge c can only create pairs whose merge id exceeds c — so no lower-id merge
can appear later. Hence one ascending pass over the merge list, applying all
non-overlapping occurrences of each, is equivalent and vectorizes cleanly.
tests/test_tokenizer.cpp checks the two agree on real text.

Usage:
  python3 tools/prepare_data.py --out data/tinystories [--vocab-size 4096]
      [--max-train-mb 100] [--max-val-mb 5]
"""
import argparse
import multiprocessing as mp
import os
import struct
import urllib.request

import numpy as np

BASE = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main"
TRAIN_FILE = "TinyStoriesV2-GPT4-train.txt"
VAL_FILE = "TinyStoriesV2-GPT4-valid.txt"
MAGIC, VERSION = 20240520, 1

_MERGES = None  # worker global: list of (idx, left, right) in ascending id order


def download(url, dest):
    if os.path.exists(dest):
        print(f"already have {dest}")
        return
    print(f"downloading {url}")
    tmp = dest + ".tmp"
    urllib.request.urlretrieve(url, tmp)
    os.rename(tmp, dest)


def load_merges(model_path):
    merges = []
    with open(model_path) as f:
        assert f.readline().strip() == "forgebpe v1"
        f.readline()  # vocab size
        for line in f:
            idx, left, right = map(int, line.split())
            merges.append((idx, left, right))
    merges.sort()
    return merges


def apply_merge(ids, idx, left, right):
    is_pair = (ids[:-1] == left) & (ids[1:] == right)
    pos = np.where(is_pair)[0]
    if pos.size == 0:
        return ids
    if left == right:  # "aaa": keep the first of each run, non-overlapping
        keep, last = [], -2
        for p in pos:
            if p != last + 1:
                keep.append(p)
                last = p
        pos = np.asarray(keep, dtype=np.int64)
    ids[pos] = idx
    return np.delete(ids, pos + 1)


def encode_block(block_bytes):
    ids = np.frombuffer(block_bytes, dtype=np.uint8).astype(np.int32)
    for idx, left, right in _MERGES:
        if ids.size < 2:
            break
        ids = apply_merge(ids, idx, left, right)
    return ids.astype(np.uint16)


def _init_worker(model_path):
    global _MERGES
    _MERGES = load_merges(model_path)


def tokenize_file(text_path, model_path, out_path, max_mb, block_mb=2.0):
    with open(text_path, "rb") as f:
        data = f.read(int(max_mb * 1024 * 1024))

    # split into blocks on story boundaries so merges never straddle documents
    sep = b"<|endoftext|>"
    blocks, cur = [], bytearray()
    limit = int(block_mb * 1024 * 1024)
    for story in data.split(sep):
        if not story.strip():
            continue
        cur += story
        if len(cur) >= limit:
            blocks.append(bytes(cur))
            cur = bytearray()
    if cur:
        blocks.append(bytes(cur))
    print(f"tokenizing {len(data) / 1e6:.1f} MB in {len(blocks)} blocks")

    with mp.Pool(initializer=_init_worker, initargs=(model_path,)) as pool:
        parts = []
        for i, ids in enumerate(pool.imap(encode_block, blocks)):
            parts.append(ids)
            if (i + 1) % 10 == 0:
                print(f"  block {i + 1}/{len(blocks)}")
    tokens = np.concatenate(parts) if parts else np.zeros(0, dtype=np.uint16)
    print(f"  -> {len(tokens) / 1e6:.2f}M tokens "
          f"({len(data) / max(1, len(tokens)):.2f} bytes/token)")

    with open(out_path, "wb") as f:
        header = np.zeros(256, dtype=np.int32)
        header[0], header[1], header[2] = MAGIC, VERSION, len(tokens)
        f.write(header.tobytes())
        f.write(tokens.tobytes())
    print(f"wrote {out_path} ({os.path.getsize(out_path) / 1e6:.1f} MB)")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", required=True)
    ap.add_argument("--vocab-size", type=int, default=4096)
    ap.add_argument("--max-train-mb", type=float, default=100.0)
    ap.add_argument("--max-val-mb", type=float, default=5.0)
    ap.add_argument("--tokenizer-mb", type=float, default=10.0,
                    help="corpus sample size for BPE training")
    ap.add_argument("--tokenizer", default=None,
                    help="existing .model; otherwise trained on the corpus")
    args = ap.parse_args()

    os.makedirs(args.out, exist_ok=True)
    train_txt = os.path.join(args.out, TRAIN_FILE)
    val_txt = os.path.join(args.out, VAL_FILE)
    download(f"{BASE}/{TRAIN_FILE}?download=true", train_txt)
    download(f"{BASE}/{VAL_FILE}?download=true", val_txt)

    model = args.tokenizer or os.path.join(args.out, f"tok{args.vocab_size}.model")
    if not os.path.exists(model):
        import sys
        sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
        import train_tokenizer as tt
        with open(train_txt, "rb") as f:
            sample = f.read(int(args.tokenizer_mb * 1024 * 1024))
        print(f"training {args.vocab_size}-vocab BPE on {len(sample) / 1e6:.1f} MB")
        tt.write_model(model, tt.train(sample, args.vocab_size))
        print(f"wrote {model}")

    tokenize_file(train_txt, model, os.path.join(args.out, "train.bin"),
                  args.max_train_mb)
    tokenize_file(val_txt, model, os.path.join(args.out, "val.bin"), args.max_val_mb)


if __name__ == "__main__":
    main()