# /// script
# requires-python = ">=3.11,<3.14"
# dependencies = [
#   "torch>=2.6,<2.8",
#   "datasets>=3.0",
#   "numpy<3",
#   "onnxruntime>=1.20",
#   "transformers>=4.46",
#   "tokenizers>=0.20",
# ]
# ///
"""Scores autocomplete models the way people use them: type real text, pause, and see if the gray suggestion is right.

Every model sees the same pauses in the same text and is decoded with the same rule as offline.page:
follow the model's top pick while the running product of probabilities stays above a cutoff, and show
the text up to the last finished word. A suggestion is right when the text really continues that way.

Two scenarios:
  typing   the cursor is at the end of what's written so far.
  editing  the cursor is mid-sentence and real text follows a gap of 2-6 words. Only fill-in-the-middle
           models are given that text; the rest see only what's before the cursor.

Usage:
  uv run eval.py --model tinynext-v1 --onnx model_quint8.onnx
  uv run eval.py --model smollm2-135m --hf HuggingFaceTB/SmolLM2-135M
  uv run eval.py --table                  # compare every metrics/*.json
"""

import argparse
import hashlib
import json
import random
import re
import statistics
import time
from collections.abc import Callable
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np

HERE = Path(__file__).parent
EVAL_SET = HERE / "eval_set.json"
METRICS = HERE / "metrics"

CUTOFFS = (0.3, 0.4, 0.5, 0.6, 0.7)
MAIN_CUTOFF = 0.5  # what offline.page uses
PAUSE_EVERY = 12  # characters between pauses
MIN_SHOWN_CHARS = 4  # offline.page hides shorter suggestions
MAX_WORDS = 6
MAX_CHARS = 48
DOMAINS = ("articles", "messages", "email")
SCENARIOS = ("typing", "editing")
PUNCTUATION = ".,!?;:"

# ---- Eval text -------------------------------------------------------------------------------------------------
# Kept apart from training: Wikipedia isn't in either model's data, OpenAssistant validation isn't in SmolTalk's
# training split, and the emails come from the senders the original enron-autocomplete held out for its test set.

CUT_REPLY = re.compile(
    r"(?im)^\s*(?:-+\s*original message|-+\s*forwarded by|-+\s*inline attachment|from:\s|to:\s|sent:\s"
    r"|subject:\s|cc:\s|.*\bwrote:\s*$|_{5,}|-{5,}|={5,}|\*{5,}|\d{1,2}/\d{1,2}/\d{2,4}\s+\d{1,2}:\d{2})"
)
SENSITIVE = re.compile(r"(?i)https?://|www\.|\S+@\S+|\(?\b\d{3}\)?[-. ]\d{3}[-. ]\d{4}\b|\b\d{3}-\d{2}-\d{4}\b")
BULK_SENDER = re.compile(r"(?i)announce|newsletter|mailer|outlook\.team|^all[._]|general\.|chairman|enron\.(?:on|net)")
SMART_CHARS = str.maketrans({"‘": "'", "’": "'", "“": '"', "”": '"', "–": "-", "—": "-", "\t": " "})


def clean_email(body: str) -> list[str]:
    """The same cleaning enron-autocomplete was trained with."""
    text = body.replace("\r\n", "\n").replace("=20", " ").replace("=\n", "").translate(SMART_CHARS)
    text = "\n".join(line for line in text.split("\n") if not line.lstrip().startswith(">"))
    cut = CUT_REPLY.search(text)
    if cut:
        text = text[: cut.start()]
    text = re.sub(r"(?m)^-- *\n[\s\S]*", "", text)
    paragraphs = []
    for block in re.split(r"\n\s*\n", text):
        lines = [re.sub(r" +", " ", line).strip() for line in block.split("\n") if line.strip()]
        if not lines:
            continue
        joined = lines[0]
        for previous, line in zip(lines, lines[1:]):
            joined += (" " if len(previous) >= 50 else "\n") + line
        if SENSITIVE.search(joined) or sum(c.isalpha() for c in joined) < 0.6 * len(joined):
            continue
        paragraphs.append(joined)
    return paragraphs


def is_held_out_sender(sender: str) -> bool:
    """enron-autocomplete put senders whose hash bucket is under 5 in its test split."""
    sender = sender.lower()
    if "@enron.com" not in sender or BULK_SENDER.search(sender):
        return False
    return int(hashlib.sha1(sender.encode()).hexdigest(), 16) % 100 < 5


