#!/usr/bin/env python3
"""teaser · 编辑 case（外观编辑 + 动作 / 轨迹控制）的合成（maliang venv，CPU）。被 edit_app.py compose / edit_sameworld.py compose
调用，也可单独跑。画法沿用 compose.py（深色底、细边框、胶囊标签、ffmpeg x264 crf 16 BT.709、16 fps、1920×1080）。

产物（outputs/teaser/edit/app/showcase/，下称 $O）
  look_<b>.mp4          同一张语义地图 × 多种首帧外观：左上 = 语义图，其余 = 原首帧 + 各外观编辑（同 ids、同按键）
  look_all_<v>.mp4      同一种编辑（如 night）在所有底子上（每格右上角是原版小窗）
  look_pairs.mp4        每个底子挑 2 种编辑，「原版 | 编辑」并排 4×… 宫格
  mirror_<b>.mp4        「原版 | 镜像（A↔D）」并排 + 按键 HUD
  text_<b>.mp4          原版 + 只改文本的 3 种（全 standing still / 反方向 / 一直攻击）2×2 + 文本按键 HUD（语义不变）
  sameworld_<g>.mp4     左 = 俯视总图（各 clip 轨迹一色 + 当前位置圆点），右 = 各 clip 的生成视频 + 按键 HUD
  sheets/*.jpg          对应的帧拼图（帧 0/20/40/60/80），给人看 / 报告用
按键 HUD（keys_hud）：W / A / S / D 方向键 + ATK，按下 = 亮。帧 f 的按键 = blocks20[max(0, ceil(f/4) − 1)]（cell j = 帧 4j−3..4j
← block j−1，cell 0 = cell 1）。

usage：$PY scripts/teaser/edit_compose.py app [--only look act]  |  sameworld <group>...
"""
from __future__ import annotations

import argparse
import json
import math
import os
import sys

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

M = "/opt/dlami/nvme/zhiyangdeng/maliang"
HERE = os.path.dirname(os.path.abspath(__file__))
for _p in (f"{M}/code/gc25d", HERE):
    if _p not in sys.path:
        sys.path.insert(0, _p)
import compose as CP  # noqa: E402

A = f"{M}/outputs/teaser/edit/app"
O = f"{A}/showcase"
W, H = CP.CANVAS


# ------------------------------------------------------------------ 按键 HUD ----
def block_of(f: int) -> int:
    return max(0, math.ceil(f / 4) - 1)


_HUD: dict = {}


def keys_hud(state, size: int = 26):
    """state = (w, a, s, d, attack) 0/1 → (rgb, alpha)：上排 W，下排 A S D，右边 ATK。按下 = 白底黑字，否则半透明黑底灰字。"""
    key = (tuple(int(x) for x in state), size)
    if key in _HUD:
        return _HUD[key]
    g = round(size * 0.18)
    aw = round(size * 2.1)
    wd, ht = 3 * size + 2 * g + g + aw, 2 * size + g
    im = Image.new("RGBA", (wd + 2, ht + 2), (0, 0, 0, 0))
    d = ImageDraw.Draw(im)
    pos = {"W": (size + g, 0), "A": (0, size + g), "S": (size + g, size + g), "D": (2 * (size + g), size + g)}
    on = dict(zip("WASD", key[0][:4]))
    fs = round(size * 0.55)
    for k, (x, y) in pos.items():
        lit = on[k]
        d.rounded_rectangle((x, y, x + size, y + size), radius=round(size * 0.22),
                            fill=(245, 245, 250, 240) if lit else (0, 0, 0, 120),
                            outline=(255, 255, 255, 200) if lit else (200, 200, 210, 110), width=1)
        tw = CP.text_w(k, fs, "SemiBold")
        CP.draw_text(d, x + (size - tw) / 2, y + (size - fs) / 2 - 1, k, fs, (20, 20, 25, 255) if lit else (210, 212, 220, 220),
                     "SemiBold")
    x0 = 3 * size + 3 * g
    lit = key[0][4]
    d.rounded_rectangle((x0, size + g, x0 + aw, 2 * size + g), radius=round(size * 0.22),
                        fill=(255, 120, 70, 240) if lit else (0, 0, 0, 120),
                        outline=(255, 255, 255, 200) if lit else (200, 200, 210, 110), width=1)
    tw = CP.text_w("ATK", fs, "SemiBold")
    CP.draw_text(d, x0 + (aw - tw) / 2, size + g + (size - fs) / 2 - 1, "ATK", fs,
                 (20, 20, 25, 255) if lit else (210, 212, 220, 220), "SemiBold")
    a = np.asarray(im, np.float32)
    _HUD[key] = (a[..., :3].astype(np.uint8), a[..., 3:] / 255.0)
    return _HUD[key]


