"""Build the pre-training corpus: clean, deduplicate, train the tokenizer, encode.

Sources (all English, all openly licensed):
  - TinyStories V2 (GPT-4 part)  - short stories in a 3-4 year old's vocabulary
  - SimpleStories                 - simple stories with many topics, themes and styles
  - SODA                          - everyday two-person dialogues with a short narrative

Output: data/train.bin and data/val.bin (uint16 token ids, every document starts
with <|endoftext|>), data/tokenizer.json and data/corpus_stats.json.
"""
import argparse
import glob
import hashlib
import json
import os
import random
import time
from multiprocessing import Pool

import numpy as np
import pyarrow.parquet as pq

from tokenizer import Tokenizer

REPLACE = str.maketrans({'‘': "'", '’': "'", '“': '"', '”': '"', '–': '-',
                         '—': ' - ', '…': '...', ' ': ' ', '\r': ''})


def clean(text):
    """Normalise quotes and dashes, drop stray spaces; None if the text should go."""
    text = text.translate(REPLACE).strip()
    lines = [' '.join(line.split()) for line in text.split('\n')]
    text = '\n'.join(lines).strip()
    while '\n\n\n' in text:
        text = text.replace('\n\n\n', '\n\n')
    if len(text) < 80:
        return None
    if sum(ord(c) > 126 for c in text) > 0.002 * len(text):   # other scripts, mojibake
        return None
    return text


def tinystories(raw):
    with open(f'{raw}/TinyStoriesV2-GPT4-train.txt', encoding='utf-8') as f:
        for story in f.read().split('<|endoftext|>'):
            yield story


def simplestories(raw):
    for path in sorted(glob.glob(f'{raw}/simplestories-[0-9].parquet')):
        yield from pq.read_table(path, columns=['story']).column('story').to_pylist()


def soda(raw):
    table = pq.read_table(f'{raw}/soda-train.parquet', columns=['narrative', 'dialogue', 'speakers'])
    for narrative, lines, speakers in zip(*(table.column(c).to_pylist() for c in table.column_names)):
        if len(speakers) != len(lines):
            continue
        dialogue = '\n'.join(f'{who}: {line}' for who, line in zip(speakers, lines))
        yield f'{narrative}\n\n{dialogue}'


SOURCES = {'tinystories': tinystories, 'simplestories': simplestories, 'soda': soda}

_tok = None


def _init(path):
    global _tok
    _tok = Tokenizer.load(path)


def _encode(texts):
    eot = _tok['<|endoftext|>']
    out = []
    for t in texts:
        out.append(eot)
        out.extend(_tok.encode(t, allow_special=False))
    return np.array(out, dtype=np.uint16)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--raw', default='raw')
    ap.add_argument('--out', default='data')
    ap.add_argument('--vocab', type=int, default=8192)
    ap.add_argument('--tok-docs', type=int, default=40_000, help='documents per source for BPE training')
    ap.add_argument('--val-frac', type=float, default=0.005)
    args = ap.parse_args()
    os.makedirs(args.out, exist_ok=True)
    rng = random.Random(1337)

    # 1. clean + exact deduplication
    docs, stats = {}, {}
    for name, read in SOURCES.items():
        seen, kept, raw_count = set(), [], 0
        for text in read(args.raw):
            raw_count += 1
            text = clean(text)
            if text is None:
                continue
            h = hashlib.md5(text.lower().encode()).digest()
            if h in seen:
                continue
            seen.add(h)
            kept.append(text)
        rng.shuffle(kept)
        docs[name] = kept
        stats[name] = {'raw': raw_count, 'kept': len(kept), 'chars': sum(map(len, kept))}
        print(f'{name}: {raw_count} -> {len(kept)} documents', flush=True)

    # 2. tokenizer on a balanced sample
    tok_path = f'{args.out}/tokenizer.json'
    if not os.path.exists(tok_path):
        t0 = time.time()
        sample = [d for name in docs for d in docs[name][:args.tok_docs]]
        tok = Tokenizer.train(sample, args.vocab, verbose=True)
        tok.save(tok_path)
        print(f'tokenizer: {tok.vocab_size} tokens in {time.time() - t0:.0f}s', flush=True)

    # 3. encode in parallel; a small slice of every source goes to validation
    tok = Tokenizer.load(tok_path)
    eot = tok['<|endoftext|>']
    train, val = [], []
    with Pool(max(1, os.cpu_count() - 2), initializer=_init, initargs=(tok_path,)) as pool:
        for name, texts in docs.items():
            n_val = int(len(texts) * args.val_frac)
            for split, part in ((val, texts[:n_val]), (train, texts[n_val:])):
                chunks = [part[i:i + 2000] for i in range(0, len(part), 2000)]
                split.append(np.concatenate(pool.map(_encode, chunks)))
            stats[name]['tokens'] = int(len(train[-1]) + len(val[-1]))
            print(f'{name}: {stats[name]["tokens"] / 1e6:.1f}M tokens', flush=True)
    # shuffle whole runs of documents between sources, so any stretch of the file is a fair mix
    blocks = []
    for arr in train:
        starts = np.flatnonzero(arr == eot)
        cuts = starts[np.searchsorted(starts, np.arange(1_000_000, len(arr), 1_000_000))]
        blocks.extend(np.split(arr, np.unique(cuts)))
    rng.shuffle(blocks)
    np.concatenate(blocks).tofile(f'{args.out}/train.bin')
    np.concatenate(val).tofile(f'{args.out}/val.bin')
    stats['train_tokens'] = int(sum(len(b) for b in blocks))
    stats['val_tokens'] = int(sum(len(v) for v in val))
    with open(f'{args.out}/corpus_stats.json', 'w') as f:
        json.dump(stats, f, indent=1)


if __name__ == '__main__':
    main()