# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""Cross-check the Python (vectorized) and C++ (greedy lowest-id) BPE
encoders on real text.

Encodes a text sample with tools/prepare_data.py's encoder, dumps the ids,
then runs test_tokenizer with FORGE_TOK_* pointing at them so the C++ side
compares. Exits nonzero on disagreement.

Usage:
  python3 tests/tokenizer_agreement.py --model data/tinystories/tok4096.model \
      --text data/tinystories/TinyStoriesV2-GPT4-valid.txt --bytes 200000 \
      --binary build/test_tokenizer
"""
import argparse
import os
import subprocess
import sys
import tempfile

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "tools"))
import prepare_data as pd  # noqa: E402


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--text", required=True)
    ap.add_argument("--bytes", type=int, default=200000)
    ap.add_argument("--binary", default="build/test_tokenizer")
    args = ap.parse_args()

    with open(args.text, "rb") as f:
        sample = f.read(args.bytes)
    # cut at a story boundary so the comparison matches how .bin is built
    sep = b"<|endoftext|>"
    if sep in sample:
        sample = sample[: sample.rindex(sep)]

    pd._MERGES = pd.load_merges(args.model)
    ids = pd.encode_block(sample)
    print(f"python: {len(sample)} bytes -> {len(ids)} tokens")

    with tempfile.TemporaryDirectory() as td:
        text_path = os.path.join(td, "text.bin")
        ids_path = os.path.join(td, "ids.txt")
        with open(text_path, "wb") as f:
            f.write(sample)
        with open(ids_path, "w") as f:
            f.write("\n".join(str(int(i)) for i in ids))

        env = dict(os.environ)
        env["FORGE_TOK_MODEL"] = os.path.abspath(args.model)
        env["FORGE_TOK_TEXT"] = text_path
        env["FORGE_TOK_IDS"] = ids_path
        rc = subprocess.call([os.path.abspath(args.binary)], env=env,
                             cwd=os.path.dirname(os.path.abspath(args.binary)))
    sys.exit(rc)


if __name__ == "__main__":
    main()