#!/usr/bin/env python
"""Prepare the general-language corpus for the 103M pilot.

Source: HuggingFaceFW/fineweb-edu, subset "sample-10BT" (English web text, ODC-By 1.0), streamed in its
canonical shard order (deterministic). Documents are exact-hash deduplicated, assigned to train / val / test
by a hash of their content (never by position), and tokenized with a byte-level BPE (vocab 32768, digits split
individually) trained on the TRAIN split only. The test split is locked: it is written but never read by any
training or development script (see protocol/preregistration.md).

Outputs (data/processed/<name>/): tokenizer.json, {train,val,test}.bin (uint16, documents separated by <|endoftext|>),
manifest.json (counts, sha256, parameters).
"""
from __future__ import annotations

import argparse, hashlib, json, os, sys, time
from pathlib import Path

import numpy as np


def doc_hash(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8", "ignore")).hexdigest()


def split_of(h: str, val_pct: float, test_pct: float) -> str:
    u = int(h[:12], 16) / 16 ** 12
    if u < test_pct:
        return "test"
    if u < test_pct + val_pct:
        return "val"
    return "train"


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--out", default="data/processed/fineweb_edu_300m")
    ap.add_argument("--target-train-tokens", type=int, default=300_000_000)
    ap.add_argument("--vocab-size", type=int, default=32768)
    ap.add_argument("--tokenizer-docs", type=int, default=60_000)
    ap.add_argument("--val-pct", type=float, default=0.01)
    ap.add_argument("--test-pct", type=float, default=0.01)
    ap.add_argument("--max-doc-chars", type=int, default=40_000)
    args = ap.parse_args()
    out = Path(args.out); out.mkdir(parents=True, exist_ok=True)
    raw_dir = Path("data/raw/fineweb_edu"); raw_dir.mkdir(parents=True, exist_ok=True)

    from datasets import load_dataset
    from tokenizers import Tokenizer, models, pre_tokenizers, decoders, trainers

    t0 = time.time()
    ds = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT", split="train", streaming=True)
    # ---- pass 1: stream documents, dedup, assign splits, spill text to jsonl per split --------------------
    seen = set()
    counts = {"train": 0, "val": 0, "test": 0}
    chars = {"train": 0, "val": 0, "test": 0}
    files = {s: open(raw_dir / f"{s}.jsonl", "w") for s in counts}
    tok_sample = []
    # ~4.3 chars per BPE token on English web text: stream until the train split has enough characters.
    target_chars = int(args.target_train_tokens * 4.3)
    n_seen = 0
    for ex in ds:
        n_seen += 1
        text = ex["text"]
        if len(text) > args.max_doc_chars:
            text = text[: args.max_doc_chars]
        h = doc_hash(text)
        if h in seen:
            continue
        seen.add(h)
        s = split_of(h, args.val_pct, args.test_pct)
        files[s].write(json.dumps({"id": ex.get("id"), "h": h, "text": text}) + "\n")
        counts[s] += 1
        chars[s] += len(text)
        if s == "train" and len(tok_sample) < args.tokenizer_docs:
            tok_sample.append(text)
        if n_seen % 20000 == 0:
            print(f"[pass1] seen={n_seen} train_chars={chars['train']/1e6:.1f}M ({time.time()-t0:.0f}s)", flush=True)
        if chars["train"] >= target_chars:
            break
    for f in files.values():
        f.close()
    print("pass1 done", counts, chars, f"{time.time()-t0:.0f}s", flush=True)

    # ---- tokenizer (train split only) ------------------------------------------------------------------
    tok = Tokenizer(models.BPE(unk_token=None))
    tok.pre_tokenizer = pre_tokenizers.Sequence([pre_tokenizers.Digits(individual_digits=True),
                                                 pre_tokenizers.ByteLevel(add_prefix_space=False)])
    tok.decoder = decoders.ByteLevel()
    trainer = trainers.BpeTrainer(vocab_size=args.vocab_size, special_tokens=["<|endoftext|>", "<|pad|>"],
                                  initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), show_progress=False)
    tok.train_from_iterator(tok_sample, trainer=trainer)
    tok.save(str(out / "tokenizer.json"))
    eot = tok.token_to_id("<|endoftext|>")
    assert tok.get_vocab_size() == args.vocab_size, tok.get_vocab_size()
    print("tokenizer trained", tok.get_vocab_size(), f"{time.time()-t0:.0f}s", flush=True)

    # ---- pass 2: tokenize each split to uint16 ---------------------------------------------------------
    manifest = {"source": "HuggingFaceFW/fineweb-edu sample-10BT (streaming, canonical shard order)",
                "license": "ODC-By 1.0 (https://opendatacommons.org/licenses/by/1-0/)",
                "prepared_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), "args": vars(args), "docs_streamed": n_seen,
                "docs": counts, "chars": chars, "splits": {}, "tokenizer_sha256": hashlib.sha256((out / "tokenizer.json").read_bytes()).hexdigest(),
                "eot_id": eot, "dtype": "uint16"}
    for s in counts:
        ids_all = []
        with open(raw_dir / f"{s}.jsonl") as f:
            batch = []
            for line in f:
                batch.append(json.loads(line)["text"])
                if len(batch) == 2000:
                    for enc in tok.encode_batch(batch):
                        ids_all.append(np.asarray(enc.ids + [eot], dtype=np.uint16))
                    batch = []
            if batch:
                for enc in tok.encode_batch(batch):
                    ids_all.append(np.asarray(enc.ids + [eot], dtype=np.uint16))
        arr = np.concatenate(ids_all) if ids_all else np.zeros(0, dtype=np.uint16)
        arr.tofile(out / f"{s}.bin")
        manifest["splits"][s] = {"tokens": int(arr.size), "docs": len(ids_all),
                                 "sha256": hashlib.sha256(arr.tobytes()).hexdigest(), "locked": s == "test"}
        print(f"[pass2] {s}: {arr.size:,} tokens ({time.time()-t0:.0f}s)", flush=True)
    (out / "manifest.json").write_text(json.dumps(manifest, indent=2))
    print("done", json.dumps(manifest["splits"]))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())