#!/usr/bin/env python3
"""Deterministic sample: 102 POPE-random images (seed 20260915) -> 3 disjoint subsets of 34.
Questions: all 6 POPE-random questions per image + all VQA-v2 val questions on the same image.
No model output is read here."""
import hashlib, json, random, sys, urllib.request, time
from pathlib import Path
PKG = Path(__file__).resolve().parents[2]; D = PKG / "work/data"
RAW = Path(sys.argv[1])  # dir holding coco_pope_random.json + VQA v2 val json (downloaded by 00_download_data.sh)
SEED, N, K = 20260915, 102, 3
pope = [json.loads(l) for l in open(RAW / "coco_pope_random.json")]
imgs = sorted({p["image"] for p in pope})
chosen = random.Random(SEED).sample(imgs, N)
subsets = {f"S{i+1}": chosen[i*N//K:(i+1)*N//K] for i in range(K)}
q = json.load(open(RAW / "v2_OpenEnded_mscoco_val2014_questions.json"))["questions"]
ann = {a["question_id"]: a for a in json.load(open(RAW / "v2_mscoco_val2014_annotations.json"))["annotations"]}
byimg = {}
for x in q: byimg.setdefault(x["image_id"], []).append(x)
items = []
for sub, names in subsets.items():
    for name in names:
        iid = int(name.split("_")[-1][:-4])
        for j, p in enumerate([p for p in pope if p["image"] == name]):
            items.append(dict(qid=f"pope-{p['question_id']}", task="pope", image=name, subset=sub, question=p["text"], gt=p["label"]))
        for x in sorted(byimg.get(iid, []), key=lambda x: x["question_id"]):
            a = ann[x["question_id"]]
            items.append(dict(qid=f"vqa-{x['question_id']}", task="vqav2", image=name, subset=sub, question=x["question"],
                              gt=a["multiple_choice_answer"], answers=[y["answer"] for y in a["answers"]], answer_type=a["answer_type"]))
out_img = D / "img_orig"; out_img.mkdir(parents=True, exist_ok=True)
shas = {}
for name in chosen:
    f = out_img / name
    if not f.exists():
        for t in range(5):
            try: urllib.request.urlretrieve(f"http://images.cocodataset.org/val2014/{name}", f); break
            except Exception: time.sleep(3)
    shas[name] = hashlib.sha256(f.read_bytes()).hexdigest()
json.dump(dict(seed=SEED, n_images=N, subsets=subsets, image_sha256=shas,
               raw_sha256={n: hashlib.sha256((RAW / n).read_bytes()).hexdigest() for n in ("coco_pope_random.json", "v2_OpenEnded_mscoco_val2014_questions.json", "v2_mscoco_val2014_annotations.json")}),
          open(D / "selection.json", "w"), indent=1)
json.dump(items, open(D / "questions.json", "w"), indent=0, ensure_ascii=False)
from collections import Counter
print("images", N, "questions", len(items), Counter(i["task"] for i in items), "pope yes/no", Counter(i["gt"] for i in items if i["task"] == "pope"),
      "per subset", {s: Counter(i["task"] for i in items if i["subset"] == s) for s in subsets})