#!/usr/bin/env python3
"""
Slice 20 tokenizer audit (§2 of the brief) and corpus sequence-length
statistics (§3). Runs before any training. If the audit shows catastrophic
Cyrillic / mixed-identifier loss it is the gate to stop.

Loads the frozen dev corpus exactly like slice 19 (same normalization),
so regime A (natural text) is `text_orig` and regime B (punct-stripped) is
`text_strip`.
"""

import json
import os
import re
import statistics

from transformers import AutoTokenizer

import slice19_main as s19

RESULTS_DIR = "/tmp/mvn-s20"
MODEL_DIR = ("/tmp/mvn-s20/hf-cache/models--cointegrated--rubert-tiny/"
             "snapshots/5441c5ea8026d4f6d7505ec004845409f1259fb1")

MIXED_CYRILLIC_LATIN = re.compile(r"[а-яёА-ЯЁ]+[a-zA-Z]+|[a-zA-Z]+[а-яёА-ЯЁ]+")
HAS_CYRILLIC = re.compile(r"[а-яёА-ЯЁ]")
HAS_LATIN = re.compile(r"[a-zA-Z]")
NUMERIC = re.compile(r"[0-9]")
TOKEN_RE = re.compile(r"[^\W\d_]+", re.UNICODE)

SAMPLES = [
    "выключи свет в спальне пожалуйста",
    "turn off the lights",
    "перезапусти сервис mavend",
    "что такое Nexus",
    "включи телевизор, пожалуйста",
    "как дела у Мэйвен",
    "поставь таймер на 5 минут",
    "кто такой Home Assistant",
    "открой настройки устройства ha_cam_12",
    "Покажи статус сервера Proxmox",
    "аутентифицируй на сайте 2fa.ru",
    "сообщи погоду завтра в 18:30",
]


def is_toolish(word):
    # Maven sibling service names / HA-like identifiers: mixed case, digits,
    # underscores, or short Latin words that are not in the vocab as whole
    # words. Rough heuristic for the fragmentation probe.
    return bool(re.search(r"[A-Z0-9_/.-]", word))


def main():
    os.makedirs(RESULTS_DIR, exist_ok=True)
    tok = AutoTokenizer.from_pretrained(MODEL_DIR)
    rows = s19.load_dev()
    texts = {
        "orig": [r["text_orig"] for r in rows],
        "nofinal": [r["text_nofinal"] for r in rows],
        "strip": [r["text_strip"] for r in rows],
    }
    vv = tok.vocab_size
    unk = tok.unk_token_id

    rep = {"model": "cointegrated/rubert-tiny",
           "sha": "5441c5ea8026d4f6d7505ec004845409f1259fb1",
           "tokenizer": type(tok).__name__,
           "vocab_size": vv}

    # per-character stats (natural texts)
    chars = [len(t) for t in texts["orig"]]
    rep["char_len"] = {
        "mean": round(statistics.mean(chars), 2),
        "p50": int(sorted(chars)[len(chars) // 2]),
        "p95": sorted(chars)[int(len(chars) * .95)],
        "p99": sorted(chars)[int(len(chars) * .99)],
        "max": max(chars),
    }

    for view in ("orig", "strip"):
        ids = tok(texts[view], add_special_tokens=True, padding=False,
                  truncation=False)["input_ids"]
        lens = [len(x) for x in ids]
        n_tok = sum(lens)
        n_unk = sum(x.count(unk) for x in ids)
        n_chars = sum(len(t) for t in texts[view])
        rep[view] = {
            "tokens_per_utt_mean": round(n_tok / len(rows), 2),
            "tokens_per_char": round(n_tok / max(n_chars, 1), 4),
            "unk_count": n_unk,
            "unk_rate": round(n_unk / max(n_tok, 1), 5),
            "seq_len_p50": int(sorted(lens)[len(lens) // 2]),
            "seq_len_p90": sorted(lens)[int(len(lens) * .90)],
            "seq_len_p95": sorted(lens)[int(len(lens) * .95)],
            "seq_len_p99": sorted(lens)[int(len(lens) * .99)],
            "seq_len_max": max(lens),
            "above_96": sum(1 for x in lens if x > 96),
            "above_128": sum(1 for x in lens if x > 128),
        }
        rep[view]["p99_plus_margin"] = rep[view]["seq_len_p99"] + 6

    # mixed Cyrillic/Latin behaviour over natural texts
    mixed_words = []
    for t in texts["orig"]:
        for w in t.split():
            if MIXED_CYRILLIC_LATIN.search(w):
                mixed_words.append(w)
    rep["mixed_cyr_lat_rows"] = len({w for w in mixed_words})
    rep["mixed_cyr_lat_stats"] = {"word_count": len(mixed_words),
                                  "unique_words": len(set(mixed_words))}

    # entity/tool-name fragmentation: unique word-like tokens containing a digit
    # or underscore, or non-trivial Latin, and how many BPE/WordPiece pieces they
    # split into. Sample the extremes.
    fragments = []
    vocab = set(tok.get_vocab().keys())
    for t in texts["orig"]:
        # split into "clean" tokens (word chars + _ / digit boundaries)
        for w in re.findall(r"[A-Za-z0-9_]+\b", t):
            w2 = re.sub(r"_\b", "", w)
            if len(w2) < 3 or not is_toolish(w2):
                continue
            n_pieces = len(tok.tokenize(w2).replace("##", "_").rstrip())
            fragments.append((n_pieces, w2))
    frag = sorted(set(fragments))[-30:]
    rep["entity_fragment_examples"] = [
        {"token": w, "pieces": n} for n, w in frag
    ]

    # representative samples: full tokenization
    rep["samples"] = []
    for s in SAMPLES:
        e = tok(s, add_special_tokens=True, padding=False, truncation=False)
        rep["samples"].append({
            "text": s,
            "tokens": tok.convert_ids_to_tokens(e["input_ids"]),
            "pieces": len(e["input_ids"]),
            "unk": e["input_ids"].count(unk),
        })

    with open(os.path.join(RESULTS_DIR, "tokenizer_audit.json"), "w") as f:
        json.dump(rep, f, indent=2, ensure_ascii=False)
    print(json.dumps({k: v for k, v in rep.items() if k not in ("samples",)}, indent=2, ensure_ascii=False))
    print("\n--- samples ---")
    for s in rep["samples"]:
        print(f'{s["pieces"]:>3} unk={s["unk"]}  {s["text"]:50} ->'
              f' {" ".join(s["tokens"])}')


if __name__ == "__main__":
    main()