#!/usr/bin/env python3
"""Download TinyStories (Eldan & Li 2023, V2 / GPT-4 split) and tokenise it for nanoGPT.

Usage (from 06-experiments/):
    .venv/bin/python data/prepare_tinystories.py            # full: ~2.2 GB train text
    .venv/bin/python data/prepare_tinystories.py --spike    # small: valid file only, split 90/10

Writes data/<name>/{train,val}.bin (uint16 GPT-2 BPE tokens, nanoGPT format) and
data/<name>/prep.json (source URLs, SHA-256 of the raw files, token counts), so every
run can record exactly which data it used. Raw text and .bin files are git-ignored.
"""
import hashlib
import json
import os
import sys
import urllib.request

import numpy as np
import tiktoken

BASE = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/"
TRAIN, VALID = "TinyStoriesV2-GPT4-train.txt", "TinyStoriesV2-GPT4-valid.txt"
HERE = os.path.dirname(os.path.abspath(__file__))
SEP = "<|endoftext|>"


def download(fname, raw_dir):
    path = os.path.join(raw_dir, fname)
    if not os.path.exists(path):
        print(f"downloading {fname} ...")
        tmp = path + ".part"
        urllib.request.urlretrieve(BASE + fname, tmp)
        os.rename(tmp, path)
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return path, h.hexdigest()


def stories(path):
    """Yield one story at a time (stories are separated by <|endoftext|>)."""
    buf = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            if SEP in line:
                head, _, tail = line.partition(SEP)
                buf.append(head)
                s = "".join(buf).strip()
                if s:
                    yield s
                buf = [tail]
            else:
                buf.append(line)
    s = "".join(buf).strip()
    if s:
        yield s


def write_bin(texts, out_path, enc):
    eot = enc.eot_token
    ids = []
    for t in texts:
        ids.extend(enc.encode_ordinary(t))
        ids.append(eot)
    arr = np.array(ids, dtype=np.uint16)
    arr.tofile(out_path)
    return len(arr)


def main():
    spike = "--spike" in sys.argv
    name = "tinystories_spike" if spike else "tinystories"
    out_dir = os.path.join(HERE, name)
    raw_dir = os.path.join(HERE, "raw")
    os.makedirs(out_dir, exist_ok=True)
    os.makedirs(raw_dir, exist_ok=True)
    enc = tiktoken.get_encoding("gpt2")
    meta = {"dataset": "TinyStories V2 (GPT-4)", "source": BASE, "tokenizer": "tiktoken gpt2", "files": {}}

    vpath, vsha = download(VALID, raw_dir)
    meta["files"][VALID] = vsha
    if spike:
        allv = list(stories(vpath))
        cut = int(len(allv) * 0.9)
        meta["split"] = f"valid file split 90/10 by story order: {cut} train / {len(allv) - cut} val stories"
        n_tr = write_bin(allv[:cut], os.path.join(out_dir, "train.bin"), enc)
        n_va = write_bin(allv[cut:], os.path.join(out_dir, "val.bin"), enc)
    else:
        tpath, tsha = download(TRAIN, raw_dir)
        meta["files"][TRAIN] = tsha
        meta["split"] = "official train / valid files"
        n_tr = write_bin(stories(tpath), os.path.join(out_dir, "train.bin"), enc)
        n_va = write_bin(stories(vpath), os.path.join(out_dir, "val.bin"), enc)
    meta["tokens"] = {"train": n_tr, "val": n_va}
    json.dump(meta, open(os.path.join(out_dir, "prep.json"), "w"), indent=2)
    print(json.dumps(meta, indent=2))


if __name__ == "__main__":
    main()