"""Reproduce the tables in this model card.

    pip install transformers datasets
    python benchmark.py            # both tables
    python benchmark.py paradigms  # UniMorph only
    python benchmark.py fertility  # corpora only

Everything it reads is public and pinned: a UniMorph commit, five Hugging Face datasets at a fixed
document count, and five published tokenizers. No part of the pipeline that produced this
vocabulary is involved, which is the point -- the numbers are meant to be checkable by someone who
has never seen it.
"""
import re
import sys
import unicodedata
import urllib.request
from collections import defaultdict

from datasets import load_dataset
from transformers import AutoTokenizer

TOKENIZERS = (
    "picur/picur-tokenizer",
    "NYTK/PULI-GPT-3SX",
    "SZTAKI-HLT/hubert-base-cc",
    "Qwen/Qwen3-8B",
    "FacebookAI/xlm-roberta-base",
)

UNIMORPH = ("https://raw.githubusercontent.com/unimorph/hun/"
            "20d5e04211bbab3dd689c99637356a8eab7c5fd7/hun")
LEMMAS = 20_000

CORPORA = (
    ("lazos/gutenberg-hu", None, "data", "text"),
    ("lazos/gazette-hu", None, "data", "text"),
    ("SZTAKI-HLT/HunSum-2-abstractive", None, "data", "article"),
    ("open-web-math/open-web-math", None, "data", "text"),
    ("bigcode/starcoderdata", None, "python", "content"),
)
DOCUMENTS = 2_000
MIN_CHARS = 200
WORD = re.compile(r"\w+", re.UNICODE)


def load():
    found = {}
    for name in TOKENIZERS:
        try:
            found[name] = AutoTokenizer.from_pretrained(name)
        except Exception as error:            # a gated model is skipped, not fatal
            print(f"skip {name}: {type(error).__name__}", file=sys.stderr)
    return found


def paradigms():
    """Does an inflected form begin with its lemma, and is that lemma a single row?"""
    with urllib.request.urlopen(UNIMORPH, timeout=120) as handle:
        raw = handle.read().decode("utf-8")
    table = defaultdict(list)
    for line in raw.splitlines():
        parts = line.split("\t")
        if len(parts) < 3 or not parts[0] or not parts[1]:
            continue
        if len(table) >= LEMMAS and parts[0] not in table:
            continue
        if parts[1] != parts[0]:
            table[parts[0]].append(parts[1])
    pairs = sum(len(v) for v in table.values())
    print(f"\nUniMorph hun: {len(table):,} lemmas, {pairs:,} lemma/form pairs")
    print(f"{'tokenizer':32s} {'stem row':>9s} {'prefix':>8s} {'lemma=1':>8s} {'tok/form':>9s}")
    for name, tok in load().items():
        cache = {}

        def cut(word, tok=tok, cache=cache):
            if word not in cache:
                cache[word] = tuple(tok.tokenize(" " + word))
            return cache[word]

        rows = shared = whole = cost = 0
        for lemma, forms in table.items():
            head = cut(lemma)
            entire = len(head) == 1
            whole += entire
            for form in forms:
                piece = cut(form)
                keeps = len(piece) >= len(head) and piece[:len(head)] == head
                shared += keeps
                rows += keeps and entire
                cost += len(piece)
        print(f"{name:32s} {rows/pairs:8.1%} {shared/pairs:8.1%} "
              f"{whole/len(table):8.1%} {cost/pairs:9.2f}")


def fertility():
    """Characters a token buys, on a fixed prefix of each dataset."""
    loaded = load()
    for repo, config, folder, column in CORPORA:
        stream = load_dataset(repo, config, data_dir=folder, split="train", streaming=True)
        texts = []
        for row in stream:
            text = (row.get(column) or "").strip()
            if len(text) >= MIN_CHARS:
                texts.append(unicodedata.normalize("NFC", text))
            if len(texts) >= DOCUMENTS:
                break
        chars = sum(map(len, texts))
        words = [w for t in texts for w in WORD.findall(t)]
        print(f"\n{repo}  {len(texts):,} documents, {chars/1e6:.2f} Mchar")
        for name, tok in loaded.items():
            total = sum(len(tok.encode(t, add_special_tokens=False)) for t in texts)
            cost = {w: len(tok.encode(" " + w, add_special_tokens=False)) for w in set(words)}
            counts = [cost[w] for w in words]
            print(f"    {name:32s} {chars/total:6.3f} chars/token   "
                  f"{sum(counts)/len(counts):4.2f} pieces/word   "
                  f"{sum(n >= 4 for n in counts)/len(counts):5.1%} need 4+")


if __name__ == "__main__":
    wanted = sys.argv[1] if len(sys.argv) > 1 else "all"
    if wanted in ("all", "paradigms"):
        paradigms()
    if wanted in ("all", "fertility"):
        fertility()