#!/usr/bin/env python3
"""Bodham Labs tokenizer cost benchmark: English vs Hindi vs Hinglish vs dev-roman.

What it measures
----------------
For each tokenizer x dataset x language version it counts how many tokens the
same content costs, with NO special tokens (no BOS/EOS/CLS/SEP). Each sentence
is encoded on its own and the counts are summed. All text is NFC-normalised.

Language versions
-----------------
  en        English
  hi        Hindi in Devanagari (technical terms also in Devanagari)
  hinglish  the same Hindi in Roman script, technical terms in English
  hi_mixed  "dev-roman": Devanagari Hindi, technical terms in English Roman

Datasets
--------
  school    texts/school.json, 12 original Class 9-10 sentences, all 4 versions
  flores+   FLORES+ devtest eng_Latn / hin_Deva, first 200 sentences (en, hi
            only), if it can be downloaded (it is gated on the HF Hub)
  in22-gen  AI4Bharat IN22-Gen eng_Latn / hin_Deva, first 200 rows, if it can
            be downloaded (also gated)

Tokenizers
----------
tiktoken o200k_base and cl100k_base, plus Hugging Face tokenizers. For HF
repos only tokenizer/config files are downloaded (allow_patterns), never model
weights. Anything that cannot be loaded is recorded as skipped with its error.

Outputs
-------
  results/token_cost.json   rows + run metadata + skipped list
  results/token_cost.md     Markdown tables
  results/spot_check.txt    token pieces for one sentence, two tokenizers

If results/token_cost.json already exists, its token counts are compared with
the new run before it is overwritten, and the result is printed.

Usage:  python3 bench.py
"""
from __future__ import annotations

import json
import platform
import sys
import traceback
import unicodedata
from dataclasses import dataclass, asdict
from datetime import datetime, timezone
from importlib import metadata
from pathlib import Path
from typing import Callable

ROOT = Path(__file__).resolve().parent
TEXTS = ROOT / "texts"
RESULTS = ROOT / "results"
LANGS = ["en", "hi", "hinglish", "hi_mixed"]
PUBLIC_N = 200

# Only these files are ever fetched from a model repo (no weights).
TOKENIZER_FILES = [
    "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json",
    "added_tokens.json", "tokenizer.model", "spiece.model", "spiece.vocab",
    "sentencepiece.bpe.model", "vocab.json", "vocab.txt", "merges.txt",
    "config.json",
]

HF_TOKENIZERS: list[tuple[str, dict]] = [
    ("Qwen/Qwen2.5-0.5B-Instruct", {}),
    ("Qwen/Qwen3-0.6B", {}),
    ("google/gemma-3-270m", {}),
    ("meta-llama/Llama-3.2-1B", {}),
    ("mistralai/Mistral-7B-v0.1", {}),
    ("sarvamai/sarvam-1", {}),
    # Settings from the IndicBARTSS model card.
    ("ai4bharat/IndicBARTSS", {"do_lower_case": False, "use_fast": False, "keep_accents": True}),
    ("bert-base-multilingual-cased", {}),
]
TIKTOKEN_ENCODINGS = ["o200k_base", "cl100k_base"]

SPOT_CHECK_ID = "s01"
SPOT_CHECK_TOKENIZERS = ["tiktoken/o200k_base", "sarvamai/sarvam-1"]


def nfc(s: str) -> str:
    return unicodedata.normalize("NFC", s.strip())


def short_err(e: BaseException) -> str:
    first = str(e).strip().splitlines()[0] if str(e).strip() else ""
    msg = f"{type(e).__name__}: {first[:300]}"
    if type(e).__name__ == "GatedRepoError":
        msg = "gated repo on the HF Hub; no HF access token in this environment. " + msg
    return msg


def pkg_version(name: str) -> str | None:
    try:
        return metadata.version(name)
    except metadata.PackageNotFoundError:
        return None


# ---------------------------------------------------------------- datasets
def load_school() -> dict[str, list[str]]:
    data = json.loads((TEXTS / "school.json").read_text(encoding="utf-8"))
    out: dict[str, list[str]] = {lang: [] for lang in LANGS}
    for item in data["items"]:
        for lang in LANGS:
            if not item.get(lang):
                raise ValueError(f"school.json item {item.get('id')} missing '{lang}'")
            out[lang].append(nfc(item[lang]))
    return out


