"""FineWeb-Edu sample-10BT -> 32k ByteLevel BPE -> train.bin/val.bin (uint16) + meta.pkl + tokenizer.json.
Same bin format as tinystories_bpe; trainers consume it via --data fineweb_edu.
Stage-2 corpus (recipe order: TinyStories epoch -> FineWeb-Edu -> OLMo2-corpus class).

Phases (all resumable-ish, markers for the watcher):
  1. snapshot_download of sample/10BT parquets (~28GB)          -> DOWNLOAD_DONE
  2. train 32k BPE on ~1.5GB sample from shard 0               -> TOKENIZER_DONE
  3. tokenize shard 0: first 20M tokens -> val.bin, rest ->
     train.bin (disjoint)                                       -> SMOKE_READY
  4. tokenize remaining shards, append to train.bin             -> DONE_FINEWEB
Docs are joined with a <|eot|> separator (id 0). vocab 32768 fits uint16.
NFS note: peak disk = raw parquet ~28GB + bins ~20GB; keep raw/ for tokenizer reruns.
"""
import pickle, time
from pathlib import Path
import numpy as np
import pyarrow.parquet as pq
from huggingface_hub import snapshot_download
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.decoders import ByteLevel as ByteLevelDec

D = Path('/home/yurenh2/ept/ep_run/data/fineweb_edu')
RAW = D / 'raw'
D.mkdir(parents=True, exist_ok=True)
VOCAB = 32768
EOT = '<|eot|>'
SAMPLE_BYTES = 1_500_000_000   # tokenizer-training text cap
VAL_TOKENS = 20_000_000

t0 = time.time()
snapshot_download(repo_id='HuggingFaceFW/fineweb-edu', repo_type='dataset',
                  allow_patterns=['sample/10BT/*'], local_dir=str(RAW))
shards = sorted((RAW / 'sample' / '10BT').glob('*.parquet'))
print(f"DOWNLOAD_DONE {len(shards)} shards {time.time()-t0:.0f}s", flush=True)

def docs(shard, batch_rows=2000):
    for b in pq.ParquetFile(shard).iter_batches(batch_size=batch_rows, columns=['text']):
        for t in b.column('text').to_pylist():
            if t: yield t

tok_path = D / 'tokenizer.json'
if tok_path.exists():
    tok = Tokenizer.from_file(str(tok_path))
    print(f"TOKENIZER_DONE (cached) vocab={tok.get_vocab_size()}", flush=True)
else:
    def sample_iter():
        n = 0
        for t in docs(shards[0]):
            yield t
            n += len(t)
            if n > SAMPLE_BYTES: return
    tok = Tokenizer(BPE(unk_token=None))
    tok.pre_tokenizer = ByteLevel(add_prefix_space=False)
    tok.decoder = ByteLevelDec()
    tok.train_from_iterator(sample_iter(), BpeTrainer(vocab_size=VOCAB, special_tokens=[EOT], show_progress=False))
    tok.save(str(tok_path))
    print(f"TOKENIZER_DONE vocab={tok.get_vocab_size()} {time.time()-t0:.0f}s", flush=True)
eot_id = tok.token_to_id(EOT)
assert tok.get_vocab_size() <= 65535 and eot_id is not None
pickle.dump({'vocab_size': tok.get_vocab_size()}, open(D / 'meta.pkl', 'wb'))

def encode_shard(shard, sink, head_sink=None, head_n=0):
    """tokenize one shard; first head_n tokens to head_sink (if given), rest to sink. Returns tokens written."""
    wrote_head, n, buf = 0, 0, []
    def flush(buf):
        nonlocal wrote_head, n
        ids = []
        for e in tok.encode_batch(buf):
            ids.extend(e.ids); ids.append(eot_id)
        a = np.array(ids, dtype=np.uint16)
        if head_sink is not None and wrote_head < head_n:
            k = min(head_n - wrote_head, len(a))
            a[:k].tofile(head_sink); wrote_head += k; a = a[k:]
        a.tofile(sink); n += len(a)
    for t in docs(shard):
        buf.append(t)
        if len(buf) >= 1000: flush(buf); buf = []
    if buf: flush(buf)
    return n + wrote_head

train_f = open(D / 'train.bin', 'wb')
val_f = open(D / 'val.bin', 'wb')
n0 = encode_shard(shards[0], train_f, head_sink=val_f, head_n=VAL_TOKENS)
val_f.close(); train_f.flush()
print(f"SMOKE_READY shard0 {n0/1e6:.0f}M tokens (val {VAL_TOKENS/1e6:.0f}M) {time.time()-t0:.0f}s", flush=True)

total = n0
for i, sh in enumerate(shards[1:], 1):
    ns = encode_shard(sh, train_f)
    total += ns
    print(f"[shard {i}/{len(shards)-1}] +{ns/1e6:.0f}M -> total {total/1e9:.2f}B tokens {time.time()-t0:.0f}s", flush=True)
train_f.close()
print(f"DONE_FINEWEB total {total/1e9:.2f}B tokens (train {(total-VAL_TOKENS)/1e9:.2f}B / val {VAL_TOKENS/1e6:.0f}M) vocab {tok.get_vocab_size()} {time.time()-t0:.0f}s", flush=True)