def ascii_paragraphs(text: str, min_chars: int = 40) -> list[str]:
    text = text.translate(SMART_CHARS).encode("ascii", "ignore").decode()
    return [p.strip() for p in re.split(r"\n\s*\n|\n", text) if len(p.strip()) >= min_chars]


def build_eval_set(docs_per_domain: int = 80) -> dict[str, list[list[str]]]:
    from datasets import load_dataset

    sets: dict[str, list[list[str]]] = {domain: [] for domain in DOMAINS}
    wiki = load_dataset(
        "wikimedia/wikipedia", data_files="20231101.en/train-00040-of-00041.parquet", split="train", streaming=True
    )
    for row in wiki:
        paragraphs = [p for p in ascii_paragraphs(row["text"], 80) if not p.endswith(":")][:6]
        if len(paragraphs) >= 3:
            sets["articles"].append(paragraphs)
        if len(sets["articles"]) >= docs_per_domain:
            break

    for row in load_dataset("OpenAssistant/oasst1", split="validation"):
        if row["lang"] == "en" and "```" not in row["text"]:
            paragraphs = ascii_paragraphs(row["text"])[:6]
            if paragraphs:
                sets["messages"].append(paragraphs)
        if len(sets["messages"]) >= docs_per_domain:
            break

    for row in load_dataset("corbt/enron-emails", split="train", streaming=True):
        if is_held_out_sender(row["from"] or ""):
            paragraphs = [p for p in clean_email(row["body"] or "") if len(p) >= 40][:6]
            if paragraphs:
                sets["email"].append(paragraphs)
        if len(sets["email"]) >= docs_per_domain:
            break
    return sets


def load_eval_set() -> dict[str, list[list[str]]]:
    if not EVAL_SET.exists():
        EVAL_SET.write_text(json.dumps(build_eval_set(), indent=1))
    return json.loads(EVAL_SET.read_text())


@dataclass
class Pause:
    domain: str
    scenario: str
    before: str  # the document up to the cursor
    after: str  # text after the cursor that a model may see (empty when typing)
    truth: str  # what the person actually typed next


def pauses(eval_set: dict[str, list[list[str]]], max_per_domain: int) -> list[Pause]:
    rng = random.Random(0)
    found: list[Pause] = []
    for domain in DOMAINS:
        in_domain: list[Pause] = []
        for paragraphs in eval_set[domain]:
            for index, paragraph in enumerate(paragraphs):
                earlier = "\n".join(paragraphs[:index])
                for at in range(PAUSE_EVERY, len(paragraph) - 8, PAUSE_EVERY):
                    before = (earlier + "\n" if earlier else "") + paragraph[:at]
                    truth = paragraph[at:]
                    in_domain.append(Pause(domain, "typing", before, "", truth))
                    # Editing: a gap of 2-6 words, then the rest of the paragraph as it was written.
                    gap_words = rng.randint(2, 6)
                    gap = re.match(r"(\S*\s+){%d}" % gap_words, truth)
                    if gap and len(truth) > gap.end() + 10:
                        in_domain.append(Pause(domain, "editing", before, truth[gap.end():], truth))
        rng.shuffle(in_domain)
        found += [p for p in in_domain if p.scenario == "typing"][:max_per_domain]
        found += [p for p in in_domain if p.scenario == "editing"][:max_per_domain]
    return found


# ---- Decoding -------------------------------------------------------------------------------------------------


@dataclass
class Trajectory:
    """Suggestions that end at a word boundary, with the model's confidence in each. Greedy decoding gives the
    same path for every cutoff, so one decode answers all of them."""

    candidates: list[tuple[str, float]] = field(default_factory=list)
    milliseconds: float = 0.0

    def at(self, cutoff: float) -> str:
        best = ""
        for text, confidence in self.candidates:
            if confidence >= cutoff:
                best = text
        return best.rstrip()


class Model:
    name: str
    sees_after: bool = False

    def suggest(self, before: str, after: str) -> Trajectory:
        raise NotImplementedError


