#!/usr/bin/env python3
"""genbench3 · 首帧候选（PCE env pce_f3，FLUX.2 klein-9B；只读复用 firstframe_flux）。每个起点（sources.json）出：
  * style_<name>_s<seed>：整帧换画风（RESTYLE_V4：纯材质 / 渲染改写，不增删物体；6 个种子，门槛挑）；
  * char_<role>_<look>_s<seed>：只换一个角色的外观 —— 按 ids0 的角色掩码取包围框（长宽比 = W:H，留边），裁出来放大到
    FLUX 网格编辑（CHAR_PROMPT：同位置 / 姿势 / 体型 / 轮廓 / 武器，只换外观），缩回原尺寸，再只在掩码内（外扩 2 px、
    羽化 1 px，扣掉另一个角色）贴回原帧 ⇒ 掩码外逐像素不变，语义图不会被改；
  * --combo picks.json：在已通过门槛的画风首帧上再换角色（char 编辑的底图换成那张画风帧，提示词带上画风）。
每张候选旁边写 <name>.json：设置 + 自动指标（layout_metrics vs ids0、struct_metrics vs 原帧、角色显著度 / 填充率 / 外观变化）。
可续跑（已有 png+json 跳过），按 --shard i --nshards n 分卡。
usage: CUDA_VISIBLE_DEVICES=0 $FPY scripts/genbench3/make_f0.py --shard 0 --nshards 8 [--combo outputs/genbench3/picks_pass1.json]"""
import argparse
import json
import os
import sys
import time

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

M = "/opt/dlami/nvme/zhiyangdeng/maliang"
sys.path.insert(0, f"{M}/code/gc25d")
sys.path.insert(0, f"{M}/scripts")
import firstframe_flux as FF  # noqa: E402

G = f"{M}/outputs/genbench3"
STYLE_SEEDS, CHAR_SEEDS = (0, 1, 2, 3, 4, 5), (0, 1)
# genbench2 的 12 种（训练外观里都没有）+ 4 种常见游戏画风（pixel / paint 与 v9 的后处理画风相近 = 近分布对照）
# 主题画风写成「已有的东西换材质」：genbench2 的原描述（「巧克力墙、棒棒糖树」）等于让 FLUX 往空地上加东西（冒烟 candy / lego）
STYLES = {
    "ukiyoe": "a Japanese ukiyo-e woodblock print: bold black outlines, flat areas of Prussian blue, vermilion and ochre, "
              "visible wood-grain and washi paper texture, no shading gradients",
    "stainedglass": "a medieval stained-glass window: every existing surface, object and character rendered as pieces of "
                    "glowing translucent coloured glass separated by thick black lead lines, jewel-toned light",
    "lego": "a world made of plastic toy bricks, where the existing ground is re-textured as flat toy baseplates with "
            "studs, existing rocks, plants, walls and buildings become brick-built versions of themselves at the same "
            "places, characters become toy minifigures, glossy ABS plastic in bright primary colours",
    "voxel": "voxel art, where every existing surface, object and character is rebuilt from small uniform cubes at the "
             "same place and size, crisp flat-shaded cube faces, simple saturated colours",
    "knitted": "a knitted wool and felt craft diorama, where every existing surface, object and character is re-made "
               "from yarn and felt at the same place, visible stitches and fuzzy fibres, soft pastel textile colours",
    "porcelain": "Chinese blue-and-white porcelain painting: cobalt-blue brush patterns on glossy white glaze, everything "
                 "rendered only in shades of blue on white",
    "blueprint": "a technical blueprint drawing: thin white line-art on deep blue grid paper, hatching for shading, no "
                 "colour fills",
    "chalkboard": "a chalk drawing on a dark green school blackboard: dusty chalk strokes in white and pastel colours, "
                  "smudges and eraser marks",
    "candy": "a candy and dessert world, where the existing ground is re-textured as smooth pink frosting, existing rocks "
             "become glossy hard candies, existing plants become small lollipops and gummy shapes, existing walls and "
             "buildings become chocolate and wafer, with pastel sugary colours",
    "lava": "a volcanic hellscape, where the existing ground is re-textured as black obsidian and ash with glowing cracks "
            "of molten lava, existing objects become dark volcanic rock, floating embers, red-orange light from below",
    "mosaic": "an ancient Roman floor mosaic: everything made of small square stone and glass tesserae with visible grout "
              "lines, earthy terracotta, ochre and teal colours",
    "papyrus": "an ancient Egyptian tomb wall painting on papyrus: flat ochre, turquoise and lapis colours, thick dark "
               "outlines, cracked aged plaster texture",
    "winter": "a snowy winter night: thick fresh snow covering the existing ground, roofs and plants, cold blue moonlight "
              "and falling snowflakes",
    "watercolor": FF.STYLE_PRESETS["watercolor painting"],
    "pixelart": FF.STYLE_PRESETS["pixel-art retro game"],
    "neon": "a neon cyberpunk night: the existing surfaces become dark, wet and reflective, lit by glowing magenta and "
            "cyan neon light, existing objects and buildings get glowing neon edges, haze and rain",
}
NEAR_DIST = ("pixelart", "watercolor")          # 训练里有 pixel / paint 后处理：近分布
LOOKS = {
    "player": {
        "silverknight": "a knight in polished silver plate armour with a bright blue cape and a plumed silver helmet",
        "firemage": "a fire mage in flowing crimson-and-gold robes with a pointed hood and glowing orange runes",
        "ninja": "a ninja dressed entirely in black with a long red scarf and a black face mask",
        "ranger": "an elven ranger in a green hooded cloak and brown leather armour with long blond hair",
    },
    "boss": {
        "stonegolem": "a hulking golem made of grey granite boulders held together by glowing cyan crystals",
        "skeletonking": "a giant skeletal undead king in bone armour with a tattered purple cloak and a golden crown",
        "icetroll": "a giant ice troll with pale blue skin, white shaggy fur and frost crystals on its shoulders",
        "demonlord": "a demon warlord in jagged black-and-crimson armour with large curved horns and glowing yellow eyes",
    },
}
WHO = {"player": "the player hero (the smaller human character)", "boss": "the boss (the large enemy monster)"}
PARTS = {"player": FF.PLAYER_PARTS, "boss": FF.BOSS_PARTS}
CHAR_PROMPT = ("This is a close-up crop of a top-down isometric action RPG game screenshot. Redesign only {who} in the "
               "centre of the image as {look}. Keep exactly the same position, pose, body size, silhouette and facing "
               "direction, and keep the weapon in the same hand at the same angle. Keep the camera angle, lighting, shadows, "
               "the background and any other character exactly unchanged.{style}")
