#!/usr/bin/env python3
"""teaser · 世界编辑选片 + 出片（maliang venv，CPU）：读 <e81>/cases/*/（推完、有 edit_score.json 的），按编辑类型分组排序，出
<e81>/showcase/：
  picks.jsonl / picks.md          选中的 case（每类前 N 条，同一底子在同一类里只取一条）+ 指标
  side/<rank>_<edit>_<name>.mp4   「原版 | 编辑版」并排（compose_edit side，1920×1080）
  grid_<family>.mp4               同一类编辑的宫格（每格「原版 | 编辑版」，3 行 × 2 列）
  reel.mp4                        精选并排串接（每类轮流取，交叉淡化）
  all_index.md / all.jsonl        全部 case 的指标表（含没选中的、生效不明显的）

排序分 = effect（编辑区域里编辑版 vs 原版的 ΔE − 区域外 ΔE）的类内名次 + 0.3·fit（区域里外观判成编辑后类的比例差）
        − 硬伤惩罚（vscore 有 flags 时；没跑 vscore 就不罚）。effect < --min_effect 的不选（= 模型没画出来）。
家族（family）= edit + 主变体：add / add-drop / grow / delete / delete-class / swap / spawn-ring / spawn-avenue / trail3d /
paint（地面）/ write-paint / write-raise / trap / boss_del / boss_clone。

usage：.venv/bin/python scripts/teaser/edit_showcase.py [--root outputs/teaser/edit/e81] [--per 8] [--reel_per 2] [--procs 16]
       [--no_video]（只出 picks / 索引）
       .venv/bin/python scripts/teaser/edit_showcase.py --hero hero.txt   # 手挑名单 → showcase/hero/ + hero_reel.mp4
"""
from __future__ import annotations

import argparse
import collections
import json
import os
import subprocess
import sys
from concurrent.futures import ProcessPoolExecutor

M = "/opt/dlami/nvme/zhiyangdeng/maliang"
HERE = os.path.dirname(os.path.abspath(__file__))
PY = f"{M}/.venv/bin/python"
HARD = {"scene_lost", "degraded", "broken", "char_melt", "blur"}
TITLE = {"add": "Add an object", "add-drop": "Drop an object from the sky", "grow": "Grow / raise an object",
         "delete": "Delete an object", "delete-class": "Delete a whole class", "swap-object": "Swap materials",
         "swap-ground": "Swap the ground material",
         "spawn-ring": "Spawn a ring of objects", "spawn-avenue": "Objects sprout along the way",
         "trail3d": "Things grow in the hero's footsteps", "paint": "Paint the ground", "write-paint": "Write on the ground",
         "write-raise": "Raise letters / symbols", "trap": "Trap the boss", "boss_del": "Remove the boss",
         "boss_clone": "Clone the boss"}


def family(e: dict) -> str:
    ed, p = e["edit"], e.get("params") or {}
    if ed == "add":
        return "add-drop" if p.get("anim") == "drop" else "add"
    if ed == "delete":
        return "delete-class" if p.get("mode") == "class" else "delete"
    if ed == "spawn":
        return "spawn-avenue" if p.get("mode") == "avenue" else "spawn-ring"
    if ed == "paint":
        return "trail3d" if p.get("mode") == "trail3d" else "paint"
    if ed == "write":
        return "write-raise" if p.get("mode") == "raise" else "write-paint"
    if ed == "swap":
        return "swap-ground" if (e.get("resolved") or {}).get("on") == "ground" else "swap-object"
    return ed


