#!/usr/bin/env python3
"""genbench4 · 换场景首帧的语义一致性门槛（PCE env pce_f3，Qwen3-VL-32B，genbench3/vlm_gate.QwenVL：Yes / No 首 token 的 P(yes)）。

整图题（图 1 = 真实帧 0，图 2 = 候选）：
  * layout：同一关卡换成新场景后布局不变，且**没有东西换了类别**（空地不许变墙 / 悬崖、石头不许变房子……；同类里换成该场景
    的东西可以 —— 冒烟 v1 的措辞「不管东西变成了哪一种」把石头 → 城墙也放过了）；
  * chars ：两个角色都在、同位置、同姿势 / 体型（只记录；门槛用下面的逐角色放大比对 charcrop）；only2：画面里只有这两个角色；
    scene：（只看图 2）是不是该场景。
逐角色放大比对 charcrop：同一个框（角色掩码包围框外扩 50 %）裁原帧 / 候选：「同一个角色还在、同位置 / 姿势 / 体型吗」；基线 =
  原帧对原帧，基线 < 0.5 的角色不判（VLM 在原帧里就认不出：被叶子挡住 / 太小 / 叠在一起）。
逐区域题（按帧 0 的语义图切，**原帧问同样的题作基线**，缓存在 sources/<tag>/regions.json）：
  * open：terrain / paving 占 ≥ 92 % 且离角色 ≥ 12 px 的 96×96 格（最多 10 个，均匀取）：「只是可走的地面 / 地板，里面没有
    立着的墙、建筑、悬崖、石阵、树、大物体」；
  * obj ：architecture / rock / vegetation / prop 的连通块（≥ 300 px，最多 8 个，按面积），裁出包围框（外扩、放大到短边 ≥ 336，
    品红描出该块的语义轮廓）：「描出来的东西是不是 {该类的通称}」。
  区域分 = 原帧答 Yes（P ≥ 0.5）的区域里候选也答 Yes 的比例（obj 按面积加权）⇒ open_ok / obj_ok。
自动：语义边界边缘召回 ≥ 0.4 × 原帧（只防糊成一片；场景换光照 / 材质后 0.4–0.7 是常态，genbench3 的 0.8 不适用）。
genbench3 的「角色仍填满语义轮廓」（bg_fill）不用：换成暗场景（熔岩 / 地牢）后深色 Boss 与背景对比变低，bg_fill 涨到 0.4–0.7
而角色其实都在（冒烟目测）⇒ 角色在不在只靠 VLM chars。门槛值见 GATE（冒烟 outputs/genbench4_smoke 目测标定：区域题抓到的
都是真违规 —— 空地角落的树篱 / 枯灌木 / 城墙块、铺地上多出的门 / 火把 / 旗帜、石头变树篱 / 木架）。
输出 <cand>.vlm.json；--merge ⇒ outputs/genbench4/gate.json（每个 (起点, 场景) 槽在通过的候选里取分最高的）[+ --jobs ⇒ jobs.json]。
usage: CUDA_VISIBLE_DEVICES=0 $FPY scripts/genbench4/gate.py --shard 0 --nshards 8 ；$FPY scripts/genbench4/gate.py --merge --jobs"""
import argparse
import glob
import json
import os
import re
import sys
import time

import numpy as np
from PIL import Image
from scipy import ndimage

M = "/opt/dlami/nvme/zhiyangdeng/maliang"
G = os.environ.get("GB4", f"{M}/outputs/genbench4")
sys.path[:0] = [f"{M}/code/gc25d", f"{M}/scripts/genbench3", f"{M}/scripts/genbench4"]
from scenes import SCENES  # noqa: E402

PLAYER, BOSS = (10, 12, 13, 14, 15, 16), (11, 17, 18, 19, 20, 21)
HEAD = ("Image 1 is a screenshot from a top-down isometric action RPG game. It shows a smaller player hero and a large "
        "boss monster. ")