COMBO_STYLE = " The whole image is in the art style of {style}; draw the redesigned character in that same art style."
# 画风提示词：firstframe_flux 的 v1 会在空地上加东西（冒烟 candy：巧克力墙 / 棒棒糖 = 与语义图矛盾）⇒ 改成纯材质 / 渲染改写
# v4（冒烟：同种子下大轮廓多出的比 v3 少，浮世绘 coarse precision 0.58 → 0.74）
RESTYLE_V4 = ("Re-render this exact top-down isometric action RPG game screenshot in a new art style: {style}. This is a "
              "pure material and rendering change of the existing image: every region, object, building, wall, rock, "
              "plant, path and character stays exactly where it is with exactly the same outline, size and pose, and the "
              "two characters keep their weapons. Do not add anything new and do not remove anything; large empty areas "
              "of ground stay empty and are only re-textured. Same camera, same composition.")
PAD, MIN_CROP_W, FEATHER = 1.7, 176, 1.5
# 贴回区域 = 胶囊掩码外扩 CORE px ∪（掩码外 REACH·角色尺寸以内、FLUX 改动 ΔE > DIFF_DE 的连通区）：语义图的角色是胶囊代理，
# 真实轮廓会伸出掩码（熔岩巨人的拳头）—— 只贴掩码内会留下旧角色的残块（冒烟 icetroll）
CORE, REACH, DIFF_DE = 2, 0.18, 14.0
AR = FF.W_OUT / FF.H_OUT


# ------------------------------------------------------------------ 指标 ----
def ring(mask, other, r0=3, r1=12):
    """角色掩码外 r0..r1 px 的一圈（扣掉另一个角色）= 角色周围的背景。"""
    out = ndimage.binary_dilation(mask, iterations=r1) & ~ndimage.binary_dilation(mask, iterations=r0)
    return out & ~other


