#!/usr/bin/env python3
"""Pre-registered analysis (numpy path). Reads answers_main.jsonl + variants.json + questions.json -> results/analysis.json."""
import json, math, re
from pathlib import Path
import numpy as np
W = Path(__file__).resolve().parents[1]; cfg = json.load(open(W / "code/run_config.json"))
qs = {q["qid"]: q for q in json.load(open(W / "data/questions.json"))}
var = json.load(open(W / "data/variants.json")); ARMS = cfg["arms"]; BASE = "orig"
ans = {}; n_dup_identical = 0
for l in open(W / "results/answers_main.jsonl"):
    r = json.loads(l); k = (r["arm"], r["qid"])
    if k in ans:
        if ans[k] != r["answer_raw"]: raise SystemExit(f"conflicting duplicate answer for {k}: {ans[k]!r} vs {r['answer_raw']!r}")
        n_dup_identical += 1; continue
    ans[k] = r["answer_raw"]
# --- POPE parsing: LLaVA official eval_pope.py rule
def pope_label(t):
    t = t.split(".")[0].replace(",", ""); w = t.split(" ")
    return "no" if ("No" in w or "not" in w or "no" in w) else "yes"
# --- VQA-v2 official normalisation (vqaEval.py)
contractions = {"aint": "ain't", "arent": "aren't", "cant": "can't", "couldve": "could've", "couldnt": "couldn't", "didnt": "didn't", "doesnt": "doesn't", "dont": "don't", "hadnt": "hadn't", "hasnt": "hasn't", "havent": "haven't", "hed": "he'd", "hes": "he's", "howd": "how'd", "howll": "how'll", "hows": "how's", "Im": "I'm", "Ive": "I've", "isnt": "isn't", "itd": "it'd", "itll": "it'll", "let's": "let's", "maam": "ma'am", "mightnt": "mightn't", "mightve": "might've", "mustnt": "mustn't", "mustve": "must've", "neednt": "needn't", "notve": "not've", "oclock": "o'clock", "oughtnt": "oughtn't", "shant": "shan't", "shed": "she'd", "shes": "she's", "shouldve": "should've", "shouldnt": "shouldn't", "thats": "that's", "thered": "there'd", "theres": "there's", "theyd": "they'd", "theyll": "they'll", "theyre": "they're", "theyve": "they've", "wasnt": "wasn't", "wed": "we'd", "weve": "we've", "werent": "weren't", "whatll": "what'll", "whatre": "what're", "whats": "what's", "whatve": "what've", "whens": "when's", "whered": "where'd", "wheres": "where's", "whereve": "where've", "whod": "who'd", "wholl": "who'll", "whos": "who's", "whove": "who've", "whyll": "why'll", "whyre": "why're", "whys": "why's", "wont": "won't", "wouldve": "would've", "wouldnt": "wouldn't", "yall": "y'all", "youd": "you'd", "youll": "you'll", "youre": "you're", "youve": "you've"}
nums = {"none": "0", "zero": "0", "one": "1", "two": "2", "three": "3", "four": "4", "five": "5", "six": "6", "seven": "7", "eight": "8", "nine": "9", "ten": "10"}
articles = {"a", "an", "the"}; P1 = re.compile(r"(?!<=\d)(\.)(?!\d)"); P2 = re.compile(r"(\d)(,)(\d)")
punct = [';', r"/", '[', ']', '"', '{', '}', '(', ')', '=', '+', '\\', '_', '-', '>', '<', '@', '`', ',', '?', '!']
def proc_punct(t):
    for p in punct:
        t = t.replace(p, "") if ((p + " " in t or " " + p in t) or (re.search(P2, t) is not None)) else t.replace(p, " ")
    return P1.sub("", t)
def proc_digit(t):
    out = [nums.get(w, w) for w in t.split() if w not in articles]
    return " ".join(contractions.get(w, w) for w in out)
def vqa_norm(t): return proc_digit(proc_punct(t.replace("\n", " ").replace("\t", " ").strip().lower()))
def vqa_acc(pred, gts):
    p = vqa_norm(pred); g = [vqa_norm(x) for x in gts]
    return float(np.mean([min(1.0, sum(1 for j, y in enumerate(g) if j != i and y == p) / 3) for i in range(len(g))]))
images = sorted({q["image"] for q in qs.values()}); img_idx = {n: i for i, n in enumerate(images)}
pope_q = [q for q in qs.values() if q["task"] == "pope"]; vqa_q = [q for q in qs.values() if q["task"] == "vqav2"]
def pope_vec(arm):  # 1/0 correctness per POPE question (fixed order)
    return np.array([pope_label(ans[(arm, q["qid"])]) == q["gt"] for q in pope_q], float)
