#!/usr/bin/env python3
"""genbench 结果拼成每个源一个对比视频（2 列网格，每格 832×480 带标签）。
usage: genbench_compose.py <tag> ...   tag ∈ train_v10 train_v9 val_v10 val_v9 L3"""
import os
import sys

import cv2
import imageio.v3 as iio
import numpy as np

G = "/opt/dlami/nvme/zhiyangdeng/maliang/outputs/genbench"
L1 = "/opt/dlami/nvme/zhiyangdeng/maliang/outputs/sft_v10/eval_step-8000/L1_cond"
W, H = 832, 480
STY = {"orig": "原首帧", "inkwash": "水墨（没见过）", "clay": "黏土（没见过）", "papercraft": "纸艺低多边形（没见过）"}


def panel(path, k):
    """k-th 832-wide panel of an mp4 (title bar 48 px cropped) → [81,480,832,3]."""
    v = iio.imread(path)
    return v[:, -H:, k * W:(k + 1) * W]


def label(fr, text):
    out = fr.copy()
    cv2.rectangle(out, (0, 0), (W, 34), (0, 0, 0), -1)
    # cv2 没有中文字体：标签用英文
    cv2.putText(out, text, (10, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255, 230, 120), 2, cv2.LINE_AA)
    return out


def grid(tiles, ncol=2):
    T = min(len(v) for _, v in tiles)
    while len(tiles) % ncol:
        tiles.append(("", np.zeros_like(tiles[0][1])))
    frames = []
    for t in range(T):
        rows = [np.concatenate([label(v[t], n) for n, v in tiles[r:r + ncol]], 1) for r in range(0, len(tiles), ncol)]
        frames.append(np.concatenate(rows, 0))
    return np.stack(frames)


def write(path, frames):
    iio.imwrite(path, frames, fps=16, codec="libx264", quality=None, macro_block_size=16,
                ffmpeg_params=["-crf", "23", "-pix_fmt", "yuv420p", "-movflags", "+faststart"])
    print(path, frames.shape, round(os.path.getsize(path) / 1e6, 1), "MB")


EN = {"orig": "original first frame", "inkwash": "UNSEEN style: ink-wash", "clay": "UNSEEN style: claymation",
      "papercraft": "UNSEEN style: papercraft"}


def l2(tag):
    d = f"{G}/L2/{tag}"
    tiles = []
    any_mp4 = next(f for f in sorted(os.listdir(d)) if f.endswith(".mp4") and not f.endswith("_gen.mp4"))
    tiles.append(("input map (numeric ids, parts coloured)", panel(f"{d}/{any_mp4}", 1)))
    if tag.startswith("val"):
        cid = open(f"{d}/src/done").read().strip()
        tiles.append(("GT (held-out run)", panel(f"{L1}/{cid}.mp4", 0)))
        tiles.append(("generated: original first frame", panel(f"{L1}/{cid}.mp4", 1)))
    for s in ("orig", "inkwash", "clay", "papercraft"):
        p = f"{d}/{tag}_{s}.mp4"
        if os.path.exists(p):
            tiles.append((f"generated: {EN[s]}", panel(p, 0)))
    write(f"{G}/compare_{tag}.mp4", grid(tiles))


def l3():
    tiles = []
    for c in sorted(os.listdir(f"{G}/L3")):
        p = f"{G}/L3/{c}/{c}.mp4"
        if not os.path.exists(p):
            continue
        theme = c.split("__")[0]
        tiles.append((f"UNSEEN map [{theme}]: numeric ids", panel(p, 1)))
        tiles.append((f"generated from FLUX first frame [{theme}]", panel(p, 0)))
    write(f"{G}/compare_L3_unseen_maps.mp4", grid(tiles))


if __name__ == "__main__":
    for tag in sys.argv[1:]:
        l3() if tag == "L3" else l2(tag)