"""Regenerate align/tokenizer/_qwen2_unicode.py -- the Unicode tables the Qwen2 byte-level
tokenizer must use to be byte-identical with the HF/Rust tokenizer.

RUN ONLY in the isolated rollout venv (it has transformers; the align venv must never):

    CUDA_VISIBLE_DEVICES= .venv-rollout/bin/python tests/data/make_qwen2_unicode_tables.py

Why measured instead of taken from `unicodedata`: the Rust tokenizer's tables are NOT the
tables of any one Python. Measured against transformers 5.19.0 / tokenizers 0.23.3
(Python 3.12.11, unicodedata 15.0.0):

  * regex classes (\\p{L}, \\p{N}) are NEWER than Python's: 5004 code points added in
    Unicode 16 (Todhri, Garay, Tulu-Tigalari, Egyptian Hieroglyphs Ext-A, CJK Ext I, ...)
    are letters/digits for the Rust side and unassigned (Cn) for Python 3.12, so "X's"
    split differently ("X" + "'s" vs "X'" + "s");
  * the NFC normalizer is OLDER than Python's: 59 combining marks assigned in
    Unicode 10-15 keep combining class 0 on the Rust side (never canonically reordered)
    and U+11935 U+11930 does not compose to U+11938.

This script probes the real HF tokenizer for every code point and freezes the answer.
"""
import os
import sys
import unicodedata as ud

REPO = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
SNAP = os.environ.get(
    "GOEDEL_SNAPSHOT",
    "/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/"
    "dfd02e6271a58375dfbf3ece0175277cf6b6a89a",
)
OUT = os.path.join(REPO, "align", "tokenizer", "_qwen2_unicode.py")

import tokenizers  # noqa: E402
import transformers  # noqa: E402
from transformers import AutoTokenizer  # noqa: E402

hf = AutoTokenizer.from_pretrained(SNAP)
pt = hf.backend_tokenizer.pre_tokenizer

WHITE_SPACE = {c for a, b in (
    (0x09, 0x0D), (0x20, 0x20), (0x85, 0x85), (0xA0, 0xA0), (0x1680, 0x1680),
    (0x2000, 0x200A), (0x2028, 0x2029), (0x202F, 0x202F), (0x205F, 0x205F), (0x3000, 0x3000),
) for c in range(a, b + 1)}


def spans(s):
    return [o for _, o in pt.pre_tokenize_str(s)]


def ranges(cps):
    out = []
    for c in sorted(cps):
        if out and c == out[-1][1] + 1:
            out[-1][1] = c
        else:
            out.append([c, c])
    return tuple((a, b) for a, b in out)


L, N = set(), set()
for cp in range(0x110000):
    if 0xD800 <= cp <= 0xDFFF or cp in WHITE_SPACE:
        continue
    c = chr(cp)
    if len(spans(c + c)) == 2:                      # \p{N}: one digit per piece
        N.add(cp)
    elif spans(c + "'s") == [(0, 1), (1, 3)]:       # \p{L}+ then the 's contraction
        L.add(cp)

# ---- NFC: code points whose combining class the Rust side ignores / compositions it lacks
def reorders(cp):
    c = chr(cp)
    k = ud.combining(c)
    s = ("́" + c) if k < 230 else (c + "̖")      # 0301 is ccc 230, 0316 is ccc 220
    want = ud.normalize("NFC", s)
    return hf.decode(hf(s, add_special_tokens=False)["input_ids"]) == want


ccc_ignored = {cp for cp in range(0x110000)
               if not 0xD800 <= cp <= 0xDFFF and ud.combining(chr(cp)) and not reorders(cp)}
comp_second = set()
for cp in range(0x110000):
    if 0xD800 <= cp <= 0xDFFF:
        continue
    c = chr(cp)
    d = ud.decomposition(c)
    if d and not d.startswith("<"):
        nfd = ud.normalize("NFD", c)
        if hf.decode(hf(nfd, add_special_tokens=False)["input_ids"]) != ud.normalize("NFC", nfd):
            comp_second.add(ord(nfd[-1]))

# sanity: where Python's table knows the code point, it must agree with the probe
py_l = {cp for cp in range(0x110000) if ud.category(chr(cp))[0] == "L"}
py_n = {cp for cp in range(0x110000) if ud.category(chr(cp))[0] == "N"}
print("L: rust", len(L), "python", len(py_l), "rust-only", len(L - py_l), "python-only", len(py_l - L))
print("N: rust", len(N), "python", len(py_n), "rust-only", len(N - py_n), "python-only", len(py_n - N))
print("ccc ignored:", len(ccc_ignored), "composition-second barriers:", sorted(map(hex, comp_second)))
assert not (py_l - L) and not (py_n - N), "Python knows letters/digits the Rust side lacks: investigate"

with open(OUT, "w") as f:
    f.write('"""GENERATED by tests/data/make_qwen2_unicode_tables.py -- do not edit.\n\n')
    f.write(f"Probed from transformers {transformers.__version__} / tokenizers {tokenizers.__version__} "
            f"against Python unicodedata {ud.unidata_version}.\nSee that script for why these are frozen.\n")
    f.write('"""\n')
    f.write(f"PYTHON_UNIDATA_VERSION = {ud.unidata_version!r}  # the unicodedata these deltas were measured against\n\n")
    f.write("#: \\p{L} / \\p{N} exactly as the Rust regex engine sees them (inclusive ranges).\n")
    f.write("L_RANGES = (\n")
    for a, b in ranges(L):
        f.write(f"    (0x{a:x}, 0x{b:x}),\n")
    f.write(")\nN_RANGES = (\n")
    for a, b in ranges(N):
        f.write(f"    (0x{a:x}, 0x{b:x}),\n")
    f.write(")\n\n")
    f.write("#: Code points the Rust NFC treats as combining-class 0 although unicodedata says otherwise,\n")
    f.write("#: plus second members of compositions the Rust side lacks. NFC is applied around them.\n")
    f.write("NFC_BARRIERS = (\n")
    for a, b in ranges(ccc_ignored | comp_second):
        f.write(f"    (0x{a:x}, 0x{b:x}),\n")
    f.write(")\n")
print("wrote", OUT)