#!/usr/bin/env python3
"""POST-HOC, DESCRIPTIVE (raised in review after the run; changes no pre-registered verdict).

LLaVA-1.5's instruction data draws its COCO images from the train2017 split, so test images that sit in
train2017 are very likely in the model's training image pool. This splits our 102 images by COCO-2017-split
membership and recomputes the paired loss vs the original image in each stratum.

IMPORTANT naming limit: the two strata are "in COCO train2017" and "not in COCO train2017" -- nothing more.
We never compared these images against the actual LLaVA-1.5 training image list, and LLaVA-1.5's mixture also
draws on Visual Genome / GQA / TextCaps, with Visual Genome spanning both COCO splits. This stratification is
therefore a public-split slice, and it cannot decide whether training overlap affected anything.
No new hypothesis test: intervals are shown to expose how wide they are.
"""
import json, re
from pathlib import Path
import numpy as np
W = Path(__file__).resolve().parents[1]
cfg = json.load(open(W / "code/run_config.json")); ARMS = cfg["arms"]; BASE = "orig"
qs = {q["qid"]: q for q in json.load(open(W / "data/questions.json"))}
val2017 = set(json.load(open(W / "data/coco_val2017_image_ids.json")))
A = {}
for l in open(W / "results/answers_main.jsonl"):
    r = json.loads(l); A.setdefault((r["arm"], r["qid"]), r["answer_raw"])
src = (W / "code/04_analyze.py").read_text()
ns = {"re": re, "np": np}; exec(src[src.index("contractions = {"):src.index("images = sorted(")], ns)
vqa_acc = ns["vqa_acc"]
def pope_label(t):
    w = t.split(".")[0].replace(",", "").split(" ")
    return "no" if ("No" in w or "not" in w or "no" in w) else "yes"
def iid(name): return int(name.split("_")[-1][:-4])
pope = [q for q in qs.values() if q["task"] == "pope"]; vqa = [q for q in qs.values() if q["task"] == "vqav2"]
images = sorted({q["image"] for q in qs.values()})
strata = {"not_in_coco_train2017": [n for n in images if iid(n) in val2017],
          "in_coco_train2017": [n for n in images if iid(n) not in val2017]}
rng = np.random.default_rng(20260917); B = 4000
def score(arm, q): return float(pope_label(A[(arm, q["qid"])]) == q["gt"]) if q["task"] == "pope" else vqa_acc(A[(arm, q["qid"])], q["answers"])
out = {"_read_this_first": {
  "stratum_naming_limit": ("两层是『属 COCO train2017』与『不属 COCO train2017』。后者不等于模型没见过："
    "LLaVA-1.5 训练混合还含 Visual Genome / GQA / TextCaps，VG 与 COCO 两个划分都重叠，且未与 LLaVA 实际图片清单比对。"
    "这只是按公开划分切出来的对照层，**不能用来判断训练重叠有没有影响**；全包一律用『不属 train2017 的 17 张』这类事实表述。"),
  "not_a_test_of_contamination": ("本文件既不能确认、也不能排除训练重叠的作用：无交互项检验、无功效分析、未按难度配对，"
    "两层规模 85 对 17，两层之差的区间（difference_between_strata）几乎都含 0，且正对照 bpg_q51 在 POPE 上两层几乎一样。"),
  "descriptive_only": "事后描述性分析，不改变 claims.json 里任何预注册判定。"},
 "_note": __doc__, "n_bootstrap": B, "seed": 20260917,
       "strata_sizes": {k: len(v) for k, v in strata.items()},
       "coco_val2017_ids_n": len(val2017), "arms": {}}
for sname, imgs in strata.items():
    idx = {n: i for i, n in enumerate(imgs)}; iset = set(imgs)
    boot = rng.integers(0, len(imgs), size=(B, len(imgs)))
    for task, qlist in (("pope", pope), ("vqav2", vqa)):
        sel = [q for q in qlist if q["image"] in iset]
        qimg = np.array([idx[q["image"]] for q in sel])
        cnt = np.bincount(qimg, None, len(imgs))
        base = np.array([score(BASE, q) for q in sel])
        out.setdefault("n_questions", {}).setdefault(sname, {})[task] = len(sel)
        out["arms"].setdefault(BASE, {}).setdefault(sname, {})[task + "_abs_acc"] = float(base.mean())
        for arm in ARMS:
            if arm == BASE: continue
            v = np.array([score(arm, q) for q in sel]); d = v - base
            s = np.bincount(qimg, d, len(imgs)); vb = s[boot].sum(1) / cnt[boot].sum(1)
            out["arms"].setdefault(arm, {}).setdefault(sname, {})[task] = dict(
                delta=float(d.mean()), ci95=[float(np.percentile(vb, 2.5)), float(np.percentile(vb, 97.5))],
                abs_acc=float(v.mean()), n_questions=len(sel))
