#!/usr/bin/env python3
"""teaser · 编辑 case「60 秒外观保持」：同一张程序化语义地图 × 不同首帧外观（原版 / 夜晚 / 下雪 / 黄昏），因果 DMD 学生跑 241 cell
（= 961 帧 = 60 s，部署回路：每 cell 4 步、sink 3 / window 17）。产物在 outputs/teaser/edit/app/long/（$L）。maliang venv。

世界：genbench3 的连续源链（outputs/genbench3/sources.json，11 条 has_ext 的 clip，start 相差 96 帧、cond_ext 与下一条开头逐位
连续）—— 每条 clip 取 meta 的前 96 帧（81 + cond_ext 的 15 帧）首尾相接，截到 961 帧 = 一条 60 s 的真实相机 + 真实胶囊轨迹。
在它的世界坐标里用 procmap.World + themes_ext 的主题按全部 961 帧的足迹搭世界（构造 World / 遮挡控制 / 渲染时临时
procmap.F = 961），渲染 961 帧 parts ids。prompts = gcdistill.eval_rollout.load_chain 给的 241 条（真实按键）。

子命令
  build  [--tags ...]（GPU）→ $L/maps/<tag>/（ids.npz [961,480,832]、prompts.json（241 条）、actions.npy（[240,5]，由 prompt 反推的
         规范按键 w a s d attack，cell k ← 行 k−1）、map0.png / map0_shaded.png / depth0.png / source.json、plan.png、contact.jpg）
  f0jobs → $L/f0/jobs.jsonl（f0gen qwen21）；f0edits → $L/f0/edits.jsonl（edit_app_f0 整图编辑：LOOKS）
  roll   （GPU）DMD 部署回路，每个 (tag, look) → $L/roll/<tag>__<look>/：gen.mp4（961 帧 16 fps，TAEHV 解码）、ids.npz、
         first.png、prompts.json、actions.npy、frames.npz（每秒 1 帧）、meta.json（edit 键同 edit_app）
  compose → $A/showcase/long60_<tag>.mp4（各外观并排 + 按键 HUD + 语义小窗）+ sheets/long60_<tag>.jpg（每 10 s 一帧）

usage（env 同 SPEC §0；在 $M 下）：
  CUDA_VISIBLE_DEVICES=6 $PY scripts/teaser/edit_long.py build
  $PY scripts/teaser/edit_long.py f0jobs && CUDA_VISIBLE_DEVICES=6 $M/envs/img/bin/python scripts/teaser/f0gen.py --backend qwen21 \\
      --jobs outputs/teaser/edit/app/long/f0/jobs.jsonl
  $PY scripts/teaser/edit_long.py f0fix && CUDA_VISIBLE_DEVICES=6 $M/envs/img/bin/python scripts/teaser/edit_app_f0.py \\
      --jobs outputs/teaser/edit/app/long/f0/fix.jsonl          # Boss 人偶 → 真怪物（局部编辑，掩码外钉回）
  $PY scripts/teaser/edit_long.py f0edits && CUDA_VISIBLE_DEVICES=6 $M/envs/img/bin/python scripts/teaser/edit_app_f0.py \\
      --jobs outputs/teaser/edit/app/long/f0/edits.jsonl
  CUDA_VISIBLE_DEVICES=6 $PY scripts/teaser/edit_long.py roll ; $PY scripts/teaser/edit_long.py compose
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time

import numpy as np

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

A = f"{M}/outputs/teaser/edit/app"
L = f"{A}/long"
DATA = f"{M}/data/gc25d_v10"
DMD = f"{M}/outputs/distill/gc25d_v10/stage3_dmd/ema/ema_step500.safetensors"
N_CELLS = 241
N_FRAMES = 961
STRIDE = 96
CHAINS = {   # 名字 → (genbench3 源链 tag, 主题, 世界种子, 首帧场景)。主题要能铺满 60 s 的整条轨迹：maze / ruins 只在帧 0 附近
    # 搭一块（第一版 maze：40 s 之后走出迷宫、全是空地）；cityblock 在近景相机（R10_s3505，视野 ≈ 12 m）下街区太大、画面多是空地；
    # orchard / village 沿整条路径铺满。
    "train_v10_03": ("train_v10_03", "cityblock", 7, "tibetan@photoreal"),
    "val_v10_01": ("val_v10_01", "orchard", 4, "farm@photoreal"),
    "train_v10_03_village": ("train_v10_03", "village", 11, "tibetan@photoreal"),
}
LOOKS = ("night", "snow", "dusk")


def chain_of(tag: str) -> list[str]:
    S = {s["tag"]: s for s in json.load(open(f"{M}/outputs/genbench3/sources.json"))}
    return S[tag]["chain"]


def chain_src(tag: str) -> dict:
    """11 条 clip 的 meta（前 96 帧）首尾相接 → procmap.load_source 同形的 src（meta 961 帧）+ 241 条 prompts。"""
    import procmap as PM
    from gcdistill import data_gc as DG
    from gcdistill import eval_rollout as ER
    clips = chain_of(tag)
    keys = ("world_to_clip", "world_to_view", "capsules", "actor_pos", "ground_z")
    parts = {k: [] for k in keys}
    for c in clips:
        m = np.load(f"{DATA}/meta/{c}.npz")
        assert m["world_to_clip"].shape[0] >= STRIDE, (c, m["world_to_clip"].shape)
        for k in keys:
            parts[k].append(np.asarray(m[k][:STRIDE]))
        ortho = bool(m["orthographic"])
    meta = {k: np.concatenate(v, 0)[:N_FRAMES] for k, v in parts.items()}
    meta["ortho"] = ortho
    src0 = PM.load_source(clips[0])
    table = DG.load_text_table(os.path.join(DG.DATA_ROOT, "text_table_gc25d.pt"))
    kid = DG.key_to_id(table)
    idx = {}
    for line in open(f"{DATA}/index/clips.jsonl"):
        r = json.loads(line)
        if r["clip_id"] in clips:
            idx[r["clip_id"]] = r
    ch = ER.load_chain(DG.DATA_ROOT, [idx[c] for c in clips], N_CELLS, table, kid, "parts")
    return dict(clip_id=clips[0], row=src0["row"], meta=meta, prompts=list(ch["prompts"]), clips=clips,
                blocks=src0["blocks"])


def prompt_to_block(p: str) -> list[int]:
    import text as TXT
    inv = {v: k for k, v in TXT.MOTIONS.items()}
    body = p.replace(TXT.TEMPLATE.split("{")[0], "").rstrip(".")
    atk = body.endswith(TXT.ATTACK_SUFFIX)
    motion = body[:-len(TXT.ATTACK_SUFFIX)] if atk else body
    up, right = inv.get(motion, (0, 0))
    return [int(up > 0), int(right < 0), int(up < 0), int(right > 0), int(atk)]


def build(tag: str, dev: str = "cuda") -> None:
    import cv2
    import procmap as PM
    import themes_ext  # noqa: F401
    import geometry as GEO
    import model as GM
    import edit_sameworld as SW
    src_tag, theme, seed, scene = CHAINS[tag]
    od = f"{L}/maps/{tag}"
    os.makedirs(od, exist_ok=True)
    t0 = time.time()
    src = chain_src(src_tag)
    PM.F = N_FRAMES
    try:
        w = PM.World(src, seed)
        PM.THEME_FN[theme](w)
        cache = [PM.char_rays(src["meta"], f) for f in range(N_FRAMES)]
        occ = PM.control_occlusion(w, dev, cache)
        print(f"[{tag}] 世界 {w.nx * PM.GRID_CM / 100:.0f}×{w.ny * PM.GRID_CM / 100:.0f} m · {len(w.prims())} 图元 · "
              f"遮挡控制删 {len(occ['removed'])} · {time.time() - t0:.0f}s", flush=True)
        ids, keep = PM.render(w, dev)
    finally:
        PM.F = 81
    assert ids.shape[0] == N_FRAMES and set(np.unique(ids).tolist()) <= GM.valid_raw_ids("parts")
    coll = GM.collapse_parts(ids)
    ppx = (coll == GEO.PLAYER_ID).reshape(N_FRAMES, -1).sum(1)
    np.savez_compressed(f"{od}/ids.npz", ids=ids)
    json.dump(src["prompts"], open(f"{od}/prompts.json", "w"), indent=1)
    blocks = np.asarray([prompt_to_block(p) for p in src["prompts"][1:]], np.uint8)
    np.save(f"{od}/actions.npy", blocks)
    json.dump(dict(chain_tag=src_tag, name=tag, clips=src["clips"], theme=theme, theme_desc=PM.THEME_DESC[theme], seed=seed,
                   f0_scene=scene, ortho=src["meta"]["ortho"], z_ground_cm=w.z0, n_frames=N_FRAMES, n_cells=N_CELLS,
                   frames_per_clip=STRIDE, generator="scripts/teaser/edit_long.py", clip_id=src["clips"][0],
                   subset=src["row"]["subset"], run=src["row"]["run"]), open(f"{od}/source.json", "w"), indent=1)
    rep = dict(tag=tag, player_px_f0=int(ppx[0]), player_px_min=int(ppx.min()), frames_player_lt200=int((ppx < 200).sum()),
               n_prims=len(w.prims()), occlusion_control=occ, sec=round(time.time() - t0, 1))
    json.dump(rep, open(f"{od}/report.json", "w"), indent=1)
    cv2.imwrite(f"{od}/map0.png", PM.frame_rgb(w, ids, 0)[..., ::-1])
    cv2.imwrite(f"{od}/map0_shaded.png", PM.shaded0(w, ids, keep[0])[..., ::-1])
    cv2.imwrite(f"{od}/depth0.png", PM.depth0(w, ids, keep[0]))
    tiles = [PM.label(cv2.resize(PM.frame_rgb(w, ids, f), (PM.W // 2, PM.H // 2), interpolation=cv2.INTER_AREA),
                      f"{f / 16:.0f}s") for f in range(0, N_FRAMES, 160)]
    cv2.imwrite(f"{od}/contact.jpg", np.concatenate(tiles, 1)[..., ::-1], [cv2.IMWRITE_JPEG_QUALITY, 90])
    plan, q, res = SW.plan_image(w, [src], dev)
    P = src["meta"]["actor_pos"][:, 0, :2]
    P = P[np.isfinite(P).all(1)]
    pp = q(P).round().astype(np.int32)
    cv2.polylines(plan, [pp], False, (0, 0, 0), 5, cv2.LINE_AA)
    cv2.polylines(plan, [pp], False, (200, 64, 255), 3, cv2.LINE_AA)
    cv2.circle(plan, tuple(pp[0]), 7, (200, 64, 255), -1, cv2.LINE_AA)
    cv2.imwrite(f"{od}/plan.png", plan)
    json.dump({"path_px": q(P).tolist(), "res_cm_per_px": res}, open(f"{od}/plan.json", "w"))
    print(f"[{tag}] 完成：帧 0 玩家 {ppx[0]} px · 最少 {ppx.min()} px · {time.time() - t0:.0f}s", flush=True)


def cmd_build(a) -> None:
    for t in a.tags or list(CHAINS):
        build(t)


def cmd_f0jobs(a) -> None:
    os.makedirs(f"{L}/f0", exist_ok=True)
    rows = [{"case": f"long_{t}", "map_dir": f"{L}/maps/{t}", "scene": CHAINS[t][3], "seed": 0,
             "out": f"{L}/f0/{t}/qwen21_s0.png"} for t in CHAINS if os.path.exists(f"{L}/maps/{t}/report.json")]
    with open(f"{L}/f0/jobs.jsonl", "w") as fh:
        for r in rows:
            fh.write(json.dumps(r) + "\n")
    print(f"{len(rows)} 个首帧任务 → {L}/f0/jobs.jsonl")


BOSS_FIX = ("Edit picture 1: the smooth pink mannequin-like figure is the boss; turn it into a huge, fully detailed living "
            "monster boss: a hulking ogre with grey-green leathery skin, rough leather armour and iron bracers, holding the "
            "same big weapon.")


def base_first(t: str) -> str:
    """链 t 的底子首帧：有角色二次编辑版（_cf，Boss 人偶修成真怪物）就用它。"""
    cf = f"{L}/f0/{t}/qwen21_s0_cf.png"
    return cf if os.path.exists(cf) else f"{L}/f0/{t}/qwen21_s0.png"


def look_first(t: str, look: str | None) -> str:
    b = base_first(t)
    return b if not look else b[:-4] + f"__{look}.png"


def cmd_f0fix(a) -> None:
    """f0gen 的 afix 没认出来的 Boss 人偶（qwen21 照抄体块图的光滑鲑红胶囊）→ edit_app_f0 的 char 局部编辑任务 $L/f0/fix.jsonl。"""
    rows = []
    for t in (a.tags or list(CHAINS)):
        src = f"{L}/f0/{t}/qwen21_s0.png"
        if os.path.exists(src):
            rows.append({"out": f"{L}/f0/{t}/qwen21_s0_cf.png", "src": src, "ids": f"{L}/maps/{t}/ids.npz", "kind": "char",
                         "target": "boss", "prompt": BOSS_FIX, "seed": 5, "chain": t, "dilate_frac": 0.3})
    with open(f"{L}/f0/fix.jsonl", "w") as fh:
        for r in rows:
            fh.write(json.dumps(r) + "\n")
    print(f"{len(rows)} 个 Boss 修补任务 → {L}/f0/fix.jsonl")


def cmd_f0edits(a) -> None:
    import edit_app as E
    rows = []
    for t in CHAINS:
        src = base_first(t)
        if not os.path.exists(src):
            continue
        for v in LOOKS:
            rows.append({"out": look_first(t, v), "src": src, "ids": f"{L}/maps/{t}/ids.npz", "best_of": 3,
                         "kind": "global", "prompt": E.GLOBAL[v], "seed": E.SEED[v], "variant": v, "chain": t})
    with open(f"{L}/f0/edits.jsonl", "w") as fh:
        for r in rows:
            fh.write(json.dumps(r) + "\n")
    print(f"{len(rows)} 个外观编辑任务 → {L}/f0/edits.jsonl")


def cmd_roll(a) -> None:
    import torch
    import cv2
    import infer as INF
    import sample as S
    from gcdistill import data_gc as DG
    from gcdistill import eval_rollout as ER
    import compose as CP
    dev = torch.device("cuda")
    vae = INF.load_vae("cuda", S.MODEL_DIR)
    table = DG.load_text_table(os.path.join(DG.DATA_ROOT, "text_table_gc25d.pt"))
    kid = DG.key_to_id(table)
    model = ER.load_student(DMD, dev)
    vocab = ER.GD.GM.sem_vocab(model.dit)
    dec = ER.Decoder("tae", dev, model_dir=S.MODEL_DIR, taehv_dir=ER.TAEHV_DIR, taehv_weights=ER.TAEHV_W)
    for t in (a.tags or list(CHAINS)):
        md = f"{L}/maps/{t}"
        if not os.path.exists(f"{md}/report.json"):
            continue
        ids = np.load(f"{md}/ids.npz")["ids"]
        prompts = json.load(open(f"{md}/prompts.json"))
        ctx, cids = DG.cell_context(prompts, table, kid)
        sem = torch.from_numpy(np.ascontiguousarray(DG.to_vocab(ids, vocab)))
        for look in (None,) + LOOKS:
            first = look_first(t, look)
            name = f"{t}__{look or 'orig'}"
            od = f"{L}/roll/{name}"
            if not os.path.exists(first) or os.path.exists(f"{od}/meta.json"):
                continue
            t0 = time.time()
            f0 = cv2.cvtColor(cv2.imread(first), cv2.COLOR_BGR2RGB)
            z0 = INF.encode_first_frame(vae, f0)
            b = DG.rollout_collate([{"sem_ids": sem, "context": ctx, "context_ids": cids, "z0": z0[0].cpu()}])
            timing = []
            lat = ER.deploy_rollout(model, b["z0"].to(dev), b["context"].to(dev), b["context_ids"].to(dev),
                                    b["sem_ids"].to(dev), N_CELLS - 1, seed=0, timing=timing)
            assert torch.isfinite(lat.float()).all()
            frames = np.asarray(dec(lat))
            os.makedirs(od, exist_ok=True)
            CP.write_video(f"{od}/gen.mp4", frames)
            np.savez_compressed(f"{od}/ids.npz", ids=ids)
            cv2.imwrite(f"{od}/first.png", cv2.imread(first))
            json.dump(prompts, open(f"{od}/prompts.json", "w"), indent=1)
            np.save(f"{od}/actions.npy", np.load(f"{md}/actions.npy"))
            idx = list(range(0, len(frames), 16))
            np.savez_compressed(f"{od}/frames.npz", idx=np.asarray(idx), rgb=frames[idx], ids=ids[idx])
            src = json.load(open(f"{md}/source.json"))
            meta = {"kind": "teaser_rollout_dmd", "name": name, "mode": "dmd_deploy", "ckpt": DMD, "cells": N_CELLS,
                    "n_frames": int(len(frames)), "seed": 0, "decoder": "taehv", "denoise_steps_per_cell": 4,
                    "sink": 3, "kv_window": 17, "dit_ms_per_cell": float(np.mean(timing[2:])),
                    "sec": round(time.time() - t0, 1), "prompts": prompts,
                    "edit": {"type": "appearance_long", "variant": look or "orig", "kind": "global" if look else None,
                             "edit_frame": 0 if look else None, "chain_tag": t, "chain_clips": src["clips"],
                             "theme": src["theme"], "world_seed": src["seed"], "f0_scene": src["f0_scene"],
                             "first_src": base_first(t), "first_raw": f"{L}/f0/{t}/qwen21_s0.png", "editor": "Qwen-Image-2.1-Turbo" if look else None,
                             "map_dir": md, "label": (look or "original").capitalize(),
                             "desc": "60 s causal DMD rollout on one procedural semantic world; only the first frame differs"},
                    "files": {"video": "gen.mp4", "cond_ids": "ids.npz (961 frames)", "actions": "actions.npy ([240,5] w a s d "
                              "attack, canonical keys of the 240 cell prompts, cell k <- row k-1)", "prompts": "prompts.json",
                              "first_frame": "first.png"}}
            json.dump(meta, open(f"{od}/meta.json", "w"), indent=1)
            print(f"[long] {name}: {meta['sec']}s · DiT {meta['dit_ms_per_cell']:.0f} ms/cell", flush=True)
            del lat, frames


def cmd_compose(a) -> None:
    import cv2
    import compose as CP
    import edit_compose as EC
    O = f"{A}/showcase"
    for t in CHAINS:
        names = [f"{t}__{v}" for v in ("orig",) + LOOKS if os.path.exists(f"{L}/roll/{t}__{v}/meta.json")]
        if len(names) < 2:
            continue
        tiles = [{"kind": "gen", "case": f"{L}/roll/{n}", "hud": True, "label": n.split("__")[1].capitalize()
                  .replace("Orig", "Original").replace("Dusk", "Sunset")} for n in names]
        nf = a.frames or N_FRAMES
        EC.tiles_video(f"{O}/long60_{t}.mp4", tiles, 2, "60 seconds, one semantic world, four first frames",
                       "Causal distilled model (4 steps / cell) · same semantic map + same keys · only the first frame is "
                       "edited (night / snow / sunset)", nf=nf)
        EC.sheet([f"{L}/roll/{n}" for n in names], names, f"{O}/sheets/long60_{t}.jpg",
                 frames=tuple(range(0, N_FRAMES, 160)), tw=240)


def main(argv=None) -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    sp = ap.add_subparsers(dest="cmd", required=True)
    b = sp.add_parser("build")
    b.add_argument("--tags", nargs="*", default=None)
    sp.add_parser("f0jobs")
    sp.add_parser("f0edits")
    x = sp.add_parser("f0fix")
    x.add_argument("--tags", nargs="*", default=None)
    r = sp.add_parser("roll")
    r.add_argument("--tags", nargs="*", default=None)
    c = sp.add_parser("compose")
    c.add_argument("--frames", type=int, default=0)
    a = ap.parse_args(argv)
    {"build": cmd_build, "f0jobs": cmd_f0jobs, "f0edits": cmd_f0edits, "f0fix": cmd_f0fix, "roll": cmd_roll, "compose": cmd_compose}[a.cmd](a)


if __name__ == "__main__":
    main()