def load_rows(root: str) -> list[dict]:
    rows = []
    cd = os.path.join(root, "cases")
    for n in sorted(os.listdir(cd)):
        d = os.path.join(cd, n)
        sp, ep = os.path.join(d, "edit_score.json"), os.path.join(d, "edit.json")
        if not (os.path.isfile(sp) and os.path.isfile(ep) and os.path.isfile(os.path.join(d, "gen.mp4"))):
            continue
        s, e = json.load(open(sp)), json.load(open(ep))
        vs = f"{M}/outputs/teaser/vscore/edit/{n}.vscore.json"
        v = json.load(open(vs)) if os.path.exists(vs) else {}
        rows.append(dict(name=n, dir=d, family=family(e), edit=e["edit"], label=e["label"], base=e["base"]["name"],
                         theme=e["base"]["theme"], scene=e["base"]["scene"], edit_frame=e["edit_frame"],
                         effect=s.get("effect"), ratio=s.get("ratio"), fit=s.get("fit"), pre_delta=s.get("pre_delta"),
                         teaser=v.get("teaser"), flags=v.get("flags") or [], weak=e.get("weak")))
    return rows


def rank(rows, min_effect):
    by = collections.defaultdict(list)
    for r in rows:
        by[r["family"]].append(r)
    for fam, rs in by.items():
        eff = sorted([r["effect"] or -99 for r in rs])
        for r in rs:
            q = (eff.index(r["effect"] or -99) + 1) / len(eff)
            hard = set(r["flags"]) & HARD
            if fam == "boss_clone":                 # 两个 Boss 时 vscore 的角色放大框会框到两个 ⇒ char_melt 误报，不罚
                hard -= {"char_melt"}
            if fam.startswith("write"):             # 地上的字会被 VLM 当成「花屏 / 文字」⇒ broken / degraded 误报，只认 scene_lost
                hard &= {"scene_lost"}
            pen = 0.25 * len(hard)
            r["rank_score"] = round(q + 0.3 * (r["fit"] or 0) - pen + (0.1 * (r["teaser"] or 0)), 4)
            r["effective"] = (r["effect"] or -99) >= min_effect
        rs.sort(key=lambda r: -r["rank_score"])
    return by


def pick(by, per):
    out = {}
    for fam, rs in by.items():
        seen, sel = set(), []
        for r in rs:
            if not r["effective"] or r["base"] in seen:
                continue
            sel.append(r)
            seen.add(r["base"])
            if len(sel) >= per:
                break
        out[fam] = sel
    return out


