#!/usr/bin/env python3
"""genbench4 · 首帧门槛 v2（PCE env pce_f3，Qwen3-VL-32B）。v1（gate.py）的通过帧经 3 个独立复核 agent 检查仍常有语义矛盾
（随机 6 张里 4 张）：train_v9_04 的 9 张里 Boss 全被画成石头 / 水晶 / 圆顶 / 蘑菇（逐角色比对被「形状像角色的物体」骗过）；多出小人
（only2 卡在 0.5 放行）；铺地变水面 / 熔岩池、画面边缘加塔 / 小屋（open 格只取 ≥ 92 % 空地的内部格，边缘几乎没查）；train_v9_10 的玻璃
建筑被画成可走的地面（全图没有空地格，open 检查根本没跑）。v2 在 v1 的整图题（复用 <cand>.vlm.json 的 p_yes / charcrop）上加：
  * 格子按语义类分题：terrain 格「只有天然地面」、paving 格「只有铺装 / 搭建的地板，没变成水 / 熔岩」、architecture 格「是建筑 / 构筑物
    的一部分，不是地面 / 植物 / 散石」；内部格纯度 ≥ 0.92，**边缘一圈的格 ≥ 0.75**（FLUX 爱在边缘加框）；离角色 ≥ 12 px；
  * 角色：每个角色的放大框再问「中间是不是一个活的角色（有身体 / 头 / 四肢），不是石头 / 植物 / 雕像 / 建筑 / 物件」；整图 chars 题恢复为
    门槛（以原帧对原帧为基线：基线 < 0.5 的源不判）；
  * 阈值收紧：layout ≥ 0.5、only2 ≥ 0.8、scene ≥ 0.6、open_ok ≥ 0.9、arch_ok ≥ 0.8、obj_ok ≥ 0.85（obj 沿用 v1 的连通块题）。
所有格 / 角色题都在原帧上问一遍作基线（sources/<tag>/regions2.json），只有原帧答 Yes 的才计分。
输出 <cand>.vlm2.json；--merge ⇒ gate2.json + jobs.json（与 v1 的 jobs 不同的槽才需要重跑 rollout）。
usage: CUDA_VISIBLE_DEVICES=0 $FPY scripts/genbench4/gate2.py --shard 0 --nshards 8 ；$FPY scripts/genbench4/gate2.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"]
import gate as G1  # noqa: E402
from scenes import SCENES  # noqa: E402

CROP = G1.CROP
Q_CELL = {
    1: CROP + "Does this crop show only natural open ground (for example dirt, sand, snow, grass, moss, mud, ash, cracked "
              "earth, rock floor, ice, or shallow water and puddles), possibly with cracks, pebbles, small debris or small "
              "tufts, with no walls, buildings, towers, fences, cliffs, rock formations, trees, large plants, large objects "
              "or characters standing in it? Answer Yes or No.",
    2: CROP + "Does this crop show only a paved or built floor (for example tiles, flagstones, cobblestones, bricks, wooden "
              "planks or a boardwalk, metal plates or concrete), possibly worn, cracked, wet or with small debris, with no "
              "walls, buildings, towers, fences, trees, large plants, large objects or characters standing on it, and not "
              "turned into open water or lava? Answer Yes or No.",
    3: CROP + "Does this crop show part of a building or built structure (for example a wall, roof, tower, pillar, stairs, "
              "bridge, glass dome or glass panels), rather than open ground, plants or loose rocks? Answer Yes or No.",
}
Q_CREATURE = ("This is a close-up crop from a top-down isometric action RPG game. Is there a living character (a person, "
              "monster or creature with a body, head and limbs) in the middle of this crop, and not just a rock, plant, "
              "statue, building or object? Answer Yes or No.")
CELL, PURE_IN, PURE_BORDER, CHAR_GAP = 96, 0.92, 0.75, 12
MAX_OPEN, MAX_ARCH = 14, 6
GATE = {"layout": 0.5, "only2": 0.8, "scene": 0.6, "open_ok": 0.9, "arch_ok": 0.8, "obj_ok": 0.85, "edge_rel": 0.4,
        "chars": 0.5, "charcrop": 0.5, "creature": 0.5}
# B 档（角色安全、布局适中）：A 档（GATE 全过）在 220 槽里只有 12 个 —— FLUX 换重主题场景几乎总在边缘 / 铺地上加装饰，补边放大（padk_test）
# 与 RePaint 修补（repair.py，klein 4 步会照周围把城墙 / 蘑菇重画回去）都没解决。B 档：角色题（creature / charcrop / only2 ≥ 0.8）、
# scene、arch 照旧严，空地格允许 ≤ 25 % 不过、layout ≥ 0.3、obj ≥ 0.7，不用整图 chars 题（失败 1184 张，多数是 VLM 在小 / 暗角色上认不出）
TIER_B = {"open_ok": 0.75, "layout": 0.3, "obj_ok": 0.7}
TIER_B_SKIP = ("chars",)


def cells(ids):
    """[{kind: cell, cls ∈ {1,2,3}, box, border}]：每格取 {terrain, paving, architecture} 里的多数类，纯度够（边缘格放宽）才收。"""
    H, W = ids.shape
    chars = ndimage.binary_dilation(np.isin(ids, G1.PLAYER + G1.BOSS), iterations=CHAR_GAP)
    oy, ox = (H % CELL) // 2, (W % CELL) // 2
    ys, xs = list(range(oy, H - CELL + 1, CELL)), list(range(ox, W - CELL + 1, CELL))
    out = {1: [], 2: [], 3: []}
    for i, y in enumerate(ys):
        for j, x in enumerate(xs):
            sl = (slice(y, y + CELL), slice(x, x + CELL))
            if chars[sl].any():
                continue
            border = i in (0, len(ys) - 1) or j in (0, len(xs) - 1)
            blk = ids[sl]
            fr = {c: float((blk == c).mean()) for c in (1, 2, 3)}
            c = max(fr, key=fr.get)
            pure = PURE_BORDER if border else PURE_IN
            # 边缘格：terrain + paving 合起来算「空地」纯度（两者都是可走地面），题按多数类问
            if c in (1, 2) and (fr[1] + fr[2] if border else fr[c]) >= pure or c == 3 and fr[3] >= pure:
                out[c].append({"kind": "cell", "cls": c, "box": (y, y + CELL, x, x + CELL), "border": border, "area": CELL * CELL})
    sel = []
    open_ = out[1] + out[2]
    if len(open_) > MAX_OPEN:                                        # 边缘格优先（FLUX 加框的地方），其余均匀取
        b = [r for r in open_ if r["border"]]
        n = [r for r in open_ if not r["border"]]
        b = [b[k] for k in np.linspace(0, len(b) - 1, min(len(b), MAX_OPEN // 2 + 2)).round().astype(int)] if b else []
        k = MAX_OPEN - len(b)
        n = [n[q] for q in np.linspace(0, len(n) - 1, min(len(n), k)).round().astype(int)] if n and k > 0 else []
        open_ = b + n
    arch = out[3]
    if len(arch) > MAX_ARCH:
        arch = [arch[q] for q in np.linspace(0, len(arch) - 1, MAX_ARCH).round().astype(int)]
    sel = open_ + arch
    return sel


def cell_items(img, cs):
    return [([G1.crop_img(img, r)], Q_CELL[r["cls"]]) for r in cs]


def creature_items(img, boxes):
    from vlm_gate import zoom
    return [([zoom(img, b, 336)], Q_CREATURE) for b in boxes.values()]


def baseline(vlm, tag, src, cs, boxes):
    path = f"{G}/sources/{tag}/regions2.json"
    if os.path.exists(path):
        return json.load(open(path))
    p = vlm.yes_prob(cell_items(src, cs) + creature_items(src, boxes) + [([src, src], G1.Q["chars"])])
    nc, nb = len(cs), len(boxes)
    out = {"cells": [{"cls": r["cls"], "box": [int(v) for v in r["box"]], "border": r["border"]} for r in cs],
           "p_cells": [round(float(x), 4) for x in p[:nc]],
           "creature": {r: round(float(x), 4) for r, x in zip(boxes, p[nc:nc + nb])},
           "chars": round(float(p[nc + nb]), 4)}
    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", ".vlm2.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]
    files = [f for f in files if os.path.exists(f[:-5] + ".vlm.json") and not os.path.exists(f[:-5] + ".vlm2.json")]
    print(f"[gate2] shard {a.shard}/{a.nshards}: {len(files)} 张", flush=True)
    if not files:
        return
    vlm = QwenVL()
    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")
            cs, boxes = cells(ids), G1.char_boxes(ids)
            cache = {tag: (cs, boxes, baseline(vlm, tag, src, cs, boxes))}
        cs, boxes, base = cache[tag]
        cand = Image.open(f[:-5] + ".png").convert("RGB")
        t0 = time.time()
        p = vlm.yes_prob(cell_items(cand, cs) + creature_items(cand, boxes))
        nc = len(cs)
        out = {"p_cells": [round(float(x), 4) for x in p[:nc]],
               "creature": {r: round(float(x), 4) for r, x in zip(boxes, p[nc:])}, "sec": round(time.time() - t0, 2)}
        json.dump(out, open(f[:-5] + ".vlm2.json", "w"), indent=1)
        print(f"[gate2] {tag} {info['name']} cells {sum(x >= 0.5 for x in out['p_cells'])}/{nc} {out['creature']}", flush=True)


def cell_ok(base, pc, kinds):
    idx = [i for i, c in enumerate(base["cells"]) if c["cls"] in kinds and base["p_cells"][i] >= 0.5]
    if not idx:
        return None, 0
    return round(float(np.mean([pc[i] >= 0.5 for i in idx])), 4), len(idx)


def merge(a):
    rows = []
    for f in sorted(glob.glob(f"{G}/cands/*/*.json")):
        if f.endswith((".vlm.json", ".vlm2.json")) or not os.path.exists(f[:-5] + ".vlm2.json"):
            continue
        info = json.load(open(f))
        tag = info["tag"]
        o = json.load(open(f"{G}/sources/{tag}/f0_metrics.json"))
        v1 = json.load(open(f[:-5] + ".vlm.json"))
        v2 = json.load(open(f[:-5] + ".vlm2.json"))
        b1 = json.load(open(f"{G}/sources/{tag}/regions.json"))
        b2 = json.load(open(f"{G}/sources/{tag}/regions2.json"))
        v, rg = v1["p_yes"], v1["region"]
        ck = {"edge": info["metrics"]["layout"].get("edge_recall", 0) >= GATE["edge_rel"] * o["layout"]["edge_recall"]}
        for q in ("layout", "only2", "scene"):
            ck[q] = v[q] >= GATE[q]
        ck["chars"] = b2["chars"] < 0.5 or v["chars"] >= GATE["chars"]
        ck["charcrop"] = all(v1["charcrop"].get(r, 0) >= GATE["charcrop"] for r, x in b1["char_base"].items() if x >= 0.5)
        ck["creature"] = all(v2["creature"].get(r, 0) >= GATE["creature"] for r, x in b2["creature"].items() if x >= 0.5)
        open_ok, n_open = cell_ok(b2, v2["p_cells"], (1, 2))
        arch_ok, n_arch = cell_ok(b2, v2["p_cells"], (3,))
        ck["open_ok"] = open_ok is None or open_ok >= GATE["open_ok"]
        ck["arch_ok"] = arch_ok is None or arch_ok >= GATE["arch_ok"]
        ck["obj_ok"] = rg.get("obj_ok") is None or rg["obj_ok"] >= GATE["obj_ok"]
        ckb = dict(ck)
        ckb["open_ok"] = open_ok is None or open_ok >= TIER_B["open_ok"]
        ckb["layout"] = v["layout"] >= TIER_B["layout"]
        ckb["obj_ok"] = rg.get("obj_ok") is None or rg["obj_ok"] >= TIER_B["obj_ok"]
        for k in TIER_B_SKIP:
            ckb.pop(k)
        tier = "A" if all(ck.values()) else "B" if all(ckb.values()) else None
        rows.append({"tag": tag, "name": info["name"], "scene": info["scene"], "gen": info["gen"], "seed": info["seed"],
                     "checks": ck, "pass": tier is not None, "tier": tier, "vlm": v, "open_ok": open_ok, "n_open": n_open, "arch_ok": arch_ok,
                     "n_arch": n_arch, "obj_ok": rg.get("obj_ok"), "creature": v2["creature"], "charcrop": v1["charcrop"],
                     "score": round(v["layout"] + v["scene"] + v["only2"] + (1 if open_ok is None else open_ok)
                                    + (1 if arch_ok is None else arch_ok) + (1 if rg.get("obj_ok") is None else rg["obj_ok"]), 4)})
    picks = {}
    for r in rows:
        k = (r["tag"], r["scene"])
        rank = lambda x: ({"A": 2, "B": 1}.get(x["tier"], 0), x["score"])
        if k not in picks or rank(r) > rank(picks[k]):
            picks[k] = r
    sel = list(picks.values())
    by_scene, fails = {}, {}
    print("passing candidates by tier:", {x: sum(1 for r in rows if r["tier"] == x) for x in ("A", "B")})
    for r in sel:
        by_scene.setdefault(r["scene"], []).append(r["pass"])
    for r in rows:
        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"]),
           "check_fail_counts": fails, "pass_cands": f"{sum(r['pass'] for r in rows)}/{len(rows)}",
           "slot_pass_by_scene": {s: f"{sum(v)}/{len(v)}" for s, v in sorted(by_scene.items())},
           "slot_tier": {f"{t}/{s}": r["tier"] for (t, s), r in picks.items() if r["pass"]},
           "tier_counts": {x: sum(1 for r in sel if r["tier"] == x) for x in ("A", "B")}}
    json.dump(out, open(f"{G}/gate2.json", "w"), indent=1)
    print(json.dumps({k: out[k] for k in ("pass_cands", "check_fail_counts", "slot_pass_by_scene", "tier_counts")}, 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"], "tier": r["tier"], "f0": f"{G}/cands/{s['tag']}/{r['name']}.png"})
        json.dump(jobs, open(f"{G}/jobs2.json", "w"), indent=1)
        print(len(jobs), "个 rollout 任务（含 orig）→ jobs2.json")


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")
    a = ap.parse_args()
    merge(a) if a.merge else run(a)


if __name__ == "__main__":
    main()