#!/usr/bin/env python3
"""teaser · 素材合成（maliang venv，CPU）：把 rollout.py 的结果（或 genbench 的旧结果）拼成能直接剪进 teaser 的视频 / 索引。
视频一律经 ffmpeg 管道编码（系统 PATH 里的 ffmpeg，没有就用 imageio_ffmpeg 自带的二进制）：libx264 · yuv420p · BT.709（limited，
标记写进码流）· crf 16 · +faststart。读视频也走 ffmpeg（按码流里的色彩标记转回 RGB，与写入对称）。

case = 下面任何一种（load_case 自动识别）：
  * rollout.py 的输出目录 <out>/<name>/（gen.mp4 + ids.npz + meta.json；meta 里有标签 scene / map / label / title）；
  * genbench3/4 的 roll_bidir|roll_dmd/<tag>/<variant>/（gen.mp4 + meta.json；语义 = <root>/sources.json 里该源链的 cond 拼接，
    bidir 只有 81 帧 = 链首 clip）；
  * sample_user.py 的输出（<name>.json + <name>_gen.mp4；给目录或 json 路径；ids = json 里 inputs.ids.path）。
标签：场景名（右 / 角标）= case 的 label，否则 SCENE_EN[scene]（scenes.py 的 36 个场景 + genbench4 的 10 个），否则把键名美化；
地图名（左）= MAP_EN[map]（procmap 主题），否则美化。命令行 --scene_label / --map_label 覆盖。

子命令
  side   单个 case：语义色图 | 生成 并排（深色底、1 px 细边框、面板上方小字标签：左 = 地图主题，右 = 场景名），--legend 加图例。
  gen    单个 case：只出生成视频（--fit crop 铺满 / pad 留边），左下角标 = 场景名，--inset 右上角语义小窗。
  grid   N 个 case 同步播放的宫格（--cols，--what gen|sem，--labels scene|map|name|none，--margin / --gap）。
  reel   若干 mp4（或 case 目录 ⇒ 它的 gen.mp4）顺序拼接：相邻之间交叉淡化（--mode cross）或经黑场（--mode black），--fade 帧数。
  index  一批结果 → <out>/contact_XX.jpg（分页缩略图：帧 40 + 右上语义小窗 + 名字）+ <out>/thumbs/*.jpg + <out>/index.md（表格 + 路径）。
默认画布 1920×1080、16 fps；side / gen / grid 都取 case 的全部帧（81），--frames N 截短，--fade_io N 首尾各 N 帧淡入淡出（默认 0，交给剪辑）；
语义色图默认描 1 px 类别边界（procmap 的干净多边形更像"地图"），--flat 关掉（真实数据的碎植被描线会显脏）。
side 想要更紧的画面：--canvas 1920x640 --margin 24 --gap 24。

usage（$PY = .venv/bin/python；不用 GPU）：
  $PY scripts/teaser/compose.py side outputs/teaser/roll/<name> -o x.mp4 [--legend] [--map_label Village] [--scene_label "Ice Cave"]
  $PY scripts/teaser/compose.py gen  outputs/teaser/roll/<name> -o x.mp4 [--inset] [--fit crop|pad] [--no_badge]
  $PY scripts/teaser/compose.py grid outputs/teaser/roll/a outputs/teaser/roll/b ... -o g.mp4 --cols 3 [--what gen] [--labels scene]
  $PY scripts/teaser/compose.py reel a.mp4 b.mp4 outputs/teaser/roll/c -o r.mp4 [--fade 8] [--mode cross|black]
  $PY scripts/teaser/compose.py index 'outputs/teaser/roll/*' -o outputs/teaser/roll/_index [--per_page 48] [--cols 8]
glob 参数要加引号（脚本自己展开）；grid / reel / index 的 case 也可以写在文件里：@list.txt（每行一个）。
"""
from __future__ import annotations

import argparse
import glob
import json
import os
import shutil
import subprocess
import sys

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont

M = "/opt/dlami/nvme/zhiyangdeng/maliang"
if f"{M}/code/gc25d" not in sys.path:
    sys.path.insert(0, f"{M}/code/gc25d")
import geometry as GEO  # noqa: E402

DATA_ROOT = f"{M}/data/gc25d_v10"
FPS = 16
CANVAS = (1920, 1080)
BG = (12, 14, 19)
PANEL_BG = (20, 23, 30)
BORDER = (64, 70, 84)
FG = (236, 238, 244)
MUTED = (146, 153, 170)
FONT_LAT = "/usr/share/fonts/google-noto-vf/NotoSans[wght].ttf"
FONT_CJK = "/usr/share/fonts/google-droid-sans-fonts/DroidSansFallbackFull.ttf"     # 只有 CJK，拉丁字母走 Noto

# 语义调色板：静态类 / 玩家 core 10 / Boss core 11 沿用 geometry.PALETTE；肢体部件 12–21 给同色系（genbench4/showcase.py 同款）
PAL = GEO.PALETTE.copy()
PAL[[12, 13, 14, 15, 16]] = [(255, 120, 255), (200, 0, 200), (255, 80, 200), (170, 60, 255), (230, 150, 255)]
PAL[[17, 18, 19, 20, 21]] = [(255, 110, 80), (200, 0, 30), (255, 70, 40), (230, 40, 90), (255, 150, 120)]
LEGEND = [("sky", (0,)), ("terrain", (1,)), ("paving", (2,)), ("architecture", (3,)), ("rock", (4,)), ("vegetation", (5,)),
          ("prop", (6,)), ("player", (10, 12, 13, 14, 15, 16)), ("boss", (11, 17, 18, 19, 20, 21))]