def load_blocks(case_dir: str) -> np.ndarray | None:
    p = f"{case_dir}/actions.npy"
    if os.path.exists(p):
        return np.load(p)
    meta = json.load(open(f"{case_dir}/meta.json"))
    act = (meta["job"].get("edit") or {}).get("actions")
    return np.load(act) if act and os.path.exists(act) else None


# ------------------------------------------------------------------ 角色跟拍放大 ----
PLAYER_IDS, BOSS_IDS = (10, 12, 13, 14, 15, 16), (11, 17, 18, 19, 20, 21)


def zoom_boxes(ids: np.ndarray, target: str, factor: float, smooth: int = 7) -> list[tuple[int, int, int, int]]:
    """逐帧跟拍框：target（player / boss）在 ids[t] 里的掩码质心（看不见就沿用上一帧），滑动平均 smooth 帧，框 = 画面 / factor。"""
    F, h, w = ids.shape
    cs, last = [], np.array([w / 2, h / 2])
    for t in range(F):
        ys, xs = np.nonzero(np.isin(ids[t], PLAYER_IDS if target == "player" else BOSS_IDS))
        if len(xs) >= 30:
            last = np.array([xs.mean(), ys.mean()])
        cs.append(last)
    cs = np.asarray(cs)
    k = np.ones(smooth) / smooth
    cs = np.stack([np.convolve(np.pad(cs[:, i], smooth // 2, mode="edge"), k, mode="valid") for i in range(2)], 1)
    bw, bh = w / factor, h / factor
    out = []
    for cx, cy in cs:
        x0 = int(round(np.clip(cx - bw / 2, 0, w - bw)))
        y0 = int(round(np.clip(cy - bh / 2, 0, h - bh)))
        out.append((x0, y0, x0 + int(bw), y0 + int(bh)))
    return out


# ------------------------------------------------------------------ 通用宫格 ----
def tiles_video(out: str, tiles: list[dict], cols: int, title: str | None = None, sub: str | None = None,
                margin: int = 40, gap: int = 14, crf: int = 16, nf: int = 81) -> None:
    """tiles：[{kind: gen|sem|img, case（dir）, label, hud（bool）, inset（dir：右上角小窗，取它的 gen）}]，同步播放。"""
    top = 0
    if title:
        top = 64 if sub else 48
    rows, tw, th, pos = CP.grid_layout(len(tiles), cols, W, H - top, 832 / 480, margin, gap)
    pos = [(x, y + top) for x, y in pos]
    data = []
    for t in tiles:
        d = dict(t)
        if t["kind"] == "gen":
            d["frames"] = CP.read_video(f"{t['case']}/gen.mp4", nf)
            d["blocks"] = load_blocks(t["case"]) if t.get("hud") else None
        elif t["kind"] == "sem":
            d["ids"] = CP.load_case(t["case"])["ids"]()[:nf]
        if t.get("inset"):
            d["inset_frames"] = CP.read_video(f"{t['inset']}/gen.mp4", nf)
        if t.get("zoom") and t["kind"] == "gen":                         # (target, factor)：跟拍放大
            ids = CP.load_case(t["case"])["ids"]()[:nf]
            boxes = zoom_boxes(ids, *t["zoom"])
            d["frames"] = np.stack([cv2.resize(f[y0:y1, x0:x1], (f.shape[1], f.shape[0]), interpolation=cv2.INTER_LANCZOS4)
                                    for f, (x0, y0, x1, y1) in zip(d["frames"], boxes)])
        data.append(d)
    fs = min(24, max(14, round(th * 0.06)))
    hud_size = min(28, max(16, round(th * 0.07)))
    for d in data:
        d["pill"] = CP.pill(d["label"], fs, alpha=0.62) if d.get("label") else None
    base = Image.new("RGB", (W, H), CP.BG)
    if title:
        dr = ImageDraw.Draw(base)
        CP.draw_text(dr, margin, 14, title, 28, CP.FG, "SemiBold")
        if sub:
            CP.draw_text(dr, margin, 14 + 34, sub, 18, CP.MUTED)
    base = np.asarray(base).copy()
    for x, y in pos:
        CP.frame_rect(base, x, y, tw, th)

    def frames():
        for i in range(nf):
            fr = base.copy()
            for d, (x, y) in zip(data, pos):
                if d["kind"] == "gen":
                    tile = CP.fit(d["frames"][min(i, len(d["frames"]) - 1)], tw, th)
                elif d["kind"] == "sem":
                    tile = CP.sem_panel(d["ids"][min(i, len(d["ids"]) - 1)], tw, th)
                else:
                    tile = CP.fit(d["img"], tw, th)
                fr[y:y + th, x:x + tw] = tile
                if d.get("inset_frames") is not None:
                    iw, ih = tw // 4, th // 4
                    fr[y + 6:y + 6 + ih, x + tw - iw - 6:x + tw - 6] = CP.fit(d["inset_frames"][i], iw, ih)
                    CP.frame_rect(fr, x + tw - iw - 6, y + 6, iw, ih, (235, 235, 240), 1)
                if d["pill"] is not None:
                    CP.paste_rgba(fr, d["pill"], x + round(fs * 0.5), y + th - d["pill"][0].shape[0] - round(fs * 0.5))
                if d.get("blocks") is not None:
                    hud = keys_hud(d["blocks"][block_of(i)], hud_size)
                    CP.paste_rgba(fr, hud, x + tw - hud[0].shape[1] - 8, y + th - hud[0].shape[0] - 8)
            yield fr
    os.makedirs(os.path.dirname(out), exist_ok=True)
    n = CP.write_video(out, frames(), crf=crf)
    print(f"→ {out}（{len(tiles)} 格，{rows}×{cols}，{n} 帧）", flush=True)


def sheet(paths: list[str], labels: list[str], out: str, frames=(0, 20, 40, 60, 80), tw: int = 312, sem: str | None = None):
    import edit_app as E
    E.sheet(paths, labels, out, frames, tw)


# ------------------------------------------------------------------ app：外观 / 镜像 / 文本 ----
def compose_app(a) -> None:
    import edit_app as E
    only = set(a.only or ["look", "act", "all_look", "gallery"])
    R = f"{A}/roll"
    bases = E.load_bases()
    done = lambda n: os.path.exists(f"{R}/{n}/meta.json")  # noqa: E731
    if "look" in only:
        for s, b in bases.items():
            vs = [v for v in ["orig"] + b["variants"] if done(f"{s}__{v}")]
            if len(vs) < 2:
                continue
            tiles = [{"kind": "sem", "case": f"{R}/{s}__orig", "label": "Semantic map (shared)"}] + \
                    [{"kind": "gen", "case": f"{R}/{s}__{v}", "label": E.EDIT_EN[v]} for v in vs]
            cols = 4 if len(tiles) > 9 else 3
            tiles_video(f"{O}/look_{s}.mp4", tiles, cols, "Same semantic map, different first frames",
                        "Layout from the semantic map · appearance from the first frame (edited with Qwen-Image) · "
                        "same keys")
            sheet([f"{R}/{s}__{v}" for v in vs], [f"{s}: {E.EDIT_EN[v]}" for v in vs], f"{O}/sheets/look_{s}.jpg")
    if "gallery" in only:
        for s, who in E.GALLERY.items():
            vs = [v for v in ["orig"] + E.gallery_variants(s) if done(f"{s}__{v}")]
            if len(vs) < 3:
                continue
            tiles = [{"kind": "sem", "case": f"{R}/{s}__orig", "label": "Semantic map (shared)"}] + \
                    [{"kind": "gen", "case": f"{R}/{s}__{v}", "label": E.EDIT_EN[v]} for v in vs]
            tiles_video(f"{O}/gallery_{s}_{who}.mp4", tiles, 4, f"Same map, different {'heroes' if who == 'hero' else 'bosses'}",
                        f"Only the {'hero' if who == 'hero' else 'boss'} in the first frame is edited (local Qwen-Image edit, "
                        "pixels outside the character mask pinned) · same semantic map and keys")
            sheet([f"{R}/{s}__{v}" for v in vs], [f"{s}: {E.EDIT_EN[v]}" for v in vs], f"{O}/sheets/gallery_{s}_{who}.jpg")
            zt = [{"kind": "gen", "case": f"{R}/{s}__{v}", "label": E.EDIT_EN[v],
                   "zoom": ("player" if who == "hero" else "boss", 2.6 if who == "hero" else 1.8)} for v in vs]
            tiles_video(f"{O}/gallery_{s}_{who}_zoom.mp4", zt, 4,
                        f"Same map, different {'heroes' if who == 'hero' else 'bosses'} (tracking close-up)",
                        "Camera crop follows the character's semantic mask · identity comes from the edited first frame, "
                        "pose and motion from the semantic capsules")
    if "all_look" in only:
        allv = sorted({v for b in bases.values() for v in b["variants"]})
        for v in allv:
            names = [s for s, b in bases.items() if v in b["variants"] and done(f"{s}__{v}")]
            if len(names) < 3:
                continue
            tiles = [{"kind": "gen", "case": f"{R}/{s}__{v}", "label": f"{s}", "inset": f"{R}/{s}__orig"}
                     for s in names[:12]]
            tiles_video(f"{O}/look_all_{v}.mp4", tiles, 4 if len(tiles) > 9 else 3, f"Edit: {E.EDIT_EN[v]}",
                        "Each tile: edited first frame + the original semantic map and keys (inset = original rollout)")
    if "act" in only:
        for s in E.MIRROR_BASES:
            if done(f"{s}__orig") and done(f"{s}__mirror"):
                tiles = [{"kind": "sem", "case": f"{R}/{s}__orig", "label": "Semantic map"},
                         {"kind": "sem", "case": f"{R}/{s}__mirror", "label": "Mirrored semantic map"},
                         {"kind": "gen", "case": f"{R}/{s}__orig", "label": "Original keys", "hud": True},
                         {"kind": "gen", "case": f"{R}/{s}__mirror", "label": "Mirrored world, A/D swapped",
                          "hud": True}]
                tiles_video(f"{O}/mirror_{s}.mp4", tiles, 2, "Mirror = action control",
                            "Semantic map + first frame mirrored left-right, keys A and D swapped: same world, opposite way")
                sheet([f"{R}/{s}__orig", f"{R}/{s}__mirror"], [f"{s}: original", f"{s}: mirrored (A<->D)"],
                      f"{O}/sheets/mirror_{s}.jpg")
        for s in E.TEXT_BASES:
            names = [f"{s}__orig"] + [f"{s}__text_{t}" for t in E.TEXT_VARIANTS]
            if all(done(n) for n in names):
                labs = ["Original text"] + [E.TEXT_EN[t] for t in E.TEXT_VARIANTS]
                tiles = [{"kind": "gen", "case": f"{R}/{n}", "label": l, "hud": True} for n, l in zip(names, labs)]
                tiles_video(f"{O}/text_{s}.mp4", tiles, 2, "Only the action text changes",
                            "Same semantic map + same first frame; keys / prompts replaced (HUD = the text given)")
                sheet([f"{R}/{n}" for n in names], [f"{s}: {l}" for l in labs], f"{O}/sheets/text_{s}.jpg")


# ------------------------------------------------------------------ sameworld ----
def compose_sameworld(name: str, a=None) -> None:
    S = f"{A}/sameworld"
    info = json.load(open(f"{S}/maps/{name}/group.json"))
    cases = [f"{S}/roll/{name}__{c['k']}_{c['start']:06d}" for c in info["clips"]]
    cases = [c for c in cases if os.path.exists(f"{c}/meta.json")]
    if len(cases) < 2:
        print(f"{name}：完成的 rollout 不够（{len(cases)}）")
        return
    clips = [c for c in info["clips"] if os.path.exists(f"{S}/roll/{name}__{c['k']}_{c['start']:06d}/meta.json")]
    plan = cv2.cvtColor(cv2.imread(info["plan"]), cv2.COLOR_BGR2RGB)
    nf = 81
    gens = [CP.read_video(f"{c}/gen.mp4", nf) for c in cases]
    blocks = [load_blocks(c) for c in cases]
    top = 64
    mg, gap = 36, 14
    # 左：plan 裁到各 clip 帧 0 视野 + 轨迹的外接框（外扩 15 %），缩到宽 ≤ PLAN_W；右：宫格
    pts = np.concatenate([np.asarray(c["foot0_px"], np.float32) for c in clips] +
                         [np.asarray(c["path_px"], np.float32) for c in clips], 0)
    lo, hi = pts.min(0), pts.max(0)
    pad = 0.15 * (hi - lo).max()
    x0c, y0c = int(max(0, lo[0] - pad)), int(max(0, lo[1] - pad))
    x1c, y1c = int(min(plan.shape[1], hi[0] + pad)), int(min(plan.shape[0], hi[1] + pad))
    plan = plan[y0c:y1c, x0c:x1c]
    PLAN_W, PLAN_H = 560, H - top - 2 * mg
    sc = min(PLAN_W / plan.shape[1], PLAN_H / plan.shape[0])
    pw, ph = round(plan.shape[1] * sc), round(plan.shape[0] * sc)
    plan_s = cv2.resize(plan, (pw, ph), interpolation=cv2.INTER_AREA if sc < 1 else cv2.INTER_LINEAR)
    py0 = top + (H - top - ph) // 2
    rx0 = mg + pw + gap
    n = len(cases)
    cols = 2 if n > 1 else 1
    rows, tw, th, pos = CP.grid_layout(n, cols, W - rx0, H - top, 832 / 480, mg, gap)
    pos = [(x + rx0, y + top) for x, y in pos]
    fs = min(24, max(14, round(th * 0.06)))
    pills = []
    for c in clips:
        col = tuple(int(v) for v in CP_COLORS[c["k"] % len(CP_COLORS)])
        pills.append(CP.pill(f"Play {c['k'] + 1}", fs, fg=(20, 20, 25), bg=col, alpha=0.9))
    mode = info.get("mode")
    title = ("Same world, different playthroughs" if mode == "union" else "Same start frame, different keys")
    sub = ("One procedural world built from the union of the clips' footprints · each clip: own camera, own keys, "
           "own first frame (same scene prompt + seed)" if mode == "union" else
           "Clips translated to one common start (boss removed) · one shared first frame · each clip: its own keys")
    base = Image.new("RGB", (W, H), CP.BG)
    dr = ImageDraw.Draw(base)
    CP.draw_text(dr, mg, 14, title, 28, CP.FG, "SemiBold")
    CP.draw_text(dr, mg, 48, sub, 17, CP.MUTED)
    base = np.asarray(base).copy()
    base[py0:py0 + ph, mg:mg + pw] = plan_s
    CP.frame_rect(base, mg, py0, pw, ph)
    for x, y in pos:
        CP.frame_rect(base, x, y, tw, th)
    hud_size = min(28, max(16, round(th * 0.07)))
    paths = [(np.asarray(c["path_px"], np.float32) - [x0c, y0c]) * sc for c in clips]

    def frames():
        for i in range(nf):
            fr = base.copy()
            for k, (c, P) in enumerate(zip(clips, paths)):
                j = min(len(P) - 1, round(i * (len(P) - 1) / (nf - 1)))
                col = tuple(int(v) for v in CP_COLORS[c["k"] % len(CP_COLORS)])
                cx, cy = int(round(mg + P[j][0])), int(round(py0 + P[j][1]))
                cv2.circle(fr, (cx, cy), 9, (0, 0, 0), -1, cv2.LINE_AA)
                cv2.circle(fr, (cx, cy), 7, col, -1, cv2.LINE_AA)
            for k, (g, b, (x, y)) in enumerate(zip(gens, blocks, pos)):
                fr[y:y + th, x:x + tw] = CP.fit(g[min(i, len(g) - 1)], tw, th)
                CP.paste_rgba(fr, pills[k], x + round(fs * 0.5), y + round(fs * 0.5))
                if b is not None:
                    hud = keys_hud(b[block_of(i)], hud_size)
                    CP.paste_rgba(fr, hud, x + tw - hud[0].shape[1] - 8, y + th - hud[0].shape[0] - 8)
            yield fr
    out = f"{O}/sameworld_{name}.mp4"
    os.makedirs(O, exist_ok=True)
    CP.write_video(out, frames())
    print(f"→ {out}", flush=True)
    sheet(cases, [f"{name} play {c['k'] + 1}" for c in clips], f"{O}/sheets/sameworld_{name}.jpg")
    cv2.imwrite(f"{O}/sheets/sameworld_{name}_plan.png", cv2.cvtColor(plan, cv2.COLOR_RGB2BGR))


def compose_sameworld_looks(name: str, looks=("night", "snow")) -> None:
    """start 组：行 = 外观（原首帧 / night / snow），列 = 各 clip 的按键 ⇒ 语义（同一世界）× 首帧（外观）× 按键（动作）三者独立。"""
    S = f"{A}/sameworld"
    info = json.load(open(f"{S}/maps/{name}/group.json"))
    rows_ = [None] + [lk for lk in looks]
    tiles = []
    for lk in rows_:
        for c in info["clips"]:
            d = f"{S}/roll/{name}__{c['k']}_{c['start']:06d}" + (f"__{lk}" if lk else "")
            if not os.path.exists(f"{d}/meta.json"):
                return
            tiles.append({"kind": "gen", "case": d, "hud": True,
                          "label": f"Play {c['k'] + 1}" + (f" · {lk.capitalize()}" if lk else " · Original")})
    tiles_video(f"{O}/sameworld_{name}_looks.mp4", tiles, len(info["clips"]),
                "One world, one start: appearance x keys",
                "Rows: first frame (original / edited to night / snow) · columns: different keys · same semantic world")
    sheet([t["case"] for t in tiles], [t["label"] for t in tiles], f"{O}/sheets/sameworld_{name}_looks.jpg")


CP_COLORS = [(255, 64, 200), (40, 220, 255), (255, 200, 40), (120, 255, 90), (255, 110, 60), (170, 120, 255)]  # = edit_sameworld.COLORS


def main(argv=None) -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    sp = ap.add_subparsers(dest="cmd", required=True)
    p = sp.add_parser("app")
    p.add_argument("--only", nargs="*", default=None)
    q = sp.add_parser("sameworld")
    q.add_argument("groups", nargs="+")
    a = ap.parse_args(argv)
    if a.cmd == "app":
        compose_app(a)
    else:
        for g in a.groups:
            compose_sameworld(g, a)
            compose_sameworld_looks(g)


if __name__ == "__main__":
    main()