"""Evaluate a tokenizer.json on the four India pages, exactly as the assignment defines it.

faithful_unit = one contiguous Unicode letter/mark/number run OR one visible non-space
                punctuation/symbol character
fertility(L)  = token_count(L) / faithful_unit_count(L)
score         = 1000 / (max_fertility - min_fertility)
hindi_penalty = exp(max(0, hindi_fertility / 1.2 - 1));  adjusted = score / hindi_penalty

Also checks faithfulness: decode(encode(text)) must keep every non-whitespace character.

Usage: python evaluate_tokenizer.py [tokenizer.json] [--json metrics.json]
"""
import json
import math
import pathlib
import sys
from fractions import Fraction

from tokenizers import Tokenizer

from eval_common import units, visible

HERE = pathlib.Path(__file__).resolve().parent
LANGS = ["en", "hi", "te", "mai"]


def count_tokens(tok, text):
    # one call on the whole file, exactly as the course's evaluator does
    return len(tok.encode(text).ids)


def evaluate(path, corpus_dir=HERE / "corpus"):
    tok = Tokenizer.from_file(str(path))
    rows = {}
    for lang in LANGS:
        text = (corpus_dir / f"{lang}.faithful.txt").read_text(encoding="utf-8")
        enc = tok.encode(text)
        n_tok = len(enc.ids)
        n_unit = units(text)
        dec = tok.decode(enc.ids, skip_special_tokens=False)
        faithful = visible(dec) == visible(text)
        rows[lang] = {"tokens": n_tok, "units": n_unit, "fertility": n_tok / n_unit,
                      "fertility_exact": f"{n_tok}/{n_unit}", "faithful": faithful}
    fr = {l: Fraction(r["tokens"], r["units"]) for l, r in rows.items()}
    spread = max(fr.values()) - min(fr.values())
    score = float(1000 / spread) if spread else math.inf
    hi = rows["hi"]["fertility"]
    penalty = math.exp(max(0.0, hi / 1.2 - 1))
    return {
        "tokenizer": str(path), "vocab_size": tok.get_vocab_size(with_added_tokens=True),
        "languages": rows,
        "max_fertility_lang": max(fr, key=fr.get), "min_fertility_lang": min(fr, key=fr.get),
        "spread_exact": f"{spread.numerator}/{spread.denominator}", "spread": float(spread),
        "raw_score": score, "hindi_penalty": penalty, "hindi_adjusted_score": score / penalty,
        "all_below_1_2": all(r["fertility"] <= 1.2 for r in rows.values()),
        "all_faithful": all(r["faithful"] for r in rows.values()),
    }


if __name__ == "__main__":
    args = [a for a in sys.argv[1:] if not a.startswith("--")]
    m = evaluate(pathlib.Path(args[0]) if args else HERE / "tokenizer.json")
    print(f"vocab size {m['vocab_size']:,}")
    print(f"{'lang':5}{'tokens':>10}{'units':>10}{'fertility':>14}  faithful")
    for l, r in m["languages"].items():
        print(f"{l:5}{r['tokens']:>10,}{r['units']:>10,}{r['fertility']:>14.9f}  {r['faithful']}")
    print(f"spread {m['spread']:.3e} ({m['spread_exact']})  raw score {m['raw_score']:,.2f}  "
          f"hindi penalty {m['hindi_penalty']:.6f}  adjusted {m['hindi_adjusted_score']:,.2f}")
    if "--json" in sys.argv:
        out = sys.argv[sys.argv.index("--json") + 1]
        pathlib.Path(out).write_text(json.dumps(m, ensure_ascii=False, indent=2), encoding="utf-8")