SCENE_EN = {
    "castle": "Medieval Castle", "dungeon": "Dungeon", "lava": "Lava Realm", "icecave": "Ice Cavern", "graveyard": "Haunted Graveyard",
    "swamp": "Swamp", "alien": "Alien Planet", "wildwest": "Wild West Town", "mushroom": "Mushroom Forest", "ruins": "Ruined City",
    "egypt": "Egyptian Temple", "chinese_palace": "Imperial Palace", "jungle_temple": "Jungle Temple", "moonbase": "Moon Base",
    "mars": "Mars Colony", "crystal_cavern": "Crystal Cavern", "candy": "Candy Land", "greek": "Greek Acropolis",
    "sky_islands": "Sky Kingdom", "halloween": "Halloween Patch", "christmas": "Christmas Village", "savanna": "Savanna",
    "beach": "Tropical Beach", "toxic": "Toxic Wasteland", "utopia": "Future City", "elven": "Elven City", "farm": "Countryside Farm",
    "pirate": "Pirate Cove", "samurai": "Samurai Town", "steppe_camp": "Steppe War Camp", "clockwork": "Clockwork World",
    "underworld": "Underworld", "bamboo": "Bamboo Valley", "coral_dry": "Dry Coral Seabed", "autumn_harvest": "Harvest Fair",
    "volcano_forge": "Dwarven Forge", "orig": "Original",
}
MAP_EN = {"village": "Village", "forest": "Forest", "plaza": "Plaza", "canyon": "Canyon", "hedgemaze": "Hedge Maze",
          "stonecircle": "Stone Circle", "pillarhall": "Pillar Hall", "skyislands": "Sky Islands", "cityblock": "City Blocks",
          "mountainpass": "Mountain Pass", "racetrack": "Race Track", "glyph": "Giant Letters", "chess": "Chessboard"}


# ------------------------------------------------------------------ 小工具 ----
def pretty(key: str | None) -> str:
    """键名 → 显示名（下划线 / 连字符 → 空格，首字母大写）。"""
    if not key:
        return ""
    return " ".join(w[:1].upper() + w[1:] for w in str(key).replace("-", "_").split("_") if w)


def even(x: float) -> int:
    return int(round(x / 2)) * 2


def expand(items: list[str]) -> list[str]:
    """命令行的 case / mp4 列表：@文件（每行一个）、glob 自己展开（保持给出的顺序，glob 内按名字排序）。"""
    out = []
    for it in items:
        if it.startswith("@"):
            with open(it[1:], encoding="utf-8") as fh:
                out += expand([ln.strip() for ln in fh if ln.strip() and not ln.startswith("#")])
        elif any(c in it for c in "*?["):
            out += sorted(glob.glob(it))
        else:
            out.append(it)
    return out


# ------------------------------------------------------------------ ffmpeg ----
def ffmpeg_exe() -> str:
    exe = shutil.which("ffmpeg")
    if exe:
        return exe
    import imageio_ffmpeg
    return imageio_ffmpeg.get_ffmpeg_exe()


class VideoWriter:
    """RGB uint8 帧 → ffmpeg 管道 → H.264 yuv420p（BT.709 limited，标记写进码流）。先写 <path>.part.mp4，close 成功才改名
    ⇒ 中途崩了不会留下看似完整的 mp4。W、H 须为偶数（yuv420p）。"""

    def __init__(self, path: str, w: int, h: int, fps: int = FPS, crf: int = 16, preset: str = "slow"):
        assert w % 2 == 0 and h % 2 == 0, (w, h)
        self.path, self.w, self.h, self.n = path, w, h, 0
        os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True)
        self.tmp = path + ".part.mp4"
        cmd = [ffmpeg_exe(), "-y", "-loglevel", "error", "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", f"{w}x{h}", "-r", str(fps),
               "-i", "-", "-vf", "scale=out_color_matrix=bt709:out_range=tv", "-c:v", "libx264", "-preset", preset,
               "-crf", str(crf), "-pix_fmt", "yuv420p", "-colorspace", "bt709", "-color_primaries", "bt709",
               "-color_trc", "bt709", "-color_range", "tv", "-movflags", "+faststart", self.tmp]
        self.p = subprocess.Popen(cmd, stdin=subprocess.PIPE)

    def write(self, fr: np.ndarray) -> None:
        assert fr.dtype == np.uint8 and fr.shape == (self.h, self.w, 3), (fr.dtype, fr.shape, (self.h, self.w))
        self.p.stdin.write(np.ascontiguousarray(fr).tobytes())
        self.n += 1

    def close(self) -> None:
        self.p.stdin.close()
        rc = self.p.wait()
        if rc != 0 or self.n == 0:
            raise RuntimeError(f"ffmpeg 写 {self.path} 失败（rc={rc}，{self.n} 帧）")
        os.replace(self.tmp, self.path)

    def __enter__(self):
        return self

    def __exit__(self, et, ev, tb):
        if et is None:
            self.close()
        else:
            self.p.kill()
            self.p.wait()
            if os.path.exists(self.tmp):
                os.remove(self.tmp)


def write_video(path: str, frames, crf: int = 16, fps: int = FPS, preset: str = "slow") -> int:
    """frames：uint8 [F,H,W,3] 或逐帧迭代器（尺寸取第一帧）。返回帧数。"""
    it = iter(frames)
    first = next(it)
    with VideoWriter(path, first.shape[1], first.shape[0], fps=fps, crf=crf, preset=preset) as vw:
        vw.write(first)
        for fr in it:
            vw.write(fr)
    return vw.n