def load_flores() -> dict[str, list[str]]:
    from huggingface_hub import hf_hub_download

    out: dict[str, list[str]] = {}
    for lang, code in [("en", "eng_Latn"), ("hi", "hin_Deva")]:
        path = hf_hub_download("openlanguagedata/flores_plus", f"devtest/{code}.jsonl", repo_type="dataset")
        rows = [json.loads(line) for line in Path(path).read_text(encoding="utf-8").splitlines() if line.strip()]
        rows.sort(key=lambda r: int(r["id"]))
        out[lang] = [nfc(r["text"]) for r in rows[:PUBLIC_N]]
    return out


def load_in22() -> dict[str, list[str]]:
    from huggingface_hub import hf_hub_download
    import pandas as pd

    path = hf_hub_download("ai4bharat/IN22-Gen", "data/train-00000-of-00001.parquet", repo_type="dataset")
    df = pd.read_parquet(path).head(PUBLIC_N)
    return {"en": [nfc(x) for x in df["eng_Latn"]], "hi": [nfc(x) for x in df["hin_Deva"]]}


PUBLIC_DATASETS: list[tuple[str, str, str, Callable[[], dict[str, list[str]]]]] = [
    ("flores+_devtest", "openlanguagedata/flores_plus", "CC BY-SA 4.0", load_flores),
    ("in22-gen", "ai4bharat/IN22-Gen", "CC BY 4.0", load_in22),
]


# -------------------------------------------------------------- tokenizers
@dataclass
class Tok:
    name: str
    version: str
    encode: Callable[[str], list[int]]
    pieces: Callable[[list[int]], list[str]]
    unk_id: int | None = None
    # Independent implementations of the same tokenizer, used only to cross-check counts.
    alt_encoders: dict[str, Callable[[str], list[int]]] | None = None


def load_tiktoken(enc_name: str) -> Tok:
    import tiktoken

    enc = tiktoken.get_encoding(enc_name)

    def pieces(ids: list[int]) -> list[str]:
        return [enc.decode_single_token_bytes(i).decode("utf-8", errors="backslashreplace") for i in ids]

    return Tok(
        name=f"tiktoken/{enc_name}",
        version=f"tiktoken {pkg_version('tiktoken')} {enc_name} (n_vocab={enc.n_vocab})",
        encode=enc.encode_ordinary,  # never adds or matches special tokens
        pieces=pieces,
    )


def load_hf(repo: str, kwargs: dict) -> Tok:
    from huggingface_hub import snapshot_download
    from transformers import AutoTokenizer

    local = snapshot_download(repo, allow_patterns=TOKENIZER_FILES)
    weights = [p.name for p in Path(local).iterdir() if p.suffix in {".bin", ".safetensors", ".pth", ".h5", ".msgpack"}]
    if weights:
        raise RuntimeError(f"unexpected weight files downloaded: {weights}")
    revision = Path(local).name  # snapshot dir name is the commit sha
    tok = AutoTokenizer.from_pretrained(local, **kwargs)

    def encode(text: str) -> list[int]:
        return list(tok.encode(text, add_special_tokens=False))

    alt: dict[str, Callable[[str], list[int]]] = {}
    if (Path(local) / "tokenizer.json").exists():
        from tokenizers import Tokenizer
        raw = Tokenizer.from_file(str(Path(local) / "tokenizer.json"))
        alt["tokenizers(tokenizer.json)"] = lambda t, raw=raw: raw.encode(t, add_special_tokens=False).ids
    for sp_file in ("tokenizer.model", "spiece.model"):
        if (Path(local) / sp_file).exists():
            import sentencepiece as spm
            sp = spm.SentencePieceProcessor(model_file=str(Path(local) / sp_file))
            alt[f"sentencepiece({sp_file})"] = lambda t, sp=sp: sp.encode(t)

    return Tok(
        name=repo,
        version=f"{repo}@{revision[:12]} ({type(tok).__name__}, vocab={len(tok)})",
        encode=encode,
        pieces=lambda ids: list(tok.convert_ids_to_tokens(ids)),
        unk_id=tok.unk_token_id,
        alt_encoders=alt,
    )


