#!/usr/bin/env python3
"""Compare normalizer models on the same sentences.

The shipped model is a 0.6B task-specific fine-tune that needs no prompt but
miscounts large numbers and occasionally drops a word. This runs any causal LM
over the same battery so the trade is measurable rather than argued about:

    python bench.py                                   # the fine-tune
    python bench.py --model models/PLLuM-4B-instruct-2512 --prompted
    python bench.py --model … --prompted --save pllum.json
    python bench.py --compare fine-tune.json pllum.json

`--prompted` switches from the fine-tune's bare "text in, text out" contract to
an instruction + few-shot prompt, which a general instruct model needs. Digit
tokens are banned from the output either way, and decoding is greedy, so a
difference between two runs is the model and not the sampling.

Automatic flags are cheap heuristics, not a grade: `digits` means an
un-normalized digit survived, `lost` means a word from the input has no
plausible inflected match in the output. Read the text before believing either.
"""

import argparse
import json
import re
import sys
import time
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

import normalizer as nz   # noqa: E402  (reuses its protect/restore and banned list)

SYSTEM_PROMPTED = (
    "Jesteś narzędziem do normalizacji tekstu dla syntezatora mowy. "
    "Zamieniasz wszystkie liczby zapisane cyframi na słowa, w poprawnej formie "
    "gramatycznej wynikającej z kontekstu zdania. Nie zmieniasz, nie dodajesz "
    "ani nie usuwasz żadnych innych słów. Zwracasz wyłącznie przekształcone "
    "zdanie, bez komentarza."
)

# Few-shot examples chosen to pin down the three things a general model gets
# wrong without them: the case must follow the governing word, the rest of the
# sentence must survive verbatim, and the answer must be the sentence alone.
FEWSHOT = [
    ("Koncert dla 8250 osób.",
     "Koncert dla ośmiu tysięcy dwustu pięćdziesięciu osób."),
    ("Populacja Polski to 38 000 000 mieszkańców.",
     "Populacja Polski to trzydzieści osiem milionów mieszkańców."),
    ("Mówimy o 25 psach.",
     "Mówimy o dwudziestu pięciu psach."),
    ("Wygrał 1247 złotych.",
     "Wygrał tysiąc dwieście czterdzieści siedem złotych."),
]

BATTERY = [
    "Sprawa dotyczy 1 350 000 ludzi.",
    "Koncert dla 8250 osób.",
    "Budżet wynosi 12 500 000 złotych.",
    "Populacja Polski to 38 000 000 mieszkańców.",
    "Zapłacili 1 500 000 000 dolarów.",
    "Zostało 999 999 sztuk.",
    "Numer seryjny 1234567.",
    "Ponad 100 000 widzów na stadionie.",
    "Mówimy o 1 350 000 ludziach.",
    "Z 8250 osobami nie da się dyskutować.",
    "Firma zatrudnia 2 500 pracowników.",
    "Dystans 42 195 metrów.",
    "Cena spadła z 15 700 000 do 8 250 000.",
    "Było tam 25 psów.",
    "Mówimy o 25 psach.",
    "Przyjechał z 25 psami.",
    "Przyglądał się 25 psom.",
    "Nie ma tu 25 psów.",
    "Urodził się 3 maja 1791 roku.",
    "Wzrost o 16 procent w 2026 roku.",
    "Spotkanie o 14:30.",
    "Godzina 23:56.",
    "Radio 95,7 FM.",
    "Około 3,5 metra.",
    "Mieszka pod 12/45.",
    "Wygrał 1247 złotych.",
    "To kosztuje 1999 dolarów i 99 centów.",
    "Mam 3 koty i 25 psów.",
    "Zostało 5 minut.",
    "Temperatura opon w porządku. Przód 85, tył 75.",
    "Louis Chiron miał 57 lat, kiedy wygrał Grand Prix Monako.",
    "Zwykłe zdanie bez żadnych liczb.",
]

_WORD_RE = re.compile(r'[^\W\d_]+', re.UNICODE)


def lost_words(src: str, out: str) -> list:
    """Input words with no plausible match in the output.

    Polish inflects, so an exact comparison is useless — `ludziach` legitimately
    becomes `ludziach` and nothing else, but `psach`/`psami` share a stem. A
    five-character prefix is a crude stand-in for a lemma and is good enough to
    notice a noun that vanished entirely, which is the failure worth catching.
    """
    def stems(text):
        return {w.lower()[:5] for w in _WORD_RE.findall(text)}
    missing = stems(src) - stems(out)
    return sorted(w for w in missing if len(w) >= 4)