def vqa_vec(arm): return np.array([vqa_acc(ans[(arm, q["qid"])], q["answers"]) for q in vqa_q])
pope_img = np.array([img_idx[q["image"]] for q in pope_q]); vqa_img = np.array([img_idx[q["image"]] for q in vqa_q])
rng = np.random.default_rng(20260915); B = 10000
boot_idx = rng.integers(0, len(images), size=(B, len(images)))
def cluster_boot_diff(vec_a, vec_b, qimg):
    """95% CI of mean(a)-mean(b), resampling images (questions clustered by image)."""
    d = vec_a - vec_b; s = np.bincount(qimg, d, len(images)); c = np.bincount(qimg, None, len(images))
    num = s[boot_idx].sum(1); den = c[boot_idx].sum(1); v = num / den
    return [float(np.percentile(v, 2.5)), float(np.percentile(v, 97.5))]
def mcnemar(a, b):
    n01 = int(((a == 1) & (b == 0)).sum()); n10 = int(((a == 0) & (b == 1)).sum()); n = n01 + n10
    if n == 0: return dict(a_only=n01, b_only=n10, p_exact_two_sided=1.0)
    k = min(n01, n10); p = min(1.0, 2 * sum(math.comb(n, i) for i in range(k + 1)) / 2 ** n)
    return dict(a_only=n01, b_only=n10, p_exact_two_sided=p)
def pope_stats(arm):
    pred = [pope_label(ans[(arm, q["qid"])]) for q in pope_q]; gt = [q["gt"] for q in pope_q]
    tp = sum(p == "yes" and g == "yes" for p, g in zip(pred, gt)); fp = sum(p == "yes" and g == "no" for p, g in zip(pred, gt))
    fn = sum(p == "no" and g == "yes" for p, g in zip(pred, gt)); prec = tp / max(1, tp + fp); rec = tp / max(1, tp + fn)
    raw_ok = sum(bool(re.match(r"^\s*(yes|no)\b", ans[(arm, q["qid"])], re.I)) for q in pope_q)
    return dict(acc=float(pope_vec(arm).mean()), precision=prec, recall=rec, f1=2 * prec * rec / max(1e-12, prec + rec),
                yes_ratio=pred.count("yes") / len(pred), raw_yes_no_format_rate=raw_ok / len(pope_q),
                acc_by_subset={s: float(np.mean([pope_label(ans[(arm, q["qid"])]) == q["gt"] for q in pope_q if q["subset"] == s])) for s in ("S1", "S2", "S3")})
rate = {}
for arm in ARMS:
    b = [r["bpp"] for r in var["rows"] if r["arm"] == arm]
    rate[arm] = dict(mean_bpp=float(np.mean(b)), median_bpp=float(np.median(b)), n=len(b))