# ------------------------------------------------------------------- main
def measure(tok: Tok, dataset: str, texts: dict[str, list[str]]) -> list[dict]:
    rows = []
    en_tokens = None
    for lang in LANGS:
        if lang not in texts:
            continue
        sents = texts[lang]
        ids = [i for s in sents for i in tok.encode(s)]
        tokens = len(ids)
        chars = sum(len(s) for s in sents)
        words = sum(len(s.split()) for s in sents)
        if lang == "en":
            en_tokens = tokens
        rows.append({
            "tokenizer": tok.name,
            "tokenizer_version": tok.version,
            "dataset": dataset,
            "lang": lang,
            "n_sentences": len(sents),
            "tokens": tokens,
            "chars": chars,
            "words": words,
            "tokens_per_char": round(tokens / chars, 4),
            "tokens_per_word": round(tokens / words, 4),
            "ratio_vs_en": round(tokens / en_tokens, 4) if en_tokens else None,
            "unk_tokens": ids.count(tok.unk_id) if tok.unk_id is not None else 0,
        })
    return rows


def write_md(payload: dict, path: Path) -> None:
    rows = payload["rows"]
    L: list[str] = []
    L.append("# Tokenizer cost: English vs Hindi vs Hinglish vs dev-roman\n")
    L.append(f"Run at {payload['run_at']} · Python {payload['python']} · "
             + ", ".join(f"{k} {v}" for k, v in payload["packages"].items() if v) + "\n")
    L.append("No special tokens. Each sentence encoded separately, counts summed. Text NFC-normalised.\n")

    for ds in dict.fromkeys(r["dataset"] for r in rows):
        ds_rows = [r for r in rows if r["dataset"] == ds]
        langs = [l for l in LANGS if any(r["lang"] == l for r in ds_rows)]
        n = ds_rows[0]["n_sentences"]
        L.append(f"\n## Ratio vs English: {ds} ({n} sentences)\n")
        L.append("Tokens for the same content divided by English tokens (1.00 = same cost as English).\n")
        hdr = ["Tokenizer", "en tokens"] + [l for l in langs if l != "en"]
        L.append("| " + " | ".join(hdr) + " |")
        L.append("|" + "|".join(["---"] + ["---:"] * (len(hdr) - 1)) + "|")
        for tk in dict.fromkeys(r["tokenizer"] for r in ds_rows):
            by = {r["lang"]: r for r in ds_rows if r["tokenizer"] == tk}
            cells = [tk, str(by["en"]["tokens"])] + [f"{by[l]['ratio_vs_en']:.2f}" for l in langs if l != "en"]
            L.append("| " + " | ".join(cells) + " |")

    L.append("\n## Full results\n")
    hdr = ["Tokenizer", "Dataset", "Lang", "Sents", "Tokens", "Chars", "Words", "Tok/char", "Tok/word", "Ratio vs en"]
    L.append("| " + " | ".join(hdr) + " |")
    L.append("|---|---|---|" + "---:|" * 7)
    for r in rows:
        L.append(f"| {r['tokenizer']} | {r['dataset']} | {r['lang']} | {r['n_sentences']} | {r['tokens']} | "
                 f"{r['chars']} | {r['words']} | {r['tokens_per_char']:.3f} | {r['tokens_per_word']:.3f} | "
                 f"{r['ratio_vs_en']:.2f} |")

    L.append("\n## Tokenizer versions\n")
    for v in dict.fromkeys(r["tokenizer_version"] for r in rows):
        L.append(f"- {v}")

    if payload["skipped"]:
        L.append("\n## Skipped\n")
        L.append("| Item | Kind | Reason |")
        L.append("|---|---|---|")
        for s in payload["skipped"]:
            L.append(f"| {s['name']} | {s['kind']} | {s['error'].replace('|', '/')} |")
    path.write_text("\n".join(L) + "\n", encoding="utf-8")


def spot_check(toks: dict[str, Tok], school: dict[str, list[str]], path: Path) -> str:
    data = json.loads((TEXTS / "school.json").read_text(encoding="utf-8"))
    idx = [it["id"] for it in data["items"]].index(SPOT_CHECK_ID)
    out = [f"Spot check: school sentence {SPOT_CHECK_ID}\n"]
    for name in SPOT_CHECK_TOKENIZERS:
        tok = toks.get(name)
        if tok is None:
            out.append(f"[{name}] not loaded, skipped\n")
            continue
        out.append(f"=== {tok.version}")
        for lang in LANGS:
            text = school[lang][idx]
            ids = tok.encode(text)
            out.append(f"--- {lang} ({len(ids)} tokens): {text}")
            out.append(" | ".join(tok.pieces(ids)))
        out.append("")
    txt = "\n".join(out)
    path.write_text(txt + "\n", encoding="utf-8")
    return txt