def probe(path: str) -> tuple[int, int]:
    """mp4 的 (W, H)（PyAV 读头，不解码）。"""
    import av
    with av.open(path) as c:
        s = c.streams.video[0]
        return s.codec_context.width, s.codec_context.height


def iter_video(path: str, n: int | None = None):
    """mp4 → 逐帧 RGB uint8（ffmpeg 解码，按码流的色彩标记转 RGB）。"""
    w, h = probe(path)
    cmd = [ffmpeg_exe(), "-loglevel", "error", "-i", path, "-f", "rawvideo", "-pix_fmt", "rgb24", "-"]
    if n is not None:
        cmd[5:5] = ["-frames:v", str(n)]
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE)
    sz = w * h * 3
    try:
        while True:
            buf = p.stdout.read(sz)
            if len(buf) < sz:
                break
            yield np.frombuffer(buf, np.uint8).reshape(h, w, 3)
    finally:
        p.stdout.close()
        p.wait()


def read_video(path: str, n: int | None = None) -> np.ndarray:
    return np.stack(list(iter_video(path, n)))


# ------------------------------------------------------------------ 语义着色 ----
def colorize(ids: np.ndarray, edges: bool = True) -> np.ndarray:
    """uint8 [H,W]（或 [...,H,W]）→ RGB：PAL 着色，255 = 黑底灰斜纹（同 geometry.colorize_ids），edges ⇒ 类别边界压暗成 1 px 描线。"""
    ids = np.asarray(ids)
    out = PAL[ids]
    unk = ids == GEO.UNKNOWN_ID
    if unk.any():
        h, w = ids.shape[-2:]
        yy, xx = np.mgrid[0:h, 0:w]
        out[unk & np.broadcast_to(((xx + yy) // 4) % 2 == 1, ids.shape)] = GEO.HATCH_GREY
    if edges:
        b = np.zeros(ids.shape, bool)
        b[..., :, :-1] |= ids[..., :, :-1] != ids[..., :, 1:]
        b[..., :-1, :] |= ids[..., :-1, :] != ids[..., 1:, :]
        out[b] = (out[b].astype(np.uint16) * 9 // 16).astype(np.uint8)
    return out


def sem_panel(ids: np.ndarray, w: int, h: int, edges: bool = True) -> np.ndarray:
    """一帧 ids → w×h 的色图：先最近邻缩放 id 再着色（描线在目标分辨率上仍是 1 px）。"""
    if ids.shape != (h, w):
        ids = cv2.resize(ids, (w, h), interpolation=cv2.INTER_NEAREST)
    return colorize(ids, edges)


def fit(img: np.ndarray, w: int, h: int) -> np.ndarray:
    if img.shape[:2] == (h, w):
        return img
    up = w > img.shape[1]
    return cv2.resize(img, (w, h), interpolation=cv2.INTER_LANCZOS4 if up else cv2.INTER_AREA)


# ------------------------------------------------------------------ 文字 ----
_FONTS: dict = {}


def font(size: int, weight: str = "Medium", cjk: bool = False):
    k = (size, weight, cjk)
    if k not in _FONTS:
        if cjk:
            f = ImageFont.truetype(FONT_CJK, size)
        else:
            f = ImageFont.truetype(FONT_LAT, size)
            f.set_variation_by_name(weight)
        _FONTS[k] = f
    return _FONTS[k]


def _runs(s: str):
    out = []
    for ch in s:
        c = ord(ch) >= 0x2E80
        if out and out[-1][0] == c:
            out[-1][1] += ch
        else:
            out.append([c, ch])
    return out


def text_w(s: str, size: int, weight: str = "Medium") -> float:
    return sum(font(size, weight, c).getlength(r) for c, r in _runs(s))


def draw_text(d: ImageDraw.ImageDraw, x: float, y: float, s: str, size: int, fill, weight: str = "Medium") -> float:
    """中英混排（CJK 走 Droid，其余 Noto Sans），(x, y) = 左上，基线 y + 0.8·size。返回右端 x。"""
    for c, r in _runs(s):
        d.text((x, y + 0.8 * size), r, font=font(size, weight, c), fill=fill, anchor="ls")
        x += font(size, weight, c).getlength(r)
    return x


def pill(s: str, size: int, fg=FG, bg=(0, 0, 0), alpha: float = 0.62, weight: str = "SemiBold", pad=None):
    """文字小胶囊 → (rgb uint8 [h,w,3], a float32 [h,w,1])，贴在视频上用 paste_rgba。"""
    px, py = pad or (round(size * 0.55), round(size * 0.28))
    w, h = int(np.ceil(text_w(s, size, weight))) + 2 * px, int(round(size * 1.05)) + 2 * py
    im = Image.new("RGBA", (w, h), (0, 0, 0, 0))
    d = ImageDraw.Draw(im)
    d.rounded_rectangle((0, 0, w - 1, h - 1), radius=h // 2 if h < 3 * size else 8, fill=(*bg, int(255 * alpha)))
    draw_text(d, px, py - round(size * 0.02), s, size, (*fg, 255), weight)
    a = np.asarray(im, np.float32)
    return a[..., :3].astype(np.uint8), a[..., 3:] / 255.0


def paste_rgba(dst: np.ndarray, patch, x: int, y: int) -> None:
    rgb, a = patch
    h, w = rgb.shape[:2]
    h, w = min(h, dst.shape[0] - y), min(w, dst.shape[1] - x)
    if h <= 0 or w <= 0:
        return
    roi = dst[y:y + h, x:x + w].astype(np.float32)
    dst[y:y + h, x:x + w] = (roi * (1 - a[:h, :w]) + rgb[:h, :w] * a[:h, :w]).astype(np.uint8)


def frame_rect(img: np.ndarray, x: int, y: int, w: int, h: int, color=BORDER, t: int = 1) -> None:
    """画在面板外侧的细边框（不压住画面）。"""
    cv2.rectangle(img, (x - t, y - t), (x + w + t - 1, y + h + t - 1), color, t, lineType=cv2.LINE_8)


# ------------------------------------------------------------------ case ----
def _sources_json(gen_dir: str) -> str | None:
    """genbench 的 <root>/roll_<mode>/<tag>/<variant> → <root>/sources.json。"""
    p = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(gen_dir))), "sources.json")
    return p if os.path.isfile(p) else None


def _chain_ids(root_sources: str, tag: str, n: int) -> np.ndarray:
    """genbench 源链的 raw id（= eval_rollout.load_chain 的拼法：cond[clip] + cond_ext[:96−81]），取前 n 帧。"""
    src = {s["tag"]: s for s in json.load(open(root_sources))}[tag]
    sem, total = [], 0
    for cid in src["chain"]:
        with np.load(f"{DATA_ROOT}/cond/{cid}.npz") as z:
            sem.append(z["ids"])
        total += len(sem[-1])
        ext = f"{DATA_ROOT}/cond_ext/{cid}.npz"
        if total < n and os.path.exists(ext):
            with np.load(ext) as z:
                sem.append(z["ids"][:96 - 81])
            total += len(sem[-1])
        if total >= n:
            break
    return np.concatenate(sem, 0)[:n]


def _load_npz_ids(path: str) -> np.ndarray:
    if path.endswith(".npy"):
        return np.load(path)
    with np.load(path) as z:
        return z["ids"] if "ids" in z.files else z[z.files[0]]


def load_case(path: str) -> dict:
    """→ {name, dir, gen（mp4）, ids()（惰性读全部帧 raw id）, scene, map, label, title, seed, sec, frames_npz, sheet, side, meta, kind}。"""
    p = os.path.abspath(path.rstrip("/"))
    if os.path.isdir(p) and os.path.isfile(os.path.join(p, "meta.json")):
        meta = json.load(open(os.path.join(p, "meta.json")))
        gen = os.path.join(p, "gen.mp4")
        if meta.get("kind") == "teaser_rollout":
            lab = meta.get("labels", {})
            ids_path = os.path.join(p, "ids.npz")
            if not os.path.exists(ids_path):                         # 早期结果没存 ids.npz ⇒ 回到输入路径，按 rollout 的规则变换
                ids_path = meta["inputs"]["ids"]["path"]
            res = meta.get("res", "832x480")

            def ids_fn(ids_path=ids_path, res=res):
                a = _load_npz_ids(ids_path)
                w, h = (int(v) for v in res.split("x"))
                if a.shape[1:] == (h, w):
                    return a
                if a.shape[1:] == (GEO.SRC_H, GEO.SRC_W) and res == GEO.RES_TAG:
                    return GEO.resample_ids(a)
                return np.stack([cv2.resize(f, (w, h), interpolation=cv2.INTER_NEAREST) for f in a])
            return {"kind": "teaser", "name": meta["name"], "dir": p, "gen": gen, "ids": ids_fn, "scene": lab.get("scene"),
                    "map": lab.get("map"), "label": lab.get("label"), "title": lab.get("title"), "seed": meta.get("seed"),
                    "sec": meta.get("timing", {}).get("gpu_sec"), "frames_npz": os.path.join(p, "frames.npz"),
                    "sheet": os.path.join(p, "sheet.jpg"), "side": os.path.join(p, "side.mp4"), "meta": meta}
        if "tag" in meta and "variant" in meta:                      # genbench3 / 4 的 roll_<mode>/<tag>/<variant>
            sj = _sources_json(p)
            src = {s["tag"]: s for s in json.load(open(sj))}[meta["tag"]] if sj else {}
            n = int(meta.get("n_frames", 81))
            scene_src = str(src.get("scene", "")).split("_", 1)[-1] if src else ""
            return {"kind": "genbench", "name": f"{meta['tag']}__{meta['variant']}", "dir": p, "gen": gen,
                    "ids": (lambda sj=sj, tag=meta["tag"], n=n: _chain_ids(sj, tag, n)), "scene": meta["variant"],
                    "map": scene_src, "label": None, "title": None, "seed": meta.get("seed"), "sec": meta.get("sec"),
                    "frames_npz": os.path.join(p, "frames.npz"), "sheet": None, "side": None, "meta": meta}
        raise ValueError(f"{p}/meta.json 认不出格式（既不是 teaser rollout 也不是 genbench roll）")
    js = p if p.endswith(".json") else None
    if js is None and os.path.isdir(p):
        cands = [f for f in sorted(glob.glob(os.path.join(p, "*.json")))
                 if os.path.exists(f[:-5] + "_gen.mp4")]
        js = cands[0] if cands else None
    if js is None:
        raise FileNotFoundError(f"{path}：不是 rollout 目录（gen.mp4 + meta.json），也没有 sample_user 的 <name>.json + <name>_gen.mp4")
    meta = json.load(open(js))
    stem = js[:-5]
    name = meta.get("name") or os.path.basename(stem)
    return {"kind": "sample_user", "name": name, "dir": os.path.dirname(js), "gen": stem + "_gen.mp4",
            "ids": (lambda ip=meta["inputs"]["ids"]["path"]: _load_npz_ids(ip)), "scene": None,
            "map": name.split("__")[0] if "__" in name else None, "label": None, "title": None, "seed": meta.get("seed"),
            "sec": meta.get("timing", {}).get("total_sec"), "frames_npz": None, "sheet": stem + "_sheet.jpg",
            "side": stem + ".mp4", "meta": meta}


SCENE_EN.update({  # 量产第二批场景（scenes.py 第二批 35 个）
    "dunes": "Sand Dunes", "oasis": "Desert Oasis", "badlands": "Red-Rock Badlands", "salt_flats": "Salt Flats",
    "arabian": "Arabian Nights", "volcano_rim": "Volcano Crater", "demon_citadel": "Demon Citadel", "ash_wasteland": "Ash Wasteland",
    "phoenix_temple": "Temple of Fire", "enchanted_forest": "Enchanted Forest", "dark_forest": "Cursed Forest",
    "rainforest": "Rainforest", "redwood": "Redwood Forest", "snow_forest": "Snowy Pine Forest", "autumn_maple": "Autumn Maples",
    "cherry_blossom": "Cherry Blossoms", "dino_jungle": "Dinosaur Jungle", "circuit_board": "Circuit Board World",
    "toy_room": "Toy Room", "origami": "Origami World", "synthwave": "Synthwave Grid", "glass_palace": "Glass Palace",
    "honey_hive": "Honey Hive", "cosmic": "Cosmic Nebula", "dragon_lair": "Dragon's Hoard", "floating_library": "Magic Library",
    "ink_landscape": "Ink-Wash Landscape", "venice": "Canal City", "rice_terraces": "Rice Terraces", "tulip_fields": "Tulip Fields",
    "arctic_base": "Arctic Station", "space_station": "Overgrown Space Station", "battlefield": "Ancient Battlefield",
    "tibetan": "Mountain Monastery", "fairground": "Haunted Fairground"})
LOOK_EN = {"photoreal": "Photoreal", "cinematic": "Cinematic", "cartoon": "Cartoon", "pixel": "Pixel Art", "anime": "Anime",
           "painterly": "Hand-Painted", "lowpoly": "Low-Poly", "voxel": "Voxel"}


def scene_label(c: dict, override: str | None = None) -> str:
    """scene 形如 "<场景>@<画风>"（scenes.py 的派生场景）⇒ "Lava Realm · Pixel Art"。"""
    if override is not None:
        return override
    sc = c.get("scene") or ""
    if "@" in sc:
        b, lk = sc.split("@", 1)
        return f"{SCENE_EN.get(b, pretty(b))} · {LOOK_EN.get(lk, pretty(lk))}"
    return c.get("label") or SCENE_EN.get(sc, pretty(sc))


def map_label(c: dict, override: str | None = None) -> str:
    if override is not None:
        return override
    return MAP_EN.get(c.get("map") or "", pretty(c.get("map")))


def case_frames(c: dict, n: int | None) -> tuple[np.ndarray, np.ndarray]:
    """→ (生成 [F,H,W,3], ids [F,h,w])，F = min(视频帧数, ids 帧数, n)。"""
    gen = read_video(c["gen"], n)
    ids = c["ids"]()
    f = min(len(gen), len(ids))
    return gen[:f], ids[:f]


def fade_io(frames, n_total: int, k: int):
    for i, fr in enumerate(frames):
        a = min(1.0, (i + 1) / k, (n_total - i) / k) if k > 0 else 1.0
        yield fr if a >= 1.0 else (fr.astype(np.float32) * a).astype(np.uint8)


# ------------------------------------------------------------------ (a) side / gen ----
def legend_strip(ids_present: set[int], size: int = 20) -> Image.Image:
    items = [(n, ids) for n, ids in LEGEND if ids_present & set(ids)]
    sw = round(size * 0.9)
    widths = [sw + 8 + text_w(n, size) + 22 for n, _ in items]
    im = Image.new("RGBA", (int(sum(widths)) + 2, size + 10), (0, 0, 0, 0))
    d = ImageDraw.Draw(im)
    x = 0
    for (n, ids), w in zip(items, widths):
        d.rounded_rectangle((x, 4, x + sw, 4 + sw), 3, fill=tuple(int(v) for v in PAL[ids[0]]))
        if len(ids) > 1:                                             # 角色：下半格 = 肢体部件色
            d.rectangle((x, 4 + sw // 2, x + sw, 4 + sw), fill=tuple(int(v) for v in PAL[ids[1]]))
        draw_text(d, x + sw + 8, 1, n, size, (*MUTED, 255))
        x += w
    return im


def cmd_side(a) -> None:
    c = load_case(a.case)
    W, H = a.canvas
    gen, ids = case_frames(c, a.frames)
    gap, mx = a.gap, a.margin
    pw = (W - 2 * mx - gap) // 2
    ph = round(pw * gen.shape[1] / gen.shape[2])
    cap = round(24 * W / 1920)                                      # 标签字号（宽 1920 ⇒ 24）
    y0 = (H - ph) // 2 + (cap if a.legend else 0) // 2
    xs = (mx, mx + pw + gap)
    base = Image.new("RGB", (W, H), BG)
    d = ImageDraw.Draw(base)
    for x, (head, sub) in zip(xs, (("SEMANTIC MAP", map_label(c, a.map_label)), ("GENERATED", scene_label(c, a.scene_label)))):
        xe = draw_text(d, x, y0 - cap * 2, head, cap, FG, "SemiBold")
        if sub:
            draw_text(d, xe + cap * 0.6, y0 - cap * 2, "·  " + sub, cap, MUTED, "Medium")
    if a.title:
        draw_text(d, (W - text_w(a.title, cap + 8, "SemiBold")) / 2, y0 - cap * 2 - (cap + 8) * 2.4, a.title, cap + 8, FG, "SemiBold")
    if a.legend:
        lg = legend_strip(set(np.unique(ids[::8]).tolist()), round(cap * 0.82))
        base.paste(lg, (xs[0], y0 + ph + round(cap * 0.8)), lg)
    if a.note:
        draw_text(d, (W - text_w(a.note, round(cap * 0.8))) / 2, H - cap * 2.6, a.note, round(cap * 0.8), MUTED)
    base = np.asarray(base).copy()
    for x in xs:
        frame_rect(base, x, y0, pw, ph)
    nf = len(gen)

    def frames():
        for i in range(nf):
            out = base.copy()
            out[y0:y0 + ph, xs[0]:xs[0] + pw] = sem_panel(ids[i], pw, ph, not a.flat)
            out[y0:y0 + ph, xs[1]:xs[1] + pw] = fit(gen[i], pw, ph)
            yield out
    n = write_video(a.out, fade_io(frames(), nf, a.fade_io), crf=a.crf)
    print(f"side → {a.out}（{W}×{H}，{n} 帧）")


def cmd_gen(a) -> None:
    c = load_case(a.case)
    W, H = a.canvas
    gen = read_video(c["gen"], a.frames)
    ids = c["ids"]()[:len(gen)] if a.inset else None
    h, w = gen.shape[1:3]
    s = max(W / w, H / h) if a.fit == "crop" else min(W / w, H / h)
    sw, sh = even(w * s), even(h * s)
    if a.fit == "crop":
        sw, sh = max(W, sw), max(H, sh)
    ox, oy = (W - sw) // 2, (H - sh) // 2
    lab = scene_label(c, a.scene_label)
    badge = pill(lab, round(H * 0.03)) if lab and not a.no_badge else None
    m = round(H * 0.035)
    iw = even(W * a.inset_frac)
    ih = even(iw * h / w)
    ins_cap = pill("semantic input", round(H * 0.017), alpha=0.55, weight="Medium") if a.inset else None

    def frames():
        for i in range(len(gen)):
            big = cv2.resize(gen[i], (sw, sh), interpolation=cv2.INTER_LANCZOS4)
            if a.fit == "crop":
                out = big[-oy:-oy + H, -ox:-ox + W].copy()
            else:
                out = np.empty((H, W, 3), np.uint8)
                out[:] = BG
                out[oy:oy + sh, ox:ox + sw] = big
            if a.inset:
                x, y = W - m - iw, m
                out[y:y + ih, x:x + iw] = sem_panel(ids[i], iw, ih, not a.flat)
                frame_rect(out, x, y, iw, ih, (235, 235, 240), 2)
                paste_rgba(out, ins_cap, x + 6, y + ih - ins_cap[0].shape[0] - 6)
            if badge is not None:
                paste_rgba(out, badge, m, H - m - badge[0].shape[0])
            yield out
    n = write_video(a.out, fade_io(frames(), len(gen), a.fade_io), crf=a.crf)
    print(f"gen → {a.out}（{W}×{H}，{n} 帧，fit={a.fit}{'，语义小窗' if a.inset else ''}）")


# ------------------------------------------------------------------ (b) grid ----
def grid_layout(n: int, cols: int, W: int, H: int, aspect: float, margin: int, gap: int):
    rows = (n + cols - 1) // cols
    tw = (W - 2 * margin - (cols - 1) * gap) / cols
    th = tw / aspect
    if rows * th + (rows - 1) * gap + 2 * margin > H:
        th = (H - 2 * margin - (rows - 1) * gap) / rows
        tw = th * aspect
    tw, th = int(tw), int(th)
    gw, gh = cols * tw + (cols - 1) * gap, rows * th + (rows - 1) * gap
    x0, y0 = (W - gw) // 2, (H - gh) // 2
    return rows, tw, th, [(x0 + (k % cols) * (tw + gap), y0 + (k // cols) * (th + gap)) for k in range(n)]


def cmd_grid(a) -> None:
    cases = [load_case(p) for p in expand(a.cases)]
    if not cases:
        raise SystemExit("grid：没有 case")
    W, H = a.canvas
    cols = a.cols or int(np.ceil(np.sqrt(len(cases) * 16 / 9 * 480 / 832)))
    gens = [read_video(c["gen"], a.frames) if a.what == "gen" else None for c in cases]
    idss = [c["ids"]() if a.what == "sem" else None for c in cases]
    lens = [len(g) if g is not None else len(i) for g, i in zip(gens, idss)]
    nf = min(lens)
    h0, w0 = (gens[0].shape[1:3] if gens[0] is not None else idss[0].shape[1:3])
    rows, tw, th, pos = grid_layout(len(cases), cols, W, H, w0 / h0, a.margin, a.gap)
    fs = max(13, round(th * 0.075))
    labels = []
    for c in cases:
        s = {"scene": scene_label(c), "map": map_label(c), "name": c["name"], "none": ""}[a.labels]
        labels.append(pill(s, fs, alpha=0.6) if s else None)

    def frames():
        for i in range(nf):
            out = np.empty((H, W, 3), np.uint8)
            out[:] = BG
            for k, (x, y) in enumerate(pos):
                tile = fit(gens[k][i], tw, th) if a.what == "gen" else sem_panel(idss[k][i], tw, th, not a.flat)
                out[y:y + th, x:x + tw] = tile
                if labels[k] is not None:
                    paste_rgba(out, labels[k], x + round(fs * 0.5), y + th - labels[k][0].shape[0] - round(fs * 0.5))
            yield out
    n = write_video(a.out, fade_io(frames(), nf, a.fade_io), crf=a.crf)
    print(f"grid → {a.out}（{len(cases)} 个 case，{rows}×{cols}，格 {tw}×{th}，{W}×{H}，{n} 帧）")


# ------------------------------------------------------------------ (c) reel ----
def n_frames(path: str) -> int:
    import av
    with av.open(path) as c:
        n = c.streams.video[0].frames
    return n or sum(1 for _ in iter_video(path))


def cmd_reel(a) -> None:
    srcs = [p if p.endswith(".mp4") else load_case(p)["gen"] for p in expand(a.clips)]
    W, H = a.canvas if a.canvas is not None else probe(srcs[0])
    k = a.fade

    def norm(fr):
        if fr.shape[:2] == (H, W):
            return fr
        h, w = fr.shape[:2]
        s = min(W / w, H / h)
        sw, sh = even(w * s), even(h * s)
        out = np.empty((H, W, 3), np.uint8)
        out[:] = 0
        out[(H - sh) // 2:(H - sh) // 2 + sh, (W - sw) // 2:(W - sw) // 2 + sw] = cv2.resize(
            fr, (sw, sh), interpolation=cv2.INTER_LANCZOS4 if s > 1 else cv2.INTER_AREA)
        return out

    def scale(fr, t):
        return (fr.astype(np.float32) * t).astype(np.uint8)

    def frames():
        tail = None
        for ci, src in enumerate(srcs):
            clip = [norm(f) for f in iter_video(src)]
            first, last = ci == 0, ci == len(srcs) - 1
            if len(clip) <= (0 if first else k) + (0 if last else k):
                raise ValueError(f"{src} 只有 {len(clip)} 帧，不够 --fade {k}")
            if a.mode == "cross":
                head = 0
                if tail is not None:
                    for j in range(k):
                        t = (j + 1) / (k + 1)
                        yield (tail[j].astype(np.float32) * (1 - t) + clip[j].astype(np.float32) * t).astype(np.uint8)
                    head = k
                end = len(clip) if last else len(clip) - k
                yield from clip[head:end]
                tail = clip[end:]
            else:                                                     # 经黑场：每段（首段除外）淡入、（末段除外）淡出
                n = len(clip)
                for j, fr in enumerate(clip):
                    t = 1.0
                    if not first and j < k:
                        t = min(t, (j + 1) / (k + 1))
                    if not last and j >= n - k:
                        t = min(t, (n - j) / (k + 1))
                    yield fr if t >= 1.0 else scale(fr, t)

    total = sum(n_frames(s) for s in srcs) - (k * (len(srcs) - 1) if a.mode == "cross" else 0)
    n = write_video(a.out, fade_io(frames(), total, a.fade_io), crf=a.crf)
    print(f"reel → {a.out}（{len(srcs)} 段，{W}×{H}，{n} 帧 = {n / FPS:.1f} s，{a.mode} 过渡 {k} 帧）")


# ------------------------------------------------------------------ (d) index ----
def thumb(c: dict, tw: int = 416, th: int = 240, frame: int = 40) -> np.ndarray:
    """帧 frame 的生成画面 + 右上角 1/3 宽语义小窗（frames.npz 有就不解码视频）。"""
    rgb = ids = None
    fz = c.get("frames_npz")
    if fz and os.path.exists(fz):
        with np.load(fz) as z:
            idx = list(z["idx"])
            if frame in idx:
                rgb, ids = z["rgb"][idx.index(frame)], z["ids"][idx.index(frame)]
    if rgb is None:
        g = read_video(c["gen"], frame + 1)
        frame = min(frame, len(g) - 1)
        rgb, ids = g[frame], c["ids"]()[frame]
    out = fit(rgb, tw, th).copy()
    iw, ih = tw // 3, th // 3
    out[3:3 + ih, tw - iw - 3:tw - 3] = sem_panel(ids, iw, ih)
    cv2.rectangle(out, (tw - iw - 4, 2), (tw - 3, 3 + ih), (235, 235, 240), 1)
    return out


def cmd_index(a) -> None:
    paths = expand(a.roots)
    cases = []
    for p in paths:
        try:
            cases.append(load_case(p))
        except (FileNotFoundError, ValueError, KeyError):
            continue
    if not cases:
        raise SystemExit(f"index：{a.roots} 下没有认得出的结果")
    cases.sort(key=lambda c: c["name"])
    os.makedirs(os.path.join(a.out, "thumbs"), exist_ok=True)
    tw, th, lab_h, cols = a.thumb_w, round(a.thumb_w * 480 / 832), 26, a.cols
    tiles = []
    for c in cases:
        t = thumb(c, tw, th)
        tp = os.path.join(a.out, "thumbs", c["name"] + ".jpg")
        cv2.imwrite(tp, cv2.cvtColor(t, cv2.COLOR_RGB2BGR), [cv2.IMWRITE_JPEG_QUALITY, 88])
        c["thumb"] = tp
        strip = Image.new("RGB", (tw, lab_h), BG)
        d = ImageDraw.Draw(strip)
        sub = " · ".join(s for s in (scene_label(c), map_label(c)) if s)
        x = draw_text(d, 6, 4, c["name"][:46], 14, FG, "SemiBold")
        if sub and x + 12 + text_w(sub, 13) < tw:
            draw_text(d, tw - 6 - text_w(sub, 13), 5, sub, 13, MUTED)
        tiles.append(np.concatenate([t, np.asarray(strip)], 0))
    pages = []
    for pi in range(0, len(tiles), a.per_page):
        chunk = tiles[pi:pi + a.per_page]
        rows = (len(chunk) + cols - 1) // cols
        sheet = np.empty((rows * (th + lab_h + 6) + 6, cols * (tw + 6) + 6, 3), np.uint8)
        sheet[:] = BG
        for k, t in enumerate(chunk):
            x, y = 6 + (k % cols) * (tw + 6), 6 + (k // cols) * (th + lab_h + 6)
            sheet[y:y + th + lab_h, x:x + tw] = t
        pp = os.path.join(a.out, f"contact_{len(pages):02d}.jpg")
        cv2.imwrite(pp, cv2.cvtColor(sheet, cv2.COLOR_RGB2BGR), [cv2.IMWRITE_JPEG_QUALITY, 88])
        pages.append(pp)
    rel = lambda p: os.path.relpath(p, a.out) if p else ""      # noqa: E731
    lines = [f"# {a.title or 'teaser 结果索引'}", "",
             f"{len(cases)} 个结果 · 缩略图 = 帧 40（2.5 s）生成画面，右上角小窗 = 同帧语义条件。contact sheet："
             + "、".join(f"[{os.path.basename(p)}]({rel(p)})" for p in pages), "",
             "| # | 缩略图 | name | 场景 | 地图 | seed | GPU s | 文件 |", "|---|---|---|---|---|---|---|---|"]
    for k, c in enumerate(cases):
        files = [f"[gen]({rel(c['gen'])})"]
        for key in ("side", "sheet"):
            if c.get(key) and os.path.exists(c[key]):
                files.append(f"[{key}]({rel(c[key])})")
        sec = c.get("sec")
        lines.append(f"| {k} | ![]({rel(c['thumb'])}) | `{c['name']}` | {scene_label(c)} | {map_label(c)} | {c.get('seed', '')} | "
                     f"{'' if sec is None else f'{sec:.0f}'} | {' · '.join(files)} |")
    lines += ["", "绝对路径：", ""] + [f"- `{c['name']}`: `{c['dir']}`" for c in cases]
    with open(os.path.join(a.out, "index.md"), "w", encoding="utf-8") as fh:
        fh.write("\n".join(lines) + "\n")
    print(f"index → {a.out}/index.md（{len(cases)} 个，{len(pages)} 页 contact sheet）")


# ------------------------------------------------------------------ CLI ----
def _wh(s: str) -> tuple[int, int]:
    w, h = (int(v) for v in s.lower().split("x"))
    assert w % 2 == 0 and h % 2 == 0, "画布宽高须为偶数"
    return w, h


def build_parser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(description="teaser 素材合成（side / gen / grid / reel / index）")
    sub = ap.add_subparsers(dest="cmd", required=True)

    def common(p, canvas=True):
        p.add_argument("-o", "--out", required=True)
        p.add_argument("--crf", type=int, default=16)
        p.add_argument("--fade_io", type=int, default=0, help="首尾各 N 帧从黑淡入 / 淡出到黑")
        p.add_argument("--frames", type=int, default=None, help="只取前 N 帧（默认全部）")
        if canvas:
            p.add_argument("--canvas", type=_wh, default=CANVAS, help="画布 WxH，默认 1920x1080")
            p.add_argument("--flat", action="store_true", help="语义色图不描类别边界（真实数据的碎植被描线会显脏）")

    p = sub.add_parser("side", help="单个 case：语义 | 生成 并排")
    p.add_argument("case")
    common(p)
    p.add_argument("--legend", action="store_true")
    p.add_argument("--map_label", default=None)
    p.add_argument("--scene_label", default=None)
    p.add_argument("--title", default=None, help="面板上方居中的一行标题（默认无）")
    p.add_argument("--note", default=None, help="底部居中的一行小字（默认无）")
    p.add_argument("--margin", type=int, default=48)
    p.add_argument("--gap", type=int, default=32)
    p = sub.add_parser("gen", help="单个 case：只出生成视频 + 角标")
    p.add_argument("case")
    common(p)
    p.add_argument("--fit", choices=("crop", "pad"), default="crop")
    p.add_argument("--scene_label", default=None)
    p.add_argument("--no_badge", action="store_true")
    p.add_argument("--inset", action="store_true", help="右上角语义小窗")
    p.add_argument("--inset_frac", type=float, default=0.24)
    p = sub.add_parser("grid", help="N×M 宫格同步播放")
    p.add_argument("cases", nargs="+")
    common(p)
    p.add_argument("--cols", type=int, default=0, help="列数（默认按画布比例自动）")
    p.add_argument("--what", choices=("gen", "sem"), default="gen")
    p.add_argument("--labels", choices=("scene", "map", "name", "none"), default="scene")
    p.add_argument("--margin", type=int, default=0)
    p.add_argument("--gap", type=int, default=6)
    p = sub.add_parser("reel", help="顺序拼接 + 淡入淡出")
    p.add_argument("clips", nargs="+")
    common(p, canvas=False)
    p.add_argument("--canvas", type=_wh, default=None, help="默认 = 第一段的尺寸；其它尺寸的段等比缩放居中")
    p.add_argument("--fade", type=int, default=8, help="相邻两段之间的过渡帧数")
    p.add_argument("--mode", choices=("cross", "black"), default="cross")
    p = sub.add_parser("index", help="contact sheet + index.md")
    p.add_argument("roots", nargs="+", help="case 目录 / glob（加引号）/ @列表")
    p.add_argument("-o", "--out", required=True, help="输出目录")
    p.add_argument("--per_page", type=int, default=48)
    p.add_argument("--cols", type=int, default=8)
    p.add_argument("--thumb_w", type=int, default=416)
    p.add_argument("--title", default=None)
    return ap


def main(argv=None) -> None:
    a = build_parser().parse_args(argv)
    {"side": cmd_side, "gen": cmd_gen, "grid": cmd_grid, "reel": cmd_reel, "index": cmd_index}[a.cmd](a)


if __name__ == "__main__":
    main()