SCOPE_AND_LIMITS = {
 "read_this_first": "本文件的每个数字都只在下面这些范围内成立。",
 "p5_saturation_is_pope_only": ("P5 的『饱和』只在 POPE 上判过。同一码率（0.5 bpp）下 VQA-v2 的点估计仍在下降："
   "jpeg +0.02、bpg -0.68、mbt2018-mean -1.29、webp -1.63 个百分点，其中 webp 的区间不含 0。"
   "所以『0.5 bpp 饱和』必须限定为 POPE（且其含义是非劣效：不差于原图超过 3 点），不能说成下游任务整体无损。"),
 "rate_asymmetry_at_t2": "mbt2018mean_t2 中位码率 0.522 bpp，比其他编码器高约 4%（均值之比 1.015）；T1 匹配良好（1.011）。0.5 bpp 档的比较要带上这个不对称。",
 "sample_is_mostly_training_images": ("102 张 COCO val2014 图里 85 张属 COCO train2017，即 LLaVA-1.5 指令数据的 COCO 图源；"
   "POPE 的 yes/no 标签也来自同一批 COCO 物体标注。因此这些数字主要反映『图与标注都可能在训练数据里』的样本，"
   "不能当作对新图的估计。分层结果见 posthoc_train_overlap.json。"),
 "ci_containing_zero_is_not_no_loss": ("区间含 0 只表示未能检出差异。注意 P5 本身就是非劣效判据（CI 下界 ≥ −3 个百分点）："
   "在 POPE 上四个编码器均不差于原图超过 3 个百分点，实测下界 ≥ −1.31 点；但这不等于无损失。"
   "本轮的实际分辨力用区间半宽表示（不是功效分析）：POPE 上『某臂 vs 原图』的 95% 区间半宽 0.98-2.45 个百分点"
   "（中位 1.47），『编码器两两比较』1.31-2.53（中位 1.76）；VQA 上『某臂 vs 原图』中位 1.94。"
   "小于这个量级的差异本轮看不出来；大于它的（如 P4 的 3.10 点）才可能落在区间之外。"),
 "no_ranking_between_codecs": ("dominance_screen_* 只是按点估计做的支配关系筛选，不是统计意义上的占优，更不是 Pareto 前沿："
   "除 mbt2018-mean vs JPEG @0.25bpp（Δ=3.10 点，区间不含 0）外，其余编码器两两差距都在各自区间半宽之内。"
   "只有两个码率点，也画不出前沿曲线。"),
 "bpp_is_at_original_resolution": "bpp 按 COCO 原分辨率计算，而模型只看 336px 缩图，送进模型的有效码率远高于标称值；orig 本身是约 4.785 bpp 的 JPEG，不是无损基线。",
 "seeds_are_subsets_not_repeats": "S1/S2/S3 是同一次抽样切出的互不重叠图子集；解码确定性，没有运行间随机性。",
 "e1_name_is_wrong": ("预注册 E1 名为『幻觉方向偏移』不准确（预注册已锁，此处纠正）：召回率下降可检出、精确率几乎不变；"
   "但 POPE 假阳性本来就极少（原图 3 道 → bpg_q51 7 道），幻觉方向有没有变化本轮无法判断，两个方向都不能断言。"),
 "conclusion_subject_is_one_model_pair": "可区分的只有 mbt2018-mean 对 JPEG @0.25bpp；神经侧只有一个 MSE 训练的预训练模型，不能写成『神经压缩优于传统编码』。",
}
res = dict(_scope_and_limits=SCOPE_AND_LIMITS, n_images=len(images), n_pope=len(pope_q), n_vqa=len(vqa_q), rate=rate, arms={}, vs_orig={}, codec_pairs={})
V = {arm: (pope_vec(arm), vqa_vec(arm)) for arm in ARMS}
for arm in ARMS:
    res["arms"][arm] = dict(pope=pope_stats(arm), vqa_acc=float(V[arm][1].mean()),
                            vqa_acc_by_subset={s: float(np.mean([V[arm][1][i] for i, q in enumerate(vqa_q) if q["subset"] == s])) for s in ("S1", "S2", "S3")})
null_all_yes = float(np.mean([q["gt"] == "yes" for q in pope_q]))
res["null_pope_all_yes_acc"] = null_all_yes
def compare(a, b):
    pa, va = V[a]; pb, vb = V[b]
    sub = {s: float(np.mean([pa[i] - pb[i] for i, q in enumerate(pope_q) if q["subset"] == s])) for s in ("S1", "S2", "S3")}
    return dict(pope_delta_acc=float(pa.mean() - pb.mean()), pope_delta_ci95_cluster_boot=cluster_boot_diff(pa, pb, pope_img),
                pope_mcnemar=mcnemar(pa, pb), pope_delta_by_subset=sub, n_subsets_negative=sum(v < 0 for v in sub.values()),
                vqa_delta_acc=float(va.mean() - vb.mean()), vqa_delta_ci95_cluster_boot=cluster_boot_diff(va, vb, vqa_img),
                bpp_ratio_a_over_b=rate[a]["mean_bpp"] / rate[b]["mean_bpp"])
for arm in ARMS:
    if arm != BASE: res["vs_orig"][arm] = compare(arm, BASE)
for t in ("t1", "t2"):
    for c in ("webp", "bpg", "mbt2018mean"):
        res["codec_pairs"][f"{c}_{t}_vs_jpeg_{t}"] = compare(f"{c}_{t}", f"jpeg_{t}")
# Pareto frontier over (mean bpp lower better, POPE acc / VQA acc higher better), codec arms only (+orig)
def frontier(metric):
    pts = [(rate[a]["mean_bpp"], metric(a), a) for a in ARMS]
    return sorted(a for b, m, a in pts if not any((b2 <= b and m2 >= m) and (b2 < b or m2 > m) for b2, m2, _ in pts))
res["dominance_screen_pope_point_estimates_only"] = frontier(lambda a: res["arms"][a]["pope"]["acc"])
res["dominance_screen_vqa_point_estimates_only"] = frontier(lambda a: res["arms"][a]["vqa_acc"])
CODEC_ARMS = [a for a in ARMS if a.endswith(("_t1", "_t2"))]
def frontier_codec(metric):
    pts = [(rate[a]["mean_bpp"], metric(a), a) for a in CODEC_ARMS]
    return sorted(a for b, m, a in pts if not any((b2 <= b and m2 >= m) and (b2 < b or m2 > m) for b2, m2, _ in pts))
