#!/usr/bin/env python3
"""genbench4 · 换场景的首帧候选（PCE env pce_f3，FLUX.2 klein-9B，4 步）。每个起点 × 每个场景（scenes.SCENES）两种生成法：
  * edit_<scene>_s<seed>：编辑真实帧 0 —— RETHEME：每个语义类按 scenes 的对照表换成该场景的同类东西（只写帧 0 里出现的类），
    布局 / 相机 / 两个角色不动，不增删物体（genbench3 的 RESTYLE_V4 句式）；
  * edit2_…：RETHEME + 帧 0 没有的类写「没有、不许加」+ 不许沿画面边缘加墙 / 悬崖；pad_…：edit2 提示词 + 原帧四周反射补边到
    4/3 倍画布再裁回中间（默认只跑 edit2 ×2 + pad ×4；冒烟见 outputs/genbench4_smoke：edit / map 违规太多，弃用）；
  * map_<scene>_s<seed>：从帧 0 的语义图直接生成（firstframe_flux.map2frame 的 v4 模板 + 白武器色图 + 铺地网格），色表里每类
    的含义换成该场景的东西 ⇒ 区域 / 角色由语义图定，整张图（含角色的外观）都是新画的。
每张候选旁边写 <name>.json：设置 + 自动指标（genbench3 make_f0.metrics：layout_metrics vs ids0、struct vs 原帧、角色显著度 / 填充率）。
可续跑（已有 json 跳过），--shard / --nshards 分卡。
usage: CUDA_VISIBLE_DEVICES=0 $FPY scripts/genbench4/make_scene.py --shard 0 --nshards 8 [--only REGEX]"""
import argparse
import json
import os
import re
import sys
import time

import numpy as np
from PIL import Image

M = "/opt/dlami/nvme/zhiyangdeng/maliang"
sys.path[:0] = [f"{M}/code/gc25d", f"{M}/scripts/genbench3", f"{M}/scripts/genbench4"]
import firstframe_flux as FF  # noqa: E402
import make_f0 as MF  # noqa: E402  genbench3：metrics()
from scenes import CLS_COLOR, CLS_SUBJ, SCENES  # noqa: E402

G = os.environ.get("GB4", f"{M}/outputs/genbench4")
GENS = tuple(os.environ.get("GB4_GENS", "edit2,pad").split(","))
SEEDS = {"edit": (0, 1, 2), "edit2": (0, 1), "map": (0, 1, 2), "pad": (0, 1, 2, 3)}
RETHEME = ("Turn this top-down isometric action RPG game screenshot into a different level of the same game: {title}. "
           "Keep exactly the same camera, composition and layout: every region keeps its exact position, outline and size "
           "and only changes what it is made of: {mapping}. The two characters stay exactly where they are with the same "
           "poses, sizes, weapons and appearance. Do not add anything new and do not remove anything: no new walls, "
           "buildings, plants, objects or characters; large open areas stay open and are only re-textured. {light}")
MAP_SCENE = " Scene: the whole level is {title}. {light}"
# edit2（冒烟 v1：FLUX 为了「像城堡 / 沼泽」把石头画成城墙 / 小屋、沿画面边缘加一圈墙 = 与语义矛盾）：帧 0 里没有的类明确写
# 「没有、也不许加」，并禁止在画面边缘加墙 / 悬崖 / 建筑
ABSENT = {3: "buildings, walls, towers, gates, huts or fences", 5: "trees or plants", 4: "rocks or boulders"}
BORDER = (" Do not frame the scene with new walls, cliffs or buildings along the image borders; the borders show the same "
          "kind of region as now.")


def pad_generate(ff, prompt, src, seed):
    """pad（冒烟 padtest：edit / edit2 都会沿画面边缘加一圈城墙 / 岩壁 / 树 = 与语义矛盾）：原帧放到 4/3 倍画布中间，四周
    反射补边（模糊补边 / 灰边仍会往中间加框），edit2 提示词生成 4/3 倍网格，裁回中间 ⇒ 「框」落在补边里被裁掉。"""
    import torch
    w, h = FF.W_GEN, FF.H_GEN
    pw, ph = w * 4 // 3, h * 4 // 3
    x0, y0 = (pw - w) // 2, (ph - h) // 2
    ref = Image.fromarray(np.pad(np.asarray(FF.to_gen(src)), ((y0, ph - h - y0), (x0, pw - w - x0), (0, 0)), mode="reflect"))
    FF.W_GEN, FF.H_GEN = pw, ph                      # _Flux2 的生成网格取模块全局
    try:
        torch.cuda.synchronize()
        t0 = time.time()
        with torch.no_grad():
            big = ff.net(prompt, [ref], FF.DEFAULT_STEPS["flux2"], seed)
        torch.cuda.synchronize()
        sec = round(time.time() - t0, 3)
    finally:
        FF.W_GEN, FF.H_GEN = w, h
    return FF.area_resize(big[y0:y0 + h, x0:x0 + w], FF.W_OUT, FF.H_OUT), sec


def present(ids):
    n = np.bincount(np.asarray(ids).reshape(-1), minlength=256)
    return [c for c in range(7) if n[c] >= FF.MIN_PX], n


def retheme_prompt(ids, scene, v2=False):
    sc = SCENES[scene]
    cls, _ = present(ids)
    mapping = "; ".join(f"{CLS_SUBJ[c]} become {sc['cls'][c]}" for c in cls)
    out = RETHEME.format(title=sc["title"], mapping=mapping, light=sc["light"])
    if v2:
        miss = [ABSENT[c] for c in ABSENT if c not in cls]
        if miss:
            out += " This part of the level contains no " + ", no ".join(miss) + "; do not add any."
        out += BORDER
    return out