Q = {
    "layout": HEAD + "Image 2 is supposed to be the same level with exactly the same layout, re-themed as {title}. Ignore "
              "the art style, materials, colours and lighting. Each thing may turn into a different thing of the same "
              "kind (a building into another building, a plant into another plant, a rock into another rock, a floor into "
              "another floor), but nothing may change its kind: open ground or floor must not become a wall, building, "
              "cliff, rock pile or tree, and a rock or plant must not become a building. Does image 2 keep every building "
              "or wall, path or floor, open ground area, rock, plant and prop of image 1 in the same place with the same "
              "footprint, size and kind, with nothing large added or removed? Answer Yes or No.",
    "chars": HEAD + "Ignore the art style and the characters' appearance. Are both characters of image 1 (the player hero "
             "and the boss monster) clearly present in image 2, at the same positions, with the same poses and sizes? "
             "Answer Yes or No.",
    "only2": HEAD + "Does image 2 contain exactly these two characters (one smaller player hero and one large boss "
             "monster) and no other people, creatures or characters? Answer Yes or No.",
    "scene": "This is a screenshot from a top-down isometric action RPG game. Does this level clearly show {vlm}? "
             "Answer Yes or No.",
}
CROP = "This is a close-up crop from a top-down isometric game level. "
Q_OPEN = (CROP + "Does this crop show only open walkable ground or floor (for example dirt, sand, snow, grass, mud, ash, "
          "shallow water, ice, tiles, flagstones or planks, possibly with cracks, puddles, pebbles, small debris or small "
          "tufts), with no walls, buildings, towers, cliffs, rock formations, trees, large plants, fences or large objects "
          "standing in it? Answer Yes or No.")
Q_OBJ = {
    3: "part of a building, wall, tower, gate, hut, pillar, stairs or other built structure",
    4: "a rock, boulder, stone block, crystal formation, ice block, or a pile of rubble or stones (a natural lump, not a "
       "building or wall)",
    5: "a plant: a tree, bush, hedge, grass, reeds, cactus, mushroom, fungus, moss or other vegetation (living, dead or "
       "frozen)",
    6: "a small object or prop, such as a crate, barrel, cart, wagon, lamp, brazier, statue, sign, fence, tombstone, boat, "
       "container or piece of debris (not a building, rock or plant)",
}
Q_OBJ_T = CROP + "Is the thing outlined in magenta {what}? Answer Yes or No."
# 逐角色放大比对（整图 chars 题对小 / 重叠 / 被遮的角色常常连原帧都答 No ⇒ 全源失败，冒烟目测：train_v9_05 原帧里两人藏在竹叶后、
# train_v9_10 两人又小又叠在一起）：同一个框裁原帧 / 候选，基线 = 原帧对原帧；基线 < 0.5 的角色（VLM 在原帧里就认不出）不判
WHO = {"player": "the player hero (a smaller human character)", "boss": "the boss (a large enemy monster)"}
Q_CHAR = ("Image 1 is a close-up of {who} in a top-down isometric action RPG game. Image 2 is a close-up of exactly the same "
          "spot after the whole level was re-themed. Ignore colours, lighting, materials and art style. Is the same "
          "character still clearly there in image 2, at the same position, with the same pose, size and body shape? "
          "Answer Yes or No.")
CELL, OPEN_FRAC, CHAR_GAP, MAX_OPEN = 96, 0.92, 12, 10
MIN_COMP, MAX_COMP = 300, 8
GATE = {"layout": 0.3, "chars": 0.5, "only2": 0.5, "scene": 0.5, "open_ok": 0.8, "obj_ok": 0.7, "edge_rel": 0.4}