def load_prompted(model_dir: str, gpu_gib: float = None,
                  system: str = None, fewshot: list = None):
    """A general instruct model: system prompt, few-shot, greedy, digits banned.

    `gpu_gib` caps GPU use and spills the rest to RAM. Needed for the int4
    "awq" builds: transformers decompresses compressed-tensors weights back to
    bfloat16 at load, so an 11B model wants ~22 GB whatever the file size says.
    Offloading is slow but it is the difference between measuring a model and
    guessing about it.
    """
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tok = AutoTokenizer.from_pretrained(model_dir)
    kwargs = {"dtype": torch.bfloat16, "device_map": "auto"}
    if gpu_gib == 0:
        # Everything in RAM. Slow, but the only way to run a compressed-tensors
        # int4 build in transformers: its decompression hook expands weights to
        # bfloat16 in place, and it cannot read the meta tensors that accelerate
        # leaves behind when max_memory forces a partial offload.
        kwargs["device_map"] = {"": "cpu"}
    elif gpu_gib:
        kwargs["max_memory"] = {0: f"{gpu_gib}GiB", "cpu": "80GiB"}
    model = AutoModelForCausalLM.from_pretrained(model_dir, **kwargs).eval()

    banned = set()
    for tid in range(len(tok)):
        piece = tok.decode([tid]).strip()
        if piece and piece.isdigit():
            banned.add(tid)
    banned = sorted(banned)
    print(f"  {len(banned)} digit tokens banned; device "
          f"{next(model.parameters()).device}", file=sys.stderr)

    def run(text: str) -> str:
        messages = [{"role": "system", "content": system or SYSTEM_PROMPTED}]
        for user, assistant in (fewshot or FEWSHOT):
            messages.append({"role": "user", "content": user})
            messages.append({"role": "assistant", "content": assistant})
        messages.append({"role": "user", "content": text})
        prompt = tok.apply_chat_template(messages, tokenize=False,
                                         add_generation_prompt=True)
        inputs = tok(prompt, return_tensors="pt").to(model.device)
        n = inputs["input_ids"].shape[1]
        with torch.no_grad():
            out = model.generate(**inputs, max_new_tokens=nz.MAX_NEW_TOKENS,
                                 do_sample=False,
                                 pad_token_id=tok.pad_token_id or tok.eos_token_id,
                                 suppress_tokens=banned)
        answer = tok.decode(out[0][n:], skip_special_tokens=True).strip()
        # Instruct models like to introduce themselves; keep the last non-empty
        # line, which is the sentence in every observed case.
        lines = [l.strip() for l in answer.splitlines() if l.strip()]
        return lines[-1] if lines else answer

    return run


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model", default=None,
                    help="Model directory (default: the shipped fine-tune)")
    ap.add_argument("--prompted", action="store_true",
                    help="Use an instruction + few-shot prompt (general models)")
    ap.add_argument("--gpu-gib", type=float, default=None, metavar="N",
                    help="Cap GPU memory and spill the rest to RAM (slow, but "
                         "the only way to run a model that does not fit)")
    ap.add_argument("--no-protect", action="store_true",
                    help="Do not shield standalone single digits — normalize the "
                         "text exactly as the model card does")
    ap.add_argument("--save", metavar="FILE", help="Write results as JSON")
    ap.add_argument("--compare", nargs=2, metavar=("A", "B"),
                    help="Print two saved runs side by side and exit")
    args = ap.parse_args()

    if args.compare:
        a, b = (json.loads(Path(f).read_text(encoding="utf-8")) for f in args.compare)
        agree = 0
        for ra, rb in zip(a["results"], b["results"]):
            if ra["out"] == rb["out"]:
                agree += 1
                continue
            print(f"  in : {ra['in']}")
            print(f"  A  : {ra['out']}")
            print(f"  B  : {rb['out']}\n")
        print(f"identical: {agree}/{len(a['results'])}   "
              f"A={a['model']}  B={b['model']}")
        return 0

    if args.prompted:
        if not args.model:
            ap.error("--prompted needs --model")
        run = load_prompted(args.model, args.gpu_gib)
        label = args.model
    else:
        if args.model:
            nz.MODEL_DIR = Path(args.model)
        if not nz.load_model():
            print("model unavailable; refusing to benchmark the fallback",
                  file=sys.stderr)
            return 1
        run = nz.normalize
        label = str(nz.MODEL_DIR)
    if args.no_protect:
        nz.KEEP_RE = re.compile(r'(?!x)x')     # matches nothing
        label += " [no single-digit protection]"

    results, t0 = [], time.time()
    for text in BATTERY:
        t = time.time()
        out = run(text)
        results.append({"in": text, "out": out, "sec": round(time.time() - t, 2),
                        "digits": bool(nz.DIGIT_RE.search(nz.KEEP_RE.sub("", out))),
                        "lost": lost_words(text, out)})
    total = time.time() - t0

    for r in results:
        flags = []
        if r["digits"]:
            flags.append("DIGITS LEFT")
        if r["lost"]:
            flags.append("LOST: " + ", ".join(r["lost"]))
        print(f"  in : {r['in']}")
        print(f"  out: {r['out']}" + (f"   [{'; '.join(flags)}]" if flags else ""))
        print()

    n_dig = sum(r["digits"] for r in results)
    n_lost = sum(bool(r["lost"]) for r in results)
    print(f"{label}\n  {len(results)} sentences in {total:.1f}s "
          f"({total/len(results):.2f}s each)\n"
          f"  {n_dig} with digits left, {n_lost} with a word apparently lost")

    if args.save:
        Path(args.save).write_text(json.dumps(
            {"model": label, "prompted": args.prompted, "seconds": round(total, 1),
             "results": results}, ensure_ascii=False, indent=1), encoding="utf-8")
        print(f"  saved to {args.save}")
    return 0


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