def char_stats(rgb, ids):
    """每个角色：显著度 = 掩码内 Lab 均值与周围一圈的 ΔE；填充率 = 掩码内与周围背景（一圈的中位色）ΔE < 12 的比例
    （越高 ⇒ 掩码里画的是背景，角色没填满语义轮廓）。"""
    L = FF._lab(rgb)
    out = {}
    for role, parts in PARTS.items():
        m = np.isin(ids, parts)
        o = np.isin(ids, PARTS["boss" if role == "player" else "player"])
        rg = ring(m, o)
        if m.sum() < 50 or rg.sum() < 50:
            continue
        bg = np.median(L[rg], 0)
        out[role] = {"salience": round(float(np.linalg.norm(L[m].mean(0) - L[rg].mean(0))), 2),
                     "bg_fill": round(float((np.linalg.norm(L[m] - bg, axis=-1) < 12).mean()), 4), "px": int(m.sum())}
    return out


def coarse_struct(src, out):
    """粗尺度（½ 分辨率 + σ2 模糊，压掉画风纹理，只剩物体轮廓）的 struct_metrics：precision 低 ⇒ 多出了物体 / 轮廓。"""
    def half(x):
        return np.asarray(Image.fromarray(x).resize((x.shape[1] // 2, x.shape[0] // 2), Image.BILINEAR))
    a, b = (ndimage.gaussian_filter(half(x).astype(np.float32), (2, 2, 0)).clip(0, 255).astype(np.uint8) for x in (src, out))
    return {f"coarse_{k}": round(v, 4) for k, v in FF.struct_metrics(a, b, tol=2).items()}


def metrics(rgb, src, ids):
    m = {"layout": {k: (round(v, 4) if isinstance(v, float) else v) for k, v in FF.layout_metrics(rgb, ids).items()},
         "struct": {k: round(v, 4) for k, v in FF.struct_metrics(src, rgb).items()}, "chars": char_stats(rgb, ids)}
    m["struct"].update(coarse_struct(src, rgb))
    return m


# ------------------------------------------------------------------ 角色编辑 ----
def crop_box(mask):
    ys, xs = np.nonzero(mask)
    cy, cx = (ys.min() + ys.max()) / 2, (xs.min() + xs.max()) / 2
    bw, bh = xs.max() - xs.min() + 1, ys.max() - ys.min() + 1
    w = int(round(min(FF.W_OUT, max(MIN_CROP_W, PAD * max(bw, bh * AR)))))
    h = int(round(w / AR))
    if h > FF.H_OUT:
        h, w = FF.H_OUT, int(round(FF.H_OUT * AR))
    x0 = int(np.clip(round(cx - w / 2), 0, FF.W_OUT - w))
    y0 = int(np.clip(round(cy - h / 2), 0, FF.H_OUT - h))
    return y0, y0 + h, x0, x0 + w


def char_edit(ff, base, ids, role, look, seed, style=None):
    import torch
    m = np.isin(ids, PARTS[role])
    other = np.isin(ids, PARTS["boss" if role == "player" else "player"])
    y0, y1, x0, x1 = crop_box(m)
    crop = base[y0:y1, x0:x1]
    prompt = CHAR_PROMPT.format(who=WHO[role], look=look, style=COMBO_STYLE.format(style=style) if style else "")
    t0 = time.time()
    with torch.no_grad():
        big = ff.net(prompt, [FF.to_gen(crop)], FF.DEFAULT_STEPS["flux2"], seed)
    torch.cuda.synchronize()
    sec = time.time() - t0
    ed = FF.area_resize(big, x1 - x0, y1 - y0)
    mc, oc = m[y0:y1, x0:x1], other[y0:y1, x0:x1]
    ys, xs = np.nonzero(mc)
    reach = max(4, int(round(REACH * max(ys.max() - ys.min(), xs.max() - xs.min()))))
    diff = ndimage.gaussian_filter(np.linalg.norm(FF._lab(ed) - FF._lab(crop), axis=-1), 1.0) > DIFF_DE
    near = ndimage.binary_dilation(mc, iterations=reach)
    lab, _ = ndimage.label(diff & near)
    touch = np.unique(lab[ndimage.binary_dilation(mc, iterations=CORE) & (lab > 0)])
    grown = np.isin(lab, touch[touch > 0])                         # 与掩码相连的改动区（旧角色伸出掩码的部分）
    reg = ndimage.binary_fill_holes(ndimage.binary_dilation(mc, iterations=CORE) | grown)
    reg &= ~ndimage.binary_dilation(oc, iterations=1)
    a = np.clip(ndimage.gaussian_filter(reg.astype(np.float32), FEATHER), 0, 1)
    a[oc] = 0.0
    a[mc] = 1.0
    out = base.copy()
    out[y0:y1, x0:x1] = (a[..., None] * ed + (1 - a[..., None]) * crop).round().clip(0, 255).astype(np.uint8)
    Lb, Lo = FF._lab(base), FF._lab(out)
    info = {"prompt": prompt, "box": [y0, y1, x0, x1], "sec": round(sec, 2), "reach_px": reach,
            "grown_px": int((grown & ~mc).sum()), "region_px": int(reg.sum()), "mask_px": int(mc.sum()),
            "change_dE": round(float(np.linalg.norm(Lo[m] - Lb[m], axis=-1).mean()), 2)}
    return out, ed, info


# ------------------------------------------------------------------ 任务 ----
def jobs(sources, combo):
    out = []
    for s in sources:
        tag = s["tag"]
        if combo is None:
            for st in STYLES:
                for seed in STYLE_SEEDS:
                    out.append({"tag": tag, "kind": "style", "style": st, "seed": seed, "name": f"style_{st}_s{seed}"})
            for role, looks in LOOKS.items():
                for lk in looks:
                    for seed in CHAR_SEEDS:
                        out.append({"tag": tag, "kind": "char", "role": role, "look": lk, "seed": seed,
                                    "name": f"char_{role}_{lk}_s{seed}"})
        else:
            for c in combo.get(tag, []):                  # {"base": style 候选名, "style": st, "role":..., "look":...}
                for seed in CHAR_SEEDS:
                    out.append({"tag": tag, "kind": "combo", "base": c["base"], "style": c["style"], "role": c["role"],
                                "look": c["look"], "seed": seed,
                                "name": f"combo_{c['style']}+{c['role']}_{c['look']}_s{seed}"})
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--shard", type=int, default=0)
    ap.add_argument("--nshards", type=int, default=1)
    ap.add_argument("--combo", default=None, help="picks json：{tag: [{base, style, role, look}]}")
    ap.add_argument("--only", default=None, help="正则，只跑 tag/name 匹配的（冒烟用）")
    a = ap.parse_args()
    sources = json.load(open(f"{G}/sources.json"))
    combo = json.load(open(a.combo)) if a.combo else None
    todo = [j for i, j in enumerate(jobs(sources, combo)) if i % a.nshards == a.shard]
    if a.only:
        import re
        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_f0] shard {a.shard}/{a.nshards}: {len(todo)} 张待生成", flush=True)
    if not todo:
        return
    ff = FF.FirstFrameFlux("klein9b")
    print(f"[make_f0] 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 = np.asarray(Image.open(f"{d}/f0.png").convert("RGB"))
            ids = np.load(f"{d}/ids0.npy")
            cache[tag] = (src, ids)
            mpath = f"{d}/f0_metrics.json"
            if not os.path.exists(mpath):
                json.dump(metrics(src, src, ids), open(mpath, "w"), indent=1)
        src, ids = cache[tag]
        od = f"{G}/cands/{tag}"
        os.makedirs(od, exist_ok=True)
        info = dict(j)
        if j["kind"] == "style":
            out = ff.generate(RESTYLE_V4.format(style=STYLES[j["style"]]), [FF.to_gen(src)], seed=j["seed"])
            info.update(prompt=ff.last["prompt"], sec=ff.last["sec"], near_dist=j["style"] in NEAR_DIST)
        else:
            base = src if j["kind"] == "char" else np.asarray(Image.open(f"{od}/{j['base']}.png").convert("RGB"))
            out, ed, ci = char_edit(ff, base, ids, j["role"], LOOKS[j["role"]][j["look"]], j["seed"],
                                    STYLES[j["style"]] if j["kind"] == "combo" else None)
            info.update(ci)
            Image.fromarray(ed).save(f"{od}/{j['name']}_crop.png")
        Image.fromarray(out).save(f"{od}/{j['name']}.png")
        info["metrics"] = metrics(out, src, ids)
        json.dump(info, open(f"{od}/{j['name']}.json", "w"), indent=1)
        print(f"[make_f0] {tag} {j['name']} {info['sec']}s", flush=True)


if __name__ == "__main__":
    main()