def main() -> int:
    RESULTS.mkdir(exist_ok=True)
    skipped: list[dict] = []
    datasets_used: list[dict] = []

    datasets: dict[str, dict[str, list[str]]] = {"school": load_school()}
    datasets_used.append({"name": "school", "source": "texts/school.json", "licence": "CC BY 4.0 (Bodham Labs, original)"})
    for name, repo, licence, loader in PUBLIC_DATASETS:
        try:
            datasets[name] = loader()
            datasets_used.append({"name": name, "source": repo, "licence": licence})
            print(f"[dataset] loaded {name}: {len(datasets[name]['en'])} sentences")
        except Exception as e:
            skipped.append({"name": repo, "kind": "dataset", "error": short_err(e)})
            print(f"[dataset] SKIP {repo}: {short_err(e)}")

    toks: dict[str, Tok] = {}
    loaders: list[tuple[str, Callable[[], Tok]]] = [
        (f"tiktoken/{e}", (lambda e=e: load_tiktoken(e))) for e in TIKTOKEN_ENCODINGS
    ] + [(repo, (lambda repo=repo, kw=kw: load_hf(repo, kw))) for repo, kw in HF_TOKENIZERS]
    for name, loader in loaders:
        try:
            toks[name] = loader()
            print(f"[tokenizer] loaded {toks[name].version}")
        except Exception as e:
            skipped.append({"name": name, "kind": "tokenizer", "error": short_err(e)})
            print(f"[tokenizer] SKIP {name}: {short_err(e)}")
            if "--debug" in sys.argv:
                traceback.print_exc()

    rows: list[dict] = []
    for ds_name, texts in datasets.items():
        for tok in toks.values():
            rows.extend(measure(tok, ds_name, texts))

    crosscheck: list[dict] = []
    for tok in toks.values():
        for impl, enc in (tok.alt_encoders or {}).items():
            for lang in LANGS:
                main_n = sum(len(tok.encode(s)) for s in datasets["school"][lang])
                alt_n = sum(len(enc(s)) for s in datasets["school"][lang])
                crosscheck.append({"tokenizer": tok.name, "impl": impl, "dataset": "school", "lang": lang,
                                   "tokens_main": main_n, "tokens_alt": alt_n, "match": main_n == alt_n})
    n_bad = sum(not c["match"] for c in crosscheck)
    print(f"[crosscheck] {len(crosscheck)} comparisons vs independent implementations, {n_bad} mismatches")

    payload = {
        "run_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "python": platform.python_version(),
        "platform": platform.platform(),
        "packages": {p: pkg_version(p) for p in
                     ["tiktoken", "tokenizers", "transformers", "huggingface_hub", "datasets", "sentencepiece", "protobuf"]},
        "method": {
            "special_tokens": "none (tiktoken encode_ordinary; HF add_special_tokens=False)",
            "unit": "each sentence encoded separately; totals summed",
            "normalisation": "Unicode NFC, leading/trailing whitespace stripped",
            "chars": "Unicode code points",
            "words": "str.split() on whitespace",
            "ratio_vs_en": "tokens(lang) / tokens(en), same tokenizer and dataset",
            "public_set_size": PUBLIC_N,
        },
        "datasets": datasets_used,
        "skipped": skipped,
        "crosscheck": crosscheck,
        "rows": rows,
    }

    out_json = RESULTS / "token_cost.json"
    if out_json.exists():
        old = json.loads(out_json.read_text(encoding="utf-8"))
        key = lambda r: (r["tokenizer"], r["tokenizer_version"], r["dataset"], r["lang"], r["tokens"], r["chars"], r["words"])
        same = sorted(map(key, old["rows"])) == sorted(map(key, rows))
        print(f"[rerun-check] counts vs previous results ({old['run_at']}): {'IDENTICAL' if same else 'DIFFERENT'} "
              f"({len(rows)} rows now, {len(old['rows'])} before)")

    out_json.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    write_md(payload, RESULTS / "token_cost.md")
    print()
    print(spot_check(toks, datasets["school"], RESULTS / "spot_check.txt"))
    print(f"[done] {len(rows)} rows -> {out_json}")
    return 0


if __name__ == "__main__":
    sys.exit(main())