"""How the exam answers are graded. Copied from the training code (finetune_search.is_correct / verdict and
searchfmt.normalize), standard library only.

  python grade.py google_big.jsonl [more.jsonl ...]     re-grades the "answer" field of each row and prints
                                                        accuracy, overall and per kind

Verdicts: answered_correct, wrong, missed (said "not found" but the results held the answer) for questions with
gold answers; trap_caught / trap_fooled for questions with no gold answer (the right reply is "not found").
Accuracy = (answered_correct + trap_caught) / all questions.
"""
import collections
import json
import re
import string
import sys

NOT_FOUND = "not found"
VERDICTS = ("answered_correct", "missed", "wrong", "trap_caught", "trap_fooled")


def normalize(s):
    # Same normalization as the official SQuAD script: lowercase, no punctuation, no a/an/the.
    s = "".join(ch for ch in s.lower() if ch not in set(string.punctuation))
    s = re.sub(r"\b(a|an|the)\b", " ", s)
    return " ".join(s.split())


def said_not_found(pred):
    return normalize(pred) in ("", normalize(NOT_FOUND))


DIGITS = "0123456789"
STOP = set("""a an the of in on at to for by with from into about between and or but is are was were be been being
    it its this that these those as than then there their his her he she they them we you your our my me i
    which who whom whose what when where why how do does did has have had will would can could should s t""".split())


def inside(text, g):
    """g is in text (both normalized) and not part of a longer number or word. Only a digit can extend a
    number, because snippets glue units and signs to it: "₨277.06", "97°f" and "rs39066ltr" hold 277.06, 97 and
    390.66, while "2,277.08" does not hold 277.08 and "jupiters" does not hold jupiter."""
    def edge(c):
        return "[0-9]" if c in DIGITS else (r"\w" if c.isalnum() else None)
    before, after = edge(g[0]), edge(g[-1])
    pat = (f"(?<!{before})" if before else "") + re.escape(g) + (f"(?!{after})" if after else "")
    return re.search(pat, text) is not None


def content(s, asked):
    # The words that carry the answer: no stopwords, nothing the question already says.
    return [w for w in re.findall(r"\w+", normalize(s)) if w not in STOP and w not in asked]


def is_correct(pred, answers, question=""):
    """A gold alias inside the prediction ("277.08 PKR" holds "277.08"), or a close paraphrase: F1 >= 0.5 over the
    words the question does not already say, so "Siege of Troy" earns nothing for "Priam of Troy". A gold of one
    or two such words must be there whole ("Andy Murray" is not "Andy Burnham"), and a paraphrase must keep
    every number ("iPhone 16" is not "iPhone 18"; picking the old model is exactly the mistake this exam is for)."""
    p = normalize(pred)
    asked = set(re.findall(r"\w+", normalize(question)))
    asked |= {w[:-1] for w in asked if w.endswith("s")}           # "Arthur's sword" asks about "Arthur"
    pc = content(pred, asked)
    for g in answers:
        gn = normalize(g)
        if gn and inside(p, gn):
            return True
        gc = content(g, asked)
        same = sum((collections.Counter(pc) & collections.Counter(gc)).values())
        if not same or (len(gc) <= 2 and same < len(gc)):
            continue
        if 2 * same / (len(pc) + len(gc)) >= 0.5 and all(inside(p, w) for w in gc if any(c in DIGITS for c in w)):
            return True
    return False


def verdict(pred, answers, question=""):
    if not answers:
        return "trap_caught" if said_not_found(pred) else "trap_fooled"
    if said_not_found(pred):
        return "missed"
    return "answered_correct" if is_correct(pred, answers, question) else "wrong"


def accuracy(verdicts):
    return 100 * sum(v in ("answered_correct", "trap_caught") for v in verdicts) / max(len(verdicts), 1)


if __name__ == "__main__":
    for path in sys.argv[1:]:
        rows = [json.loads(x) for x in open(path, encoding="utf-8") if x.strip()]
        vs = [verdict(r["answer"], r["answers"], r["question"]) for r in rows]
        by_kind = collections.defaultdict(list)
        for r, v in zip(rows, vs):
            by_kind[r["kind"]].append(v)
        ok = sum(v in ("answered_correct", "trap_caught") for v in vs)
        kinds = "  ".join(f"{k} {sum(v in ('answered_correct', 'trap_caught') for v in x)}/{len(x)}"
                          for k, x in sorted(by_kind.items()))
        print(f"{path}: {ok}/{len(rows)} = {accuracy(vs):.1f}%  | {kinds}")