#!/usr/bin/env python3
"""Independent recheck, pure standard library (no numpy), own parsing/scoring code.
Recomputes from raw answers: per-arm POPE accuracy, VQA-v2 accuracy, mean bpp, and POPE McNemar exact p vs orig."""
import json, math, re, string
from pathlib import Path
W = Path(__file__).resolve().parents[1]
arms = json.load(open(W / "code/run_config.json"))["arms"]
Q = json.load(open(W / "data/questions.json")); V = json.load(open(W / "data/variants.json"))["rows"]
A = {}
with open(W / "results/answers_main.jsonl") as fh:
    for line in fh:
        r = json.loads(line); A[r["arm"] + "|" + r["qid"]] = r["answer_raw"]
def says_no(text):  # same published rule as LLaVA eval_pope.py, re-implemented
    head = text.split(".", 1)[0].replace(",", "")
    return any(tok in ("No", "no", "not") for tok in head.split(" "))
NUM = dict(none="0", zero="0", one="1", two="2", three="3", four="4", five="5", six="6", seven="7", eight="8", nine="9", ten="10")
def norm(s):  # VQA official normaliser, re-implemented independently
    s = s.replace("\n", " ").replace("\t", " ").strip().lower()
    for p in ';/[]"{}()=+\\_-><@`,?!':
        if (p + " " in s or " " + p in s) or re.search(r"(\d)(,)(\d)", s): s = s.replace(p, "")
        else: s = s.replace(p, " ")
    s = re.sub(r"(?!<=\d)(\.)(?!\d)", "", s)
    words = [NUM.get(w, w) for w in s.split() if w not in ("a", "an", "the")]
    C = {"dont": "don't", "doesnt": "doesn't", "isnt": "isn't", "arent": "aren't", "cant": "can't", "wont": "won't", "didnt": "didn't", "im": "i'm", "thats": "that's", "whats": "what's", "theres": "there's", "shes": "she's", "hes": "he's"}
    return " ".join(C.get(w, w) for w in words)
def vqa_score(pred, gts):
    p = norm(pred); g = [norm(x) for x in gts]; tot = 0.0
    for i in range(len(g)):
        tot += min(1.0, sum(1 for j in range(len(g)) if j != i and g[j] == p) / 3.0)
    return tot / len(g)
out = {"arms": {}, "mcnemar_vs_orig": {}}
correct = {}
for arm in arms:
    pope = [q for q in Q if q["task"] == "pope"]; vqa = [q for q in Q if q["task"] == "vqav2"]
    c = [(("no" if says_no(A[arm + "|" + q["qid"]]) else "yes") == q["gt"]) for q in pope]; correct[arm] = c
    bpps = [r["bpp"] for r in V if r["arm"] == arm]
    out["arms"][arm] = dict(pope_acc=sum(c) / len(c), vqa_acc=sum(vqa_score(A[arm + "|" + q["qid"]], q["answers"]) for q in vqa) / len(vqa),
                            mean_bpp=sum(bpps) / len(bpps))
for arm in arms:
    if arm == "orig": continue
    b = sum(1 for x, y in zip(correct[arm], correct["orig"]) if x and not y); c = sum(1 for x, y in zip(correct[arm], correct["orig"]) if y and not x)
    n = b + c
    p = 1.0 if n == 0 else min(1.0, 2.0 * sum(math.comb(n, i) for i in range(min(b, c) + 1)) / 2 ** n)
    out["mcnemar_vs_orig"][arm] = dict(arm_only=b, orig_only=c, p=p, delta_acc=out["arms"][arm]["pope_acc"] - out["arms"]["orig"]["pope_acc"])
# descriptive label split, recomputed independently (no test, no p-values)
pope_qs = [q for q in Q if q["task"] == "pope"]
out["descriptive_by_gt_label"] = {}
for arm in arms:
    ys = [(("no" if says_no(A[arm + "|" + q["qid"]]) else "yes") == q["gt"]) for q in pope_qs if q["gt"] == "yes"]
    ns = [(("no" if says_no(A[arm + "|" + q["qid"]]) else "yes") == q["gt"]) for q in pope_qs if q["gt"] == "no"]
    preds = [("no" if says_no(A[arm + "|" + q["qid"]]) else "yes") for q in pope_qs]
    gts = [q["gt"] for q in pope_qs]
    tp = sum(p == "yes" and g == "yes" for p, g in zip(preds, gts)); fp = sum(p == "yes" and g == "no" for p, g in zip(preds, gts))
    fn = sum(p == "no" and g == "yes" for p, g in zip(preds, gts))
    out["descriptive_by_gt_label"][arm] = dict(acc_on_gt_yes=sum(ys) / len(ys), acc_on_gt_no=sum(ns) / len(ns),
        precision=tp / max(1, tp + fp), recall=tp / max(1, tp + fn),
        yes_ratio=preds.count("yes") / len(preds))
json.dump(out, open(W / "results/recheck.json", "w"), indent=1)
print("recheck ok", {a: round(v["pope_acc"], 4) for a, v in out["arms"].items()})