res["dominance_screen_pope_codec_arms_only_point_estimates"] = frontier_codec(lambda a: res["arms"][a]["pope"]["acc"])
res["dominance_screen_vqa_codec_arms_only_point_estimates"] = frontier_codec(lambda a: res["arms"][a]["vqa_acc"])
for k, v in res["codec_pairs"].items():
    v["rate_confound_flag"] = not (0.85 <= v["bpp_ratio_a_over_b"] <= 1.18)
# P6: is VQA more rate-sensitive than POPE for JPEG at t1?  diff-of-diffs with image bootstrap
def did(arm):
    pa, va = V[arm]; pb, vb = V[BASE]
    dp = np.bincount(pope_img, pa - pb, len(images)); cp = np.bincount(pope_img, None, len(images))
    dv = np.bincount(vqa_img, va - vb, len(images)); cv = np.bincount(vqa_img, None, len(images))
    v = dv[boot_idx].sum(1) / cv[boot_idx].sum(1) - dp[boot_idx].sum(1) / cp[boot_idx].sum(1)
    return dict(vqa_delta_minus_pope_delta=float((va - vb).mean() - (pa - pb).mean()), ci95=[float(np.percentile(v, 2.5)), float(np.percentile(v, 97.5))])
res["vqa_vs_pope_sensitivity"] = {arm: did(arm) for arm in ARMS if arm != BASE}
res["compute"] = {}
rows = [json.loads(l) for l in open(W / "results/answers_main.jsonl")]
res["compute"] = dict(total_request_wall_h=sum(r["wall_s"] for r in rows) / 3600, n_requests=len(rows))
def _mean(x): return float(np.mean(x)) if len(x) else None
res["compute"]["mean_first_question_s"] = _mean([r["wall_s"] for r in rows if (r["prompt_n"] or 0) > 200])
res["compute"]["mean_cached_question_s"] = _mean([r["wall_s"] for r in rows if (r["prompt_n"] or 0) <= 200])
res["n_duplicate_rows_identical_ignored"] = n_dup_identical
def verdict_ci(ci, lo_thr, hi_thr):
    return "supported" if ci[0] > lo_thr else ("refuted" if ci[1] < hi_thr else "inconclusive")
o = res["arms"]["orig"]["pope"]; ver = {}
ver["P1_sanity"] = dict(acc=o["acc"], format_rate=o["raw_yes_no_format_rate"],
                        verdict="supported" if (0.80 <= o["acc"] <= 0.93 and o["raw_yes_no_format_rate"] >= 0.98) else "refuted")
c = res["vs_orig"]["bpg_q51"]
ver["P2_positive_control"] = dict(delta=c["pope_delta_acc"], p=c["pope_mcnemar"]["p_exact_two_sided"], ci95=c["pope_delta_ci95_cluster_boot"],
                                  verdict="supported" if (c["pope_delta_acc"] <= -0.05 and c["pope_mcnemar"]["p_exact_two_sided"] < 0.01) else "refuted")
ver["P3_null"] = dict(margin=o["acc"] - null_all_yes, verdict="supported" if o["acc"] - null_all_yes >= 0.25 else "refuted")
gate_ok = ver["P1_sanity"]["verdict"] == "supported" and ver["P2_positive_control"]["verdict"] == "supported"
def gated(v): return v if gate_ok else "not_judged_gate_failed"
def _qualify(v, p):
    """A 'supported' that does not clear the pre-registered Bonferroni line must say so in its own name."""
    return v + "_by_own_ci_rule_but_fails_bonferroni_0.01" if (v == "supported" and p >= 0.01) else v
for c_ in ("mbt2018mean", "bpg"):
    cp = res["codec_pairs"][f"{c_}_t1_vs_jpeg_t1"]
    ver[f"P4_{c_}_t1_gt_jpeg_t1"] = dict(delta=cp["pope_delta_acc"], ci95=cp["pope_delta_ci95_cluster_boot"], p=cp["pope_mcnemar"]["p_exact_two_sided"],
                                        by_subset=cp["pope_delta_by_subset"], rate_confound_flag=cp["rate_confound_flag"],
                                        role="primary" if c_ == "mbt2018mean" else "secondary",
                                        verdict=gated(_qualify(verdict_ci(cp["pope_delta_ci95_cluster_boot"], 0.0, 0.0),
                                                                cp["pope_mcnemar"]["p_exact_two_sided"])))
