"""Train the ZENITH BPE tokenizer on a text sample and export the fast (HF) copy.

Usage: python -m zenith.tokenizer.train_tokenizer --sample data/interim/tokenizer_sample.txt \
           --out zenith/tokenizer/zenith-16k.json
"""

from __future__ import annotations

import argparse
import random
from pathlib import Path

from zenith.tokenizer.bpe import train_bpe
from zenith.tokenizer.export import encode_fast, export_hf, load_fast


def main(argv=None) -> None:
    ap = argparse.ArgumentParser(prog="python -m zenith.tokenizer.train_tokenizer")
    ap.add_argument("--sample", type=Path, required=True)
    ap.add_argument("--out", type=Path, required=True, help="Python tokenizer path; the HF copy goes next to it as .hf.json")
    ap.add_argument("--vocab-size", type=int, default=16384)
    ap.add_argument("--parity-lines", type=int, default=10_000)
    args = ap.parse_args(argv)

    text = args.sample.read_text()
    docs = [d for d in text.split("\n\n") if d.strip()]
    tok = train_bpe(docs, vocab_size=args.vocab_size)
    tok.save(args.out)
    hf_path = args.out.with_name(args.out.stem + ".hf.json")
    export_hf(tok, hf_path)

    lines = [l for l in text.splitlines() if l.strip()]
    random.Random(0).shuffle(lines)
    lines = lines[: args.parity_lines]
    fast_ids = encode_fast(load_fast(hf_path), lines)
    for line, ids in zip(lines, fast_ids):
        if ids != tok.encode(line):
            raise SystemExit(f"parity failure on line: {line[:80]!r}")
    n_tokens = sum(len(ids) for ids in fast_ids)
    n_bytes = sum(len(l.encode()) for l in lines)
    print(f"parity ok: {len(lines)} lines")
    print(f"bytes/token: {n_bytes / max(n_tokens, 1):.2f}")


if __name__ == "__main__":
    main()