def run(cmd):
    p = subprocess.run(cmd, capture_output=True, text=True)
    if p.returncode:
        print("FAIL", " ".join(cmd[:4]), p.stderr[-500:], flush=True)
    return p.returncode


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", default=f"{M}/outputs/teaser/edit/e81")
    ap.add_argument("--per", type=int, default=8)
    ap.add_argument("--reel_per", type=int, default=2)
    ap.add_argument("--min_effect", type=float, default=8.0)
    ap.add_argument("--procs", type=int, default=16)
    ap.add_argument("--no_video", action="store_true")
    ap.add_argument("--hero", default=None, help="手挑的 case 名单（每行一个 case 名或目录）⇒ showcase/hero/NN_<name>.mp4 + hero_reel.mp4")
    a = ap.parse_args(argv)
    if a.hero:
        return hero(a)
    sd = os.path.join(a.root, "showcase")
    os.makedirs(os.path.join(sd, "side"), exist_ok=True)
    rows = load_rows(a.root)
    by = rank(rows, a.min_effect)
    picks = pick(by, a.per)
    with open(os.path.join(sd, "all.jsonl"), "w") as fh:
        for fam in sorted(by):
            for r in by[fam]:
                fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    with open(os.path.join(sd, "all_index.md"), "w") as fh:
        fh.write(f"# 世界编辑：全部 {len(rows)} 条（按类、类内按排序分）\n\n")
        fh.write("| 类 | 条数 | 生效（effect ≥ %.0f） | effect 中位 |\n|---|---|---|---|\n" % a.min_effect)
        for fam in sorted(by):
            rs = by[fam]
            eff = sorted(r["effect"] or -99 for r in rs)
            fh.write(f"| {fam} | {len(rs)} | {sum(r['effective'] for r in rs)} | {eff[len(eff) // 2]:.1f} |\n")
        for fam in sorted(by):
            fh.write(f"\n## {fam}\n\n| 排序分 | effect | fit | pre | teaser | 标签 | case |\n|---|---|---|---|---|---|---|\n")
            for r in by[fam]:
                fh.write(f"| {r['rank_score']} | {r['effect']} | {r['fit']} | {r['pre_delta']} | {r['teaser']} | {r['label']} | "
                         f"`{r['name']}` |\n")
    sel = [(fam, i, r) for fam in sorted(picks) for i, r in enumerate(picks[fam])]
    with open(os.path.join(sd, "picks.jsonl"), "w") as fh:
        for fam, i, r in sel:
            r["side"] = os.path.join(sd, "side", f"{fam}_{i:02d}__{r['name']}.mp4")
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    with open(os.path.join(sd, "picks.md"), "w") as fh:
        fh.write(f"# 世界编辑精选（{len(sel)} 条，每类 ≤ {a.per}）\n\n")
        for fam in sorted(picks):
            fh.write(f"\n## {fam} — {TITLE.get(fam, fam)}（{len(picks[fam])} / {len(by[fam])}）\n\n")
            for i, r in enumerate(picks[fam]):
                fh.write(f"- {r['label']} · effect {r['effect']} · fit {r['fit']} · `{os.path.basename(r['side'])}`\n")
    print(f"{len(rows)} 条 → 精选 {len(sel)}；每类：" + ", ".join(f"{f} {len(picks[f])}/{len(by[f])}" for f in sorted(by)), flush=True)
    if a.no_video:
        return
    ce = os.path.join(HERE, "compose_edit.py")
    jobs = [[PY, ce, "side", r["dir"], "-o", r["side"]] for _, _, r in sel if not os.path.exists(r["side"])]
    for fam in sorted(picks):
        if len(picks[fam]) >= 2:
            cs = [r["dir"] for r in picks[fam][:6]]
            jobs.append([PY, ce, "grid", *cs, "-o", os.path.join(sd, f"grid_{fam}.mp4"), "--cols", "2", "--pairs",
                         "--title", f"{TITLE.get(fam, fam)} — original | edited"])
    with ProcessPoolExecutor(a.procs) as ex:
        list(ex.map(run, jobs))
    reel = []
    for i in range(a.reel_per):
        for fam in ("add", "grow", "delete", "delete-class", "swap-object", "spawn-avenue", "paint", "add-drop", "boss_del",
                    "boss_clone", "trap", "write-raise", "spawn-ring", "trail3d", "write-paint", "swap-ground"):
            if fam in picks and i < len(picks[fam]):
                reel.append(picks[fam][i]["side"])
    if reel:
        run([PY, os.path.join(HERE, "compose.py"), "reel", *reel, "-o", os.path.join(sd, "reel.mp4"), "--fade", "8"])
    print(f"出片 → {sd}", flush=True)


def hero(a):
    """手挑名单 → 逐条并排（showcase/hero/NN_<name>.mp4）+ 串接 hero_reel.mp4（交叉淡化 8 帧）。"""
    hd = os.path.join(a.root, "showcase", "hero")
    os.makedirs(hd, exist_ok=True)
    names = [ln.strip() for ln in open(a.hero) if ln.strip() and not ln.startswith("#")]
    ce = os.path.join(HERE, "compose_edit.py")
    jobs, outs = [], []
    for i, n in enumerate(names):
        d = n if os.path.isdir(n) else os.path.join(a.root, "cases", n)
        o = os.path.join(hd, f"{i:02d}_{os.path.basename(d.rstrip('/'))}.mp4")
        outs.append(o)
        if not os.path.exists(o):
            jobs.append([PY, ce, "side", d, "-o", o])
    with ProcessPoolExecutor(a.procs) as ex:
        list(ex.map(run, jobs))
    run([PY, os.path.join(HERE, "compose.py"), "reel", *outs, "-o", os.path.join(a.root, "showcase", "hero_reel.mp4"), "--fade", "8"])
    print(f"hero：{len(outs)} 条 → {hd}，hero_reel.mp4", flush=True)


if __name__ == "__main__":
    main()