for c_ in ("jpeg", "webp", "bpg", "mbt2018mean"):
    cv = res["vs_orig"][f"{c_}_t2"]; ci = cv["pope_delta_ci95_cluster_boot"]
    ver[f"P5_{c_}_t2_within_3pts"] = dict(delta=cv["pope_delta_acc"], ci95=ci, p=cv["pope_mcnemar"]["p_exact_two_sided"], by_subset=cv["pope_delta_by_subset"],
                                         verdict=gated("supported" if ci[0] >= -0.03 else ("refuted" if ci[1] < -0.03 else "inconclusive")))
s6 = res["vqa_vs_pope_sensitivity"]["jpeg_t1"]
ver["P6_vqa_more_sensitive_jpeg_t1"] = dict(**s6, verdict=gated("supported" if s6["ci95"][1] < 0 else ("refuted" if s6["ci95"][0] > 0 else "inconclusive")))
ver["E1_yes_ratio_by_arm"] = {a: res["arms"][a]["pope"]["yes_ratio"] for a in ARMS}
res["verdicts"] = ver; res["bonferroni_threshold_primary"] = 0.01
# --- DESCRIPTIVE ONLY (not pre-registered, not a hypothesis test): split each arm's POPE accuracy by
# ground-truth label, and report how many questions actually changed on each side. On a 50/50 truth a shift
# toward answering "no" lowers accuracy by itself, so the split shows WHERE a drop sits -- it does NOT
# identify the mechanism. Any p values in this block are resolution aids only (is a side's change
# distinguishable from zero at n=306?), never evidence for a claim.
yes_mask = np.array([q["gt"] == "yes" for q in pope_q]); no_mask = ~yes_mask
def label_split(arm):
    v = V[arm][0]
    return float(v[yes_mask].mean()), float(v[no_mask].mean())
oy, on = label_split(BASE); desc = {}
for arm in ARMS:
    ay, an = label_split(arm); dy, dn = ay - oy, an - on
    yr = res["arms"][arm]["pope"]["yes_ratio"]
    my = mcnemar(V[arm][0][yes_mask], V[BASE][0][yes_mask]); mn = mcnemar(V[arm][0][no_mask], V[BASE][0][no_mask])
    # p values here are RESOLUTION AIDS ONLY: they say whether a side's change is distinguishable from zero
    # at this sample size. They are not hypothesis tests and carry no claim.
    sig_y = my["p_exact_two_sided"] < 0.01; sig_n = mn["p_exact_two_sided"] < 0.01
    if arm == BASE: src = "baseline (reference for the other arms)"
    elif sig_y and sig_n: src = "distinguishable_change_on_both_sides"
    elif sig_y: src = "change_concentrated_on_gt_yes_side; gt_no side not distinguishable from zero here"
    elif sig_n: src = "change_concentrated_on_gt_no_side; gt_yes side not distinguishable from zero here"
    else: src = "no_distinguishable_change_on_either_side"
    desc[arm] = dict(acc_on_gt_yes=ay, acc_on_gt_no=an, delta_vs_orig_on_gt_yes=dy, delta_vs_orig_on_gt_no=dn,
                     yes_ratio=yr, yes_ratio_delta_vs_orig=yr - res["arms"][BASE]["pope"]["yes_ratio"],
                     n_gt_yes=int(yes_mask.sum()), n_gt_no=int(no_mask.sum()),
                     gt_yes_n_orig_correct_arm_wrong=my["b_only"], gt_yes_n_orig_wrong_arm_correct=my["a_only"],
                     gt_no_n_orig_correct_arm_wrong=mn["b_only"], gt_no_n_orig_wrong_arm_correct=mn["a_only"],
                     p_exact_two_sided_gt_yes_for_resolution_only=my["p_exact_two_sided"],
                     p_exact_two_sided_gt_no_for_resolution_only=mn["p_exact_two_sided"],
                     n_wrong_gt_no_orig=int((V[BASE][0][no_mask] == 0).sum()), n_wrong_gt_no_arm=int((V[arm][0][no_mask] == 0).sum()),
                     where_the_change_sits=src,
                     mechanism_note="Descriptive only. This split cannot separate loss of visual evidence from an overall shift toward answering 'no': both predict a drop concentrated on gt=yes. The gt=no side sits at a ceiling (orig 0.9902) where this sample cannot resolve changes of a few questions, so it supports and excludes nothing.")
res["descriptive_by_gt_label_not_a_test"] = desc
json.dump(res, open(W / "results/analysis.json", "w"), indent=1)
print(json.dumps({a: (round(res["arms"][a]["pope"]["acc"], 4), round(res["arms"][a]["vqa_acc"], 4), round(rate[a]["mean_bpp"], 4)) for a in ARMS}, indent=0))