#!/usr/bin/env python3
"""双向 SFT 的泛化测试 L2 / L3（L1 = sample.py --val_n 20，已跑）。
L2 见过的地图 + 没见过的首帧风格：clip 的真 ids / prompts，首帧 = FLUX restyle 成训练里没有的画风。
L3 没见过的地图 + 没见过的首帧：procmap 程序生成的新世界（真相机 + 真胶囊运动）+ FLUX map2frame 生成的首帧。
usage: genbench_run.py [--ckpt ...] [--gpus 4,5] [--flux_gpu 3]"""
import argparse
import json
import os
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor

M = "/opt/dlami/nvme/zhiyangdeng/maliang"
G = f"{M}/outputs/genbench"
PY, FPY = f"{M}/.venv/bin/python", "/opt/dlami/nvme/zhiyangdeng/PCE/envs/pce_f3/bin/python"
ENV = dict(os.environ, DIFFSYNTH_SKIP_DOWNLOAD="True", PYTHONPATH=f"{M}/code/gc25d:{M}/code/arope",
           HF_HOME="/opt/dlami/nvme/zhiyangdeng/.cache/huggingface", TMPDIR=f"{M}/.tmp", OMP_NUM_THREADS="8")
L2_CLIPS = {  # 2 条训练 clip（地图严格见过）+ 2 条 val clip（L1 里有原首帧的结果可比）
    "train_v10": "01_Sunlit_Mediterranean_Courtyard__R10_s1501__000672",
    "train_v9": "06_Overgrown_Tropical_Water_Temple__P9C_golden_s6942__003648",
    "val_v10": "01_Sunlit_Mediterranean_Courtyard__R10_s1512__001152",
    "val_v9": "02_Autumn_Woodland_Shrine__P9_golden_s2909__001152",
}
STYLES = {  # 训练数据里没有的画风（v9 的 sketch/pixel/comic/toon/paint/雪/夜 都见过，不用）
    "inkwash": "traditional Chinese ink-wash painting (shuimo) on textured rice paper: black and grey ink strokes, soft ink "
               "bleeding, sparse muted colour washes, empty paper tones",
    "clay": "claymation stop-motion scene: everything sculpted from coloured plasticine and modelling clay, visible "
            "fingerprints, rounded soft shapes, studio lighting",
    "papercraft": "low-poly papercraft diorama: every object folded from coloured cardstock with flat polygon facets, "
                  "visible paper folds and cut edges, soft paper shadows",
}


def run(cmd, log, env=ENV):
    os.makedirs(os.path.dirname(log), exist_ok=True)
    with open(log, "w") as fh:
        r = subprocess.run(cmd, stdout=fh, stderr=subprocess.STDOUT, env=env, cwd=f"{M}/.tmp")
    if r.returncode:
        raise RuntimeError(f"失败（{r.returncode}）：{' '.join(cmd[:4])}… 见 {log}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default=f"{M}/outputs/sft_v10/step-8000.safetensors")
    ap.add_argument("--gpus", default="4,5")
    ap.add_argument("--flux_gpu", default="3")
    a = ap.parse_args()
    gpus = a.gpus.split(",")

    # 1) 导出 L2 clip 的首帧 / ids / prompts
    for tag, cid in L2_CLIPS.items():
        d = f"{G}/L2/{tag}/src"
        if not os.path.exists(f"{d}/done"):
            run([PY, f"{M}/code/gc25d/sample_user.py", "--export_clip", cid, "--out", d],
                f"{G}/logs/export_{tag}.log", dict(ENV, CUDA_VISIBLE_DEVICES=gpus[0]))
            open(f"{d}/done", "w").write(cid)
    print("export ok", flush=True)

    # 2) FLUX：L2 restyle ×12 + L3 map2frame ×4（一个进程，模型只加载一次）
    jobs = []
    for tag in L2_CLIPS:
        src = f"{G}/L2/{tag}/src"
        f0 = [f for f in os.listdir(src) if f.endswith("frame0.png")][0]
        for s, text in STYLES.items():
            jobs.append({"mode": "restyle", "image": f"{src}/{f0}", "style": text, "out": f"{G}/L2/{tag}/{s}_f0.png"})
    cases = sorted(d for d in os.listdir(f"{G}/procmap") if os.path.isdir(f"{G}/procmap/{d}"))
    for c in cases:
        jobs.append({"mode": "map2frame", "map": f"{G}/procmap/{c}/ids.npz", "frame": 0, "out": f"{G}/L3/{c}/f0.png"})
    todo = [j for j in jobs if not os.path.exists(j["out"])]
    if todo:
        for j in todo:
            os.makedirs(os.path.dirname(j["out"]), exist_ok=True)
        json.dump(todo, open(f"{G}/flux_jobs.json", "w"), indent=1)
        run([FPY, f"{M}/code/gc25d/firstframe_flux.py", "batch", "--jobs", f"{G}/flux_jobs.json"], f"{G}/logs/flux.log",
            dict(ENV, CUDA_VISIBLE_DEVICES=a.flux_gpu, PYTHONPATH=f"{M}/code/gc25d"))
    print("flux ok", flush=True)

    # 3) 生成（50 步，同 L1）
    runs = []
    for tag in L2_CLIPS:
        src = f"{G}/L2/{tag}/src"
        ids = [f for f in os.listdir(src) if f.endswith("ids.npz")][0]
        pr = [f for f in os.listdir(src) if f.endswith("prompts.json")][0]
        f0 = [f for f in os.listdir(src) if f.endswith("frame0.png")][0]
        firsts = {s: f"{G}/L2/{tag}/{s}_f0.png" for s in STYLES}
        if tag.startswith("train"):
            firsts["orig"] = f"{src}/{f0}"
        for s, ff in firsts.items():
            runs.append(dict(name=f"{tag}_{s}", first=ff, ids=f"{src}/{ids}", prompts=f"{src}/{pr}", out=f"{G}/L2/{tag}"))
    for c in cases:
        p = f"{G}/procmap/{c}"
        runs.append(dict(name=c, first=f"{G}/L3/{c}/f0.png", ids=f"{p}/ids.npz", prompts=f"{p}/prompts.json", out=f"{G}/L3/{c}"))

    def gen(i_r):
        i, r = i_r
        if os.path.exists(f"{r['out']}/{r['name']}.mp4"):
            return r["name"]
        run([PY, f"{M}/code/gc25d/sample_user.py", "--ckpt", a.ckpt, "--first_frame", r["first"], "--ids", r["ids"],
             "--prompts", r["prompts"], "--out", r["out"], "--name", r["name"], "--thumb"],
            f"{G}/logs/gen_{r['name']}.log", dict(ENV, CUDA_VISIBLE_DEVICES=gpus[i % len(gpus)]))
        return r["name"]

    # 每张卡一个串行队列
    queues = [[(i, r) for i, r in enumerate(runs) if i % len(gpus) == g] for g in range(len(gpus))]
    with ThreadPoolExecutor(len(gpus)) as ex:
        for done in ex.map(lambda q: [gen(x) for x in q], queues):
            print("gen ok", done, flush=True)
    print("ALL DONE", flush=True)


if __name__ == "__main__":
    sys.exit(main())