def scene_legend(ids, scene):
    """firstframe_flux.map_legend（v4）的场景版：类的含义换成该场景里的东西，角色句子照抄。"""
    sc = SCENES[scene]
    cls, n = present(ids)
    txt = "Color legend of the layout map: " + "; ".join(f"{CLS_COLOR[c]} = {sc['cls'][c]}" for c in cls)
    if n[FF.UNKNOWN_ID] >= FF.MIN_PX:
        txt += "; black = unknown area, fill it with scenery that continues the surroundings"
    txt += "."
    has_p, has_b = n[list(FF.PLAYER_PARTS)].sum() >= FF.MIN_PX, n[list(FF.BOSS_PARTS)].sum() >= FF.MIN_PX
    if has_p:
        txt += (" The blue humanoid figure is the player character, an armored human warrior"
                + (" holding a sword (the white stroke)" if n[16] else "")
                + "; draw him exactly where the blue figure is, at the same size and in the same pose.")
    if has_b:
        txt += (" The red humanoid figure is the boss, a hulking monster"
                + (" wielding a weapon (the white stroke)" if n[21] else "")
                + "; draw it exactly where the red figure is, at the same size and in the same pose.")
    if has_p or has_b:
        txt += (" These are the only characters; do not add any other people or creatures." if has_p and has_b else
                " This is the only character; do not add any other people or creatures.")
        txt += (" The flat blue and red colors only mark who is who: give the characters natural colors and materials "
                "that fit the scene, and render weapons as ordinary steel or wood, not glowing.")
    return txt


def map_prompt(ids, scene):
    return FF.MAP_PROMPTS["v4"].format(legend=scene_legend(ids, scene)) + MAP_SCENE.format(**SCENES[scene])


def jobs(sources):
    return [{"tag": s["tag"], "kind": "scene", "gen": g, "scene": sc, "seed": seed, "name": f"{g}_{sc}_s{seed}"}
            for s in sources for sc in SCENES for g in GENS for seed in SEEDS[g]]


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, help="正则，只跑 tag/name 匹配的（冒烟用）")
    ap.add_argument("--extra", default=None, help="补种子：gate.json 里没通过的 (起点, 场景) 槽再出 pad 种子 A..B（如 4-9）")
    a = ap.parse_args()
    sources = json.load(open(f"{G}/sources.json"))
    if a.extra:
        lo, hi = map(int, a.extra.split("-"))
        gj = json.load(open(f"{G}/gate.json"))
        expr = {f"{r['tag']}/{r['scene']}" for r in gj["rows"] if r["vlm"]["scene"] >= 0.5}   # 场景表达得出来的槽才补
        failed = set(gj["failed_slots"]) & expr
        all_jobs = [{"tag": t, "kind": "scene", "gen": "pad", "scene": sc, "seed": seed, "name": f"pad_{sc}_s{seed}"}
                    for t, sc in (x.split("/") for x in sorted(failed)) for seed in range(lo, hi + 1)]
    else:
        all_jobs = jobs(sources)
    todo = [j for i, j in enumerate(all_jobs) if i % a.nshards == a.shard]
    if a.only:
        todo = [j for j in todo if re.search(a.only, f"{j['tag']}/{j['name']}")]
    todo = [j for j in todo if not os.path.exists(f"{G}/cands/{j['tag']}/{j['name']}.json")]
    print(f"[make_scene] shard {a.shard}/{a.nshards}: {len(todo)} 张待生成", flush=True)
    if not todo:
        return
    ff = FF.FirstFrameFlux("klein9b")
    print(f"[make_scene] FLUX 加载 {ff.load_sec:.0f}s", flush=True)
    cache = {}
    for j in todo:
        tag = j["tag"]
        if tag not in cache:
            d = f"{G}/sources/{tag}"
            src, ids = np.asarray(Image.open(f"{d}/f0.png").convert("RGB")), np.load(f"{d}/ids0.npy")
            cache[tag] = (src, ids, FF.to_gen(src), FF.to_gen(FF.flux_map(ids, "white", tiles=True), nearest=True))
            if not os.path.exists(f"{d}/f0_metrics.json"):
                json.dump(MF.metrics(src, src, ids), open(f"{d}/f0_metrics.json", "w"), indent=1)
                Image.fromarray(FF.flux_map(ids, "white", tiles=True)).save(f"{d}/map0.png")
        src, ids, src_g, map_g = cache[tag]
        od = f"{G}/cands/{tag}"
        os.makedirs(od, exist_ok=True)
        if j["gen"] == "pad":
            prompt = retheme_prompt(ids, j["scene"], True)
            out, sec = pad_generate(ff, prompt, src, j["seed"])
        elif j["gen"] in ("edit", "edit2"):
            out = ff.generate(retheme_prompt(ids, j["scene"], j["gen"] == "edit2"), [src_g], seed=j["seed"])
            prompt, sec = ff.last["prompt"], ff.last["sec"]
        else:
            out = ff.generate(map_prompt(ids, j["scene"]), [map_g], seed=j["seed"])
            prompt, sec = ff.last["prompt"], ff.last["sec"]
        info = dict(j, prompt=prompt, sec=sec)
        Image.fromarray(out).save(f"{od}/{j['name']}.png")
        info["metrics"] = MF.metrics(out, src, ids)
        json.dump(info, open(f"{od}/{j['name']}.json", "w"), indent=1)
        print(f"[make_scene] {tag} {j['name']} {info['sec']}s", flush=True)


if __name__ == "__main__":
    main()