# ------------------------------------------------------------------ 区域 ----
def regions(ids):
    """[{kind: open|obj, cls, box (y0,y1,x0,x1), area, mask (obj 才有，box 内的 bool)}]。"""
    H, W = ids.shape
    chars = ndimage.binary_dilation(np.isin(ids, PLAYER + BOSS), iterations=CHAR_GAP)
    openm = np.isin(ids, (1, 2)) & ~chars
    cells = []
    oy, ox = (H % CELL) // 2, (W % CELL) // 2
    for y in range(oy, H - CELL + 1, CELL):
        for x in range(ox, W - CELL + 1, CELL):
            if openm[y:y + CELL, x:x + CELL].mean() >= OPEN_FRAC:
                cells.append({"kind": "open", "cls": 1, "box": (y, y + CELL, x, x + CELL), "area": CELL * CELL})
    if len(cells) > MAX_OPEN:
        cells = [cells[i] for i in np.linspace(0, len(cells) - 1, MAX_OPEN).round().astype(int)]
    comps = []
    for c in (3, 4, 5, 6):
        lab, n = ndimage.label(ids == c, structure=np.ones((3, 3)))
        for k, sl in enumerate(ndimage.find_objects(lab), 1):
            m = lab[sl] == k
            if m.sum() < MIN_COMP:
                continue
            y0, y1, x0, x1 = sl[0].start, sl[0].stop, sl[1].start, sl[1].stop
            py, px = max(12, (y1 - y0) // 4), max(12, (x1 - x0) // 4)
            Y0, Y1, X0, X1 = max(0, y0 - py), min(H, y1 + py), max(0, x0 - px), min(W, x1 + px)
            mm = np.zeros((Y1 - Y0, X1 - X0), bool)
            mm[y0 - Y0:y1 - Y0, x0 - X0:x1 - X0] = m
            comps.append({"kind": "obj", "cls": c, "box": (Y0, Y1, X0, X1), "area": int(m.sum()), "mask": mm})
    comps = sorted(comps, key=lambda r: -r["area"])[:MAX_COMP]
    return cells + comps


def crop_img(img, r, min_side=336):
    y0, y1, x0, x1 = r["box"]
    c = img.crop((x0, y0, x1, y1))
    if r["kind"] == "obj":                                  # 品红描出该块的语义轮廓（原帧 / 候选同一条线）
        edge = r["mask"] & ~ndimage.binary_erosion(r["mask"], iterations=1)
        edge = ndimage.binary_dilation(edge, iterations=1)
        a = np.asarray(c).copy()
        a[edge] = (255, 0, 255)
        c = Image.fromarray(a)
    s = max(1.0, min_side / max(1, min(c.size)))
    return c.resize((int(c.width * s), int(c.height * s)), Image.LANCZOS)


def char_boxes(ids, pad=0.5, min_side=64):
    H, W = ids.shape
    out = {}
    for role, parts in (("player", PLAYER), ("boss", BOSS)):
        m = np.isin(ids, parts)
        if m.sum() < 150:
            continue
        ys, xs = np.nonzero(m)
        cy, cx = (ys.min() + ys.max()) / 2, (xs.min() + xs.max()) / 2
        h = max(min_side, (ys.max() - ys.min() + 1) * (1 + pad))
        w = max(min_side, (xs.max() - xs.min() + 1) * (1 + pad))
        out[role] = (int(max(0, cy - h / 2)), int(min(H, cy + h / 2)), int(max(0, cx - w / 2)), int(min(W, cx + w / 2)))
    return out


def char_items(src, img, boxes):
    from vlm_gate import zoom
    return [([zoom(src, b, 336), zoom(img, b, 336)], Q_CHAR.format(who=WHO[r])) for r, b in boxes.items()]


def region_items(img, regs):
    return [([crop_img(img, r)], Q_OPEN if r["kind"] == "open" else Q_OBJ_T.format(what=Q_OBJ[r["cls"]])) for r in regs]


def region_score(regs, p_orig, p_cand):
    out = {}
    for kind in ("open", "obj"):
        idx = [i for i, r in enumerate(regs) if r["kind"] == kind and p_orig[i] >= 0.5]
        if not idx:
            out[f"{kind}_ok"], out[f"{kind}_n"] = None, 0
            continue
        w = np.array([regs[i]["area"] if kind == "obj" else 1.0 for i in idx], float)
        out[f"{kind}_ok"] = round(float((w * (np.asarray(p_cand)[idx] >= 0.5)).sum() / w.sum()), 4)
        out[f"{kind}_n"] = len(idx)
    return out


def orig_regions(vlm, tag, src, regs, boxes):
    """原帧的区域题 + 角色比对题（原帧对原帧）基线，缓存。"""
    path = f"{G}/sources/{tag}/regions.json"
    if os.path.exists(path):
        out = json.load(open(path))
        if "char_base" in out:
            return out
        p = []
    else:
        p = vlm.yes_prob(region_items(src, regs)) if regs else []
        out = {"regions": [{k: ([int(x) for x in v] if k == "box" else v) for k, v in r.items() if k != "mask"} for r in regs],
               "p_orig": [round(float(x), 4) for x in p]}
    pc = vlm.yes_prob(char_items(src, src, boxes)) if boxes else []
    out["char_base"] = {r: round(float(v), 4) for r, v in zip(boxes, pc)}
    tmp = f"{path}.{os.getpid()}.tmp"                         # 多个分片同时算同一个源：原子写，免得读到半个文件
    json.dump(out, open(tmp, "w"), indent=1)
    os.replace(tmp, path)
    return out


# ------------------------------------------------------------------ 跑 ----
def run(a):
    from vlm_gate import QwenVL
    files = sorted(f for f in glob.glob(f"{G}/cands/*/*.json") if not f.endswith(".vlm.json"))
    if a.only:
        files = [f for f in files if re.search(a.only, f)]
    files = [f for i, f in enumerate(files) if i % a.nshards == a.shard]
    def todo(f):
        v = f[:-5] + ".vlm.json"
        return not os.path.exists(v) or "charcrop" not in json.load(open(v))
    files = [f for f in files if todo(f)]
    print(f"[gate] shard {a.shard}/{a.nshards}: {len(files)} 张", flush=True)
    if not files:
        return
    vlm = QwenVL()
    print(f"[gate] Qwen3-VL-32B 加载 {vlm.load_s:.0f}s", flush=True)
    cache = {}
    for f in files:
        info = json.load(open(f))
        tag = info["tag"]
        if tag not in cache:
            src = Image.open(f"{G}/sources/{tag}/f0.png").convert("RGB")
            ids = np.load(f"{G}/sources/{tag}/ids0.npy")
            regs, boxes = regions(ids), char_boxes(ids)
            cache = {tag: (src, regs, boxes, orig_regions(vlm, tag, src, regs, boxes))}
        src, regs, boxes, ro = cache[tag]
        cand = Image.open(f[:-5] + ".png").convert("RGB")
        vpath = f[:-5] + ".vlm.json"
        t0 = time.time()
        if os.path.exists(vpath):                                   # 只补角色比对题
            out = json.load(open(vpath))
        else:
            sc = SCENES[info["scene"]]
            glob_it = [(q, [cand] if q == "scene" else [src, cand], Q[q].format(**sc)) for q in Q]
            p = vlm.yes_prob([(ims, tx) for _, ims, tx in glob_it] + region_items(cand, regs))
            res = {q: round(float(v), 4) for (q, _, _), v in zip(glob_it, p)}
            pc = [round(float(x), 4) for x in p[len(glob_it):]]
            out = {"p_yes": res, "region": region_score(regs, ro["p_orig"], pc), "p_regions": pc}
        pch = vlm.yes_prob(char_items(src, cand, boxes)) if boxes else []
        out["charcrop"] = {r: round(float(v), 4) for r, v in zip(boxes, pch)}
        out["sec"] = round(time.time() - t0, 2)
        json.dump(out, open(vpath, "w"), indent=1)
        print(f"[gate] {tag} {info['name']} {out['p_yes']} {out['region']} {out['charcrop']}", flush=True)


def auto_ok(info, o):
    m = info["metrics"]
    return {"edge": m["layout"].get("edge_recall", 0) >= GATE["edge_rel"] * o["layout"]["edge_recall"]}


def merge(a):
    rows = []
    for f in sorted(glob.glob(f"{G}/cands/*/*.json")):
        if f.endswith(".vlm.json") or not os.path.exists(f[:-5] + ".vlm.json"):
            continue
        info = json.load(open(f))
        o = json.load(open(f"{G}/sources/{info['tag']}/f0_metrics.json"))
        vj = json.load(open(f[:-5] + ".vlm.json"))
        if "region" not in vj:                                      # 冒烟 v1 的旧格式（没有区域题）
            continue
        v, rg = vj["p_yes"], vj["region"]
        base = json.load(open(f"{G}/sources/{info['tag']}/regions.json")).get("char_base", {})
        why = auto_ok(info, o)
        for q in ("layout", "only2", "scene"):
            why[q] = v[q] >= GATE[q]
        cc = vj.get("charcrop", {})
        why["chars"] = all(cc.get(r, 0) >= GATE["chars"] for r, b in base.items() if b >= 0.5)
        for k in ("open_ok", "obj_ok"):
            why[k] = rg.get(k) is None or rg[k] >= GATE[k]
        ro, ob = rg.get("open_ok"), rg.get("obj_ok")
        rows.append({"tag": info["tag"], "name": info["name"], "scene": info["scene"], "gen": info["gen"], "seed": info["seed"],
                     "checks": why, "vlm": v, "region": rg, "charcrop": vj.get("charcrop"), "pass": all(why.values()),
                     "coarse_p": info["metrics"]["struct"]["coarse_edge_precision"],
                     "edge_rel": round(info["metrics"]["layout"].get("edge_recall", 0) / o["layout"]["edge_recall"], 4),
                     "score": round(v["layout"] + v["scene"] + min(list(vj.get("charcrop", {}).values()) or [1]) + (1 if ro is None else ro)
                                    + (1 if ob is None else ob), 4)})
    picks = {}
    for r in rows:
        k = (r["tag"], r["scene"])
        if k not in picks or (r["pass"], r["score"]) > (picks[k]["pass"], picks[k]["score"]):
            picks[k] = r
    sel = list(picks.values())
    by_scene, by_gen, fails = {}, {}, {}
    for r in sel:
        by_scene.setdefault(r["scene"], []).append(r["pass"])
    for r in rows:
        by_gen.setdefault(r["gen"], []).append(r["pass"])
        for k, ok in r["checks"].items():
            fails[k] = fails.get(k, 0) + (not ok)
    out = {"gate": GATE, "rows": rows, "picks": {f"{t}/{s}": r["name"] for (t, s), r in picks.items() if r["pass"]},
           "failed_slots": sorted(f"{t}/{s}" for (t, s), r in picks.items() if not r["pass"]),
           "pass_rate_cands_by_gen": {g: f"{sum(v)}/{len(v)}" for g, v in sorted(by_gen.items())},
           "check_fail_counts": fails,
           "slot_pass_by_scene": {s: f"{sum(v)}/{len(v)}" for s, v in sorted(by_scene.items())},
           "picked_gen": {g: sum(1 for r in sel if r["pass"] and r["gen"] == g) for g in sorted(by_gen)}}
    json.dump(out, open(f"{G}/gate.json", "w"), indent=1)
    print(json.dumps({k: out[k] for k in ("pass_rate_cands_by_gen", "check_fail_counts", "slot_pass_by_scene", "picked_gen")},
                     indent=1))
    print("slots pass", len(out["picks"]), "/", len(sel))
    if a.jobs:
        jobs = []
        for s in json.load(open(f"{G}/sources.json")):
            jobs.append({"tag": s["tag"], "variant": "orig", "kind": "orig", "f0": f"{G}/sources/{s['tag']}/f0.png"})
            for sc in SCENES:
                r = picks.get((s["tag"], sc))
                if r and r["pass"]:
                    jobs.append({"tag": s["tag"], "variant": sc, "kind": "scene", "scene": sc, "gen": r["gen"],
                                 "seed": r["seed"], "f0": f"{G}/cands/{s['tag']}/{r['name']}.png"})
        json.dump(jobs, open(f"{G}/jobs.json", "w"), indent=1)
        print(len(jobs), "个 rollout 任务（含 orig）")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--shard", type=int, default=0)
    ap.add_argument("--nshards", type=int, default=1)
    ap.add_argument("--only", default=None)
    ap.add_argument("--merge", action="store_true")
    ap.add_argument("--jobs", action="store_true", help="--merge 时顺便写 jobs.json")
    a = ap.parse_args()
    merge(a) if a.merge else run(a)


if __name__ == "__main__":
    main()