"""Standalone nano-proofread benchmark: the published weights vs the best context-free
script, overall and on the context-dependent slice the script cannot do.

The naive script (`naive_fix` in data_proofread.py) is the strongest *context-free*
fixer a developer would write: remove adjacent doubled words, and normalise every
homophone to the most common member of its family. It fixes doubled words and
`could of`->`could have` perfectly, but it cannot know whether `their`/`there`/
`they're` is right — that needs the surrounding words. The hard slice = examples where
the correct homophone is NOT the script's default, i.e. exactly where context decides.

    python eval_nano_proofread.py
"""

from __future__ import annotations

import torch

from data_proofread import OOD_CASES, naive_fix, proofread_pairs
from modeling_nano_proofread import load

_HOLDOUT_SEED = 987_654_321


@torch.no_grad()
def _greedy_batch(model, prompts, max_new=48, device="cpu"):
    """Greedy-decode `max_new` tokens per prompt, then cut at the newline EOS.

    Prompts are bucketed by EXACT byte length so each batch is a dense, padding-free
    (B, L) tensor — bucketing (not padding) keeps the result bit-identical to decoding
    one prompt at a time."""
    enc = [list(p.encode("utf-8")) for p in prompts]
    order = sorted(range(len(prompts)), key=lambda i: len(enc[i]))
    max_seq = model.cfg["max_seq_len"]
    out = [None] * len(prompts)
    i = 0
    while i < len(order):
        L = len(enc[order[i]])
        j = i
        while j < len(order) and len(enc[order[j]]) == L:
            j += 1
        rows = order[i:j]
        toks = torch.tensor([enc[k] for k in rows], dtype=torch.long, device=device)
        for _ in range(max_new):
            nxt = model(toks[:, -max_seq:])[:, -1, :].argmax(-1, keepdim=True)
            toks = torch.cat([toks, nxt], dim=1)
        tail = toks[:, L:L + max_new]
        for r, k in enumerate(rows):
            s = bytes(int(b) & 0xFF for b in tail[r].tolist()).decode("utf-8", "replace")
            out[k] = s.split("\n", 1)[0]
        i = j
    return out


def main():
    m = load()
    n_params = sum(p.numel() for p in m.parameters())
    pairs = proofread_pairs(_HOLDOUT_SEED, 4000)
    preds = _greedy_batch(m, [p for p, _ in pairs])

    mo = nv = 0
    hn = hm = hnv = 0
    for (prompt, target), pred in zip(pairs, preds):
        inp = prompt[:-4]
        mo += pred == target
        nv += naive_fix(inp) == target
        if naive_fix(target) != target:  # correct answer isn't the script default
            hn += 1
            hm += pred == target
            hnv += naive_fix(inp) == target
    N = len(pairs)
    print(f"params : {n_params:,}")
    print(f"held-out: N={N} (seed {_HOLDOUT_SEED})\n")
    print(f"  {'overall':<28}model {mo/N:>6.1%}   naive-script {nv/N:>6.1%}")
    print(f"  context slice (N={hn}):   model {hm/hn:>6.1%}   naive-script {hnv/hn:>6.1%}")
    print("  (context slice = correct homophone is NOT the script's default —")
    print("   exactly where the surrounding words decide the answer)")

    # ---- the honest test: hand-written natural phrases NOT drawn from the frames ----
    ood_preds = _greedy_batch(m, [f"{i} => " for i, _ in OOD_CASES])
    om = onv = 0
    print(f"\nOUT-OF-DISTRIBUTION (hand-written, N={len(OOD_CASES)}):")
    for (inp, gold), pred in zip(OOD_CASES, ood_preds):
        om += pred == gold
        onv += naive_fix(inp) == gold
        flag = "ok " if pred == gold else "XX "
        print(f"  {flag}{inp:<28}-> {pred:<28} (gold: {gold})")
    O = len(OOD_CASES)
    print(f"\n  OOD exact-match: model {om/O:>6.1%}   naive-script {onv/O:>6.1%}")


if __name__ == "__main__":
    main()