# is the between-strata difference itself distinguishable?  (resample images within each stratum)
out["difference_between_strata"] = {}
for task, qlist in (("pope", pope), ("vqav2", vqa)):
    for arm in ARMS:
        if arm == BASE: continue
        parts = {}
        for sname, imgs in strata.items():
            idx = {n: i for i, n in enumerate(imgs)}; iset = set(imgs)
            sel = [q for q in qlist if q["image"] in iset]
            qimg = np.array([idx[q["image"]] for q in sel]); cnt = np.bincount(qimg, None, len(imgs))
            d = np.array([score(arm, q) - score(BASE, q) for q in sel])
            bt = rng.integers(0, len(imgs), size=(B, len(imgs))); s = np.bincount(qimg, d, len(imgs))
            parts[sname] = (d.mean(), s[bt].sum(1) / cnt[bt].sum(1))
        diff = parts["in_coco_train2017"][1] - parts["not_in_coco_train2017"][1]
        out["difference_between_strata"].setdefault(task, {})[arm] = dict(
            point=float(parts["in_coco_train2017"][0] - parts["not_in_coco_train2017"][0]),
            ci95=[float(np.percentile(diff, 2.5)), float(np.percentile(diff, 97.5))])
# ---- caveats that must travel with these numbers (machine-readable, not prose) ----
def g(arm, task, st): return out["arms"][arm][st][task]
ctrl_pope = (g("bpg_q51", "pope", "in_coco_train2017")["delta"], g("bpg_q51", "pope", "not_in_coco_train2017")["delta"])
ctrl_vqa = (g("bpg_q51", "vqav2", "in_coco_train2017")["delta"], g("bpg_q51", "vqav2", "not_in_coco_train2017")["delta"])
doubles = [a for a in ARMS if a not in (BASE, "bpg_q51")
           and any(abs(g(a, t, "not_in_coco_train2017")["delta"]) >= 1.8 * abs(g(a, t, "in_coco_train2017")["delta"]) for t in ("pope", "vqav2"))]
out["caveats"] = {
 "this_is_descriptive_not_a_test_of_contamination": (
   "分层只是把同一批答案按『图是否属 COCO train2017』切开重算，不是对『训练重叠导致损失变小』的检验。"
   "下面三条都说明它不能被当成因果结论。"),
 "control_arm_shows_no_split": {
   "bpg_q51_pope_in_vs_not_in_train2017": list(ctrl_pope),
   "bpg_q51_vqav2_in_vs_not_in_train2017": list(ctrl_vqa),
   "why_it_matters": ("极端对照 bpg_q51 在 POPE 上两层几乎一样（%.4f vs %.4f）——POPE 这一侧看不到分层差异。"
                      "但要说清另一半：在 VQA 上它两层差距很大（%.4f vs %.4f），方向与其他臂一致。"
                      "所以这条反面证据只成立于 POPE，不能当成『分层差异整体不存在』。"
                      % (ctrl_pope[0], ctrl_pope[1], ctrl_vqa[0], ctrl_vqa[1]))},
 "strata_differ_in_baseline_difficulty": {
   "orig_pope_abs_acc": {k: out["arms"][BASE][k]["pope_abs_acc"] for k in out["strata_sizes"]},
   "orig_vqav2_abs_acc": {k: out["arms"][BASE][k]["vqav2_abs_acc"] for k in out["strata_sizes"]},
   "why_it_matters": "两层的原图准确率本来就不同：不属 train2017 的 17 张 POPE 为 0.9118，属 train2017 的 85 张为 0.8510（VQA 则相反）。没有做难度或基线配对，两层不可直接视为同质对照——这本身就是反证之一。"},
 "intervals_overlap_and_difference_not_significant": (
   "不属 train2017 的那层只有 %d 张图（POPE %d 题、VQA %d 题），区间很宽；两层之差的区间几乎都包含 0（见 difference_between_strata）。"
   % (out["strata_sizes"]["not_in_coco_train2017"], out["n_questions"]["not_in_coco_train2017"]["pope"], out["n_questions"]["not_in_coco_train2017"]["vqav2"])),
 "arms_where_not_in_train2017_loss_is_at_least_1_8x": doubles,
 "arms_where_it_does_not_hold_pope_only": ["bpg_q51 (POPE)；同一臂的 VQA 仍然是不在 train2017 的那层掉得更多"] + [a for a in ARMS if a not in (BASE, "bpg_q51") and a not in doubles],
}
json.dump(out, open(W / "results/posthoc_train_overlap.json", "w"), indent=1, ensure_ascii=False)
print("caveats:", json.dumps({k: (v if isinstance(v, (list, str)) else "...") for k, v in out["caveats"].items()}, ensure_ascii=False)[:600])
print(json.dumps(out["strata_sizes"]), json.dumps(out["n_questions"]))
for arm in ("jpeg_t1", "webp_t1", "bpg_t1", "mbt2018mean_t1", "jpeg_t2", "mbt2018mean_t2"):
    for task in ("pope", "vqav2"):
        a = out["arms"][arm]["in_coco_train2017"][task]; b = out["arms"][arm]["not_in_coco_train2017"][task]
        dd = out["difference_between_strata"][task][arm]
        print(f"{arm:16s} {task:6s} in_train2017 {a['delta']*100:+6.2f} [{a['ci95'][0]*100:+6.2f},{a['ci95'][1]*100:+6.2f}] | not_in {b['delta']*100:+6.2f} [{b['ci95'][0]*100:+6.2f},{b['ci95'][1]*100:+6.2f}] | diff {dd['point']*100:+6.2f} [{dd['ci95'][0]*100:+6.2f},{dd['ci95'][1]*100:+6.2f}]")
o = out["arms"]["orig"]
print("orig abs acc  in_train2017:", {k: round(v, 4) for k, v in o["in_coco_train2017"].items()}, " not_in_train2017:", {k: round(v, 4) for k, v in o["not_in_coco_train2017"].items()})