class CharGRU(Model):
    """enron-autocomplete / tinynext v1: one character per step, reading only the current paragraph."""

    HIDDEN, END, NEWLINE, FIRST, OFFSET = 512, 1, 2, 3, 29

    def __init__(self, name: str, onnx_path: str):
        import onnxruntime as ort

        self.name = name
        self.session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
        self.read_text: str | None = None
        self.reading = None

    def step(self, char_id: int, state):
        logits, state = self.session.run(None, {"input_id": np.array([char_id], np.int64), "state": state})
        return logits[0], state

    def encode(self, char: str) -> int | None:
        if char == "\n":
            return self.NEWLINE
        code = ord(char)
        return code - self.OFFSET if 32 <= code <= 126 else None

    def read(self, text: str):
        if self.read_text is None or not text.startswith(self.read_text):
            self.reading = self.step(self.END, np.zeros((1, self.HIDDEN), np.float32))
            self.read_text = ""
        for char in text[len(self.read_text):]:
            char_id = self.encode(char)
            if char_id is not None:
                self.reading = self.step(char_id, self.reading[1])
        self.read_text = text
        return self.reading

    def suggest(self, before: str, after: str) -> Trajectory:
        began = time.perf_counter()
        logits, state = self.read(before.split("\n")[-1])
        trajectory, run, confidence, words = Trajectory(), "", 1.0, 0
        while len(run) < MAX_CHARS and confidence >= min(CUTOFFS):
            probabilities = np.exp(logits - logits.max())
            probabilities /= probabilities.sum()
            char_id = int(probabilities.argmax())
            if char_id < self.FIRST:
                if run:
                    trajectory.candidates.append((run, confidence))
                break
            char = chr(char_id + self.OFFSET)
            if char == " ":
                trajectory.candidates.append((run, confidence))
                words += 1
                if words >= MAX_WORDS:
                    break
            confidence *= float(probabilities[char_id])
            run += char
            if char in PUNCTUATION:
                trajectory.candidates.append((run, confidence))
            logits, state = self.step(char_id, state)
        trajectory.milliseconds = (time.perf_counter() - began) * 1000
        return trajectory


class TokenLM(Model):
    """A word-piece model. `next_logits(ids, cache)` runs one forward pass and returns (last logits, cache).

    The word being typed is "healed": the prompt stops before it, and the first generated pieces must spell it
    out, so a half-typed word like "resig" can still become "resignation" in one piece.
    """

    def __init__(self, name: str, next_logits: Callable, encode: Callable[[str], list[int]], token_texts: list[str],
                 stop_ids: set[int], prompt: Callable[[list[int], list[int]], list[int]] | None = None,
                 context_chars: int = 1500, suffix_chars: int = 300):
        self.name = name
        self.next_logits = next_logits
        self.encode = encode
        self.token_texts = token_texts
        self.stop_ids = stop_ids
        self.prompt = prompt
        self.sees_after = prompt is not None
        self.context_chars = context_chars
        self.suffix_chars = suffix_chars
        lengths = np.array([len(t) for t in token_texts])
        self.nonempty = lengths > 0

    def allowed(self, need: str) -> np.ndarray:
        """Pieces consistent with the rest of the half-typed word."""
        return np.array([bool(t) and (need.startswith(t) or t.startswith(need)) for t in self.token_texts])

    def suggest(self, before: str, after: str) -> Trajectory:
        began = time.perf_counter()
        before = before[-self.context_chars:]
        # The half-typed word keeps its leading space, since word pieces start with one. A line break stays in the prompt.
        space, newline = before.rfind(" "), before.rfind("\n")
        if space > newline:
            stem, need = before[:space], before[space:]
        else:
            stem, need = before[: newline + 1], before[newline + 1:]
        prefix_ids = self.encode(stem)
        if self.prompt:
            ids = self.prompt(prefix_ids, self.encode(after[: self.suffix_chars]) if after else [])
        else:
            ids = prefix_ids
        logits, cache = self.next_logits(ids, None)

        trajectory, generated, confidence = Trajectory(), "", 1.0
        for _ in range(16):
            logits = logits.astype(np.float64)
            if need:
                mask = self.allowed(need)
                if not mask.any():
                    break
                logits = np.where(mask, logits, -np.inf)
            probabilities = np.exp(logits - logits.max())
            probabilities /= probabilities.sum()
            token = int(probabilities.argmax())
            text = self.token_texts[token]
            if token in self.stop_ids or "\n" in text:
                if generated:
                    trajectory.candidates.append((generated, confidence))
                break
            if need:
                consumed = min(len(need), len(text))
                text, need = text[consumed:], need[consumed:]
            if text[:1].isspace() and generated:
                trajectory.candidates.append((generated, confidence))
            confidence *= float(probabilities[token])
            generated += text
            if generated.rstrip()[-1:] in PUNCTUATION and generated.strip():
                trajectory.candidates.append((generated, confidence))
            if confidence < min(CUTOFFS) or len(generated) > MAX_CHARS or len(generated.split()) > MAX_WORDS:
                break
            logits, cache = self.next_logits([token], cache)
        trajectory.milliseconds = (time.perf_counter() - began) * 1000
        return trajectory


def hf_causal_lm(name: str, repo: str) -> TokenLM:
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    device = "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"
    tokenizer = AutoTokenizer.from_pretrained(repo)
    model = AutoModelForCausalLM.from_pretrained(repo, torch_dtype=torch.float32).to(device).eval()

    @torch.no_grad()
    def next_logits(ids, cache):
        out = model(torch.tensor([ids], device=device), past_key_values=cache, use_cache=True)
        return out.logits[0, -1].float().cpu().numpy(), out.past_key_values

    texts = [tokenizer.decode([i]) for i in range(len(tokenizer))]
    stops = {tokenizer.eos_token_id} if tokenizer.eos_token_id is not None else set()
    encode = lambda text: tokenizer.encode(text, add_special_tokens=False) or [tokenizer.eos_token_id]
    return TokenLM(name, next_logits, encode, texts, stops)


def onnx_token_lm(name: str, folder: str, fim: bool) -> TokenLM:
    """A tinynext v2 export: graph.onnx + weights.f32 (compressed, expanded) or model.onnx, plus tokenizer.json."""
    import onnxruntime as ort
    from tokenizers import Tokenizer

    folder = Path(folder)
    graph = folder / "graph.onnx" if (folder / "graph.onnx").exists() else folder / "model.onnx"
    session = ort.InferenceSession(str(graph), providers=["CPUExecutionProvider"])
    config = json.loads((folder / "config.json").read_text())
    tokenizer = Tokenizer.from_file(str(folder / "tokenizer.json"))
    empty = np.zeros((config["layers"], 1, config["heads"], 0, config["dim"] // config["heads"]), np.float32)

    def next_logits(ids, cache):
        past_k, past_v = cache if cache else (empty, empty)
        logits, k, v = session.run(None, {"input_ids": np.array([ids], np.int64), "past_k": past_k, "past_v": past_v})
        return logits[0], (k, v)

    special = len(config["special_tokens"])
    texts = [tokenizer.decode([i], skip_special_tokens=False) if i >= special else "" for i in range(config["vocab"])]
    encode = lambda text: tokenizer.encode(text).ids or [0]
    prompt = (lambda prefix, suffix: prefix if not suffix else [2, *suffix, 1, *prefix, 3]) if fim else None
    return TokenLM(name, next_logits, encode, texts, set(range(special)), prompt=prompt)


# ---- Scoring ---------------------------------------------------------------------------------------------------


def word_by_word(suggestion: str, truth: str) -> int:
    """Characters you'd keep accepting a word at a time until a word is wrong."""
    kept = 0
    for match in re.finditer(r"\s*\S+", suggestion):
        if truth.startswith(suggestion[: match.end()]):
            kept = match.end()
        else:
            break
    return kept


def score(model: Model, all_pauses: list[Pause]) -> dict:
    results: dict = {"model": model.name, "sees_after": model.sees_after, "cutoffs": {}, "examples": [], "latency_ms": {}}
    latencies: list[float] = []
    raw = []
    for index, pause in enumerate(all_pauses):
        trajectory = model.suggest(pause.before, pause.after if model.sees_after else "")
        latencies.append(trajectory.milliseconds)
        raw.append((pause, trajectory))
        if index % 500 == 0:
            print(f"  {model.name}: {index}/{len(all_pauses)} pauses, median {statistics.median(latencies):.1f} ms", flush=True)

    for cutoff in CUTOFFS:
        table: dict = {}
        for domain in DOMAINS:
            for scenario in SCENARIOS:
                rows = [(p, t.at(cutoff)) for p, t in raw if p.domain == domain and p.scenario == scenario]
                if not rows:
                    continue
                shown = [(p, s) for p, s in rows if len(s) >= MIN_SHOWN_CHARS]
                right = [(p, s) for p, s in shown if p.truth.startswith(s)]
                multi = [(p, s) for p, s in shown if len(s.split()) >= 2]
                multi_right = [(p, s) for p, s in multi if p.truth.startswith(s)]
                typed = len(rows) * PAUSE_EVERY
                table[f"{domain}/{scenario}"] = {
                    "pauses": len(rows),
                    "shown": round(len(shown) / len(rows), 4),
                    "right": round(len(right) / max(1, len(shown)), 4),
                    "avg_right_chars": round(sum(len(s) for _, s in right) / max(1, len(right)), 2),
                    "multiword_shown": round(len(multi) / len(rows), 4),
                    "multiword_right": round(len(multi_right) / max(1, len(multi)), 4),
                    "saved_with_tab": round(sum(len(s) for _, s in right) / typed, 4),
                    "saved_word_by_word": round(sum(word_by_word(s, p.truth) for p, s in shown) / typed, 4),
                }
        results["cutoffs"][str(cutoff)] = table

    for pause, trajectory in raw[:: max(1, len(raw) // 60)]:
        suggestion = trajectory.at(MAIN_CUTOFF)
        if len(suggestion) >= MIN_SHOWN_CHARS:
            results["examples"].append({
                "domain": pause.domain, "scenario": pause.scenario, "before": pause.before[-60:],
                "suggestion": suggestion, "truth": pause.truth[: len(suggestion) + 12], "right": pause.truth.startswith(suggestion),
            })
    results["latency_ms"] = {
        "p50": round(statistics.median(latencies), 2),
        "p95": round(float(np.percentile(latencies, 95)), 2),
        "note": "Measured in this eval process, not in a browser.",
    }
    return results


# ---- Report ----------------------------------------------------------------------------------------------------


def table(paths: list[Path]) -> str:
    runs = [json.loads(p.read_text()) for p in paths]
    lines = []
    for scenario in SCENARIOS:
        lines.append(f"\n### {scenario.capitalize()}, {int(MAIN_CUTOFF * 100)}% cutoff, suggestions of {MIN_SHOWN_CHARS}+ characters\n")
        lines.append("| Model | Domain | Shown | Right | Avg right | 2+ words shown | 2+ words right | Saved (Tab) | Saved (word by word) |")
        lines.append("|---|---|---|---|---|---|---|---|---|")
        for run in runs:
            for domain in DOMAINS:
                m = run["cutoffs"][str(MAIN_CUTOFF)].get(f"{domain}/{scenario}")
                if not m:
                    continue
                lines.append(
                    f"| {run['model']} | {domain} | {m['shown']:.0%} | {m['right']:.0%} | {m['avg_right_chars']:.1f} ch "
                    f"| {m['multiword_shown']:.0%} | {m['multiword_right']:.0%} | {m['saved_with_tab']:.1%} | {m['saved_word_by_word']:.1%} |"
                )
    lines.append("\n| Model | Latency p50 | p95 |\n|---|---|---|")
    for run in runs:
        lines.append(f"| {run['model']} | {run['latency_ms']['p50']} ms | {run['latency_ms']['p95']} ms |")
    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", help="Name for the metrics file")
    parser.add_argument("--onnx", help="A character-level GRU exported like enron-autocomplete")
    parser.add_argument("--hf", help="A Hugging Face causal LM to evaluate, e.g. HuggingFaceTB/SmolLM2-135M")
    parser.add_argument("--onnx-lm", help="Folder with a tinynext v2 ONNX export, config.json and tokenizer.json")
    parser.add_argument("--no-fim", action="store_true", help="With --onnx-lm: the model wasn't trained to fill in the middle")
    parser.add_argument("--max-pauses", type=int, default=1000, help="Per domain and scenario")
    parser.add_argument("--table", action="store_true", help="Print a comparison of every metrics/*.json")
    args = parser.parse_args()

    if args.table:
        print(table(sorted(METRICS.glob("*.json"))))
        return
    if args.onnx:
        model: Model = CharGRU(args.model, args.onnx)
    elif args.hf:
        model = hf_causal_lm(args.model, args.hf)
    elif args.onnx_lm:
        model = onnx_token_lm(args.model, args.onnx_lm, fim=not args.no_fim)
    else:
        parser.error("Pass --onnx, --hf or --table. train.py evaluates its own model through score().")
    all_pauses = pauses(load_eval_set(), args.max_pauses)
    print(f"{len(all_pauses)} pauses", flush=True)
    results = score(model, all_pauses)
    METRICS.mkdir(exist_ok=True)
    (METRICS / f"{args.model}.json").write_text(json.dumps(results, indent=1))
    print(table([METRICS / f"{args.model}.json"]))


if __name__ == "__main__":
    main()