#!/usr/bin/env python3
"""teaser · 编辑 case「同一世界、不同游玩」（maliang venv；build 要 GPU，其余 CPU）。产物在 outputs/teaser/edit/app/sameworld/（$S）。

两种组（GROUPS）
----------------
* union（任务原意）：同一 run（同一游戏世界坐标）里相邻的几条 clip（不同轨迹 / 不同按键），世界坐标原样不动。用它们**足迹的并集**
  搭一个 procmap 世界（procmap.World 的足迹 / 净空 / 视野按所有 clip 的全部帧算 —— 把各 clip 的 meta 沿帧轴拼起来，构造 World
  时临时把 procmap.F 设成总帧数；主题 = themes_ext 的布局函数，原样调用），遮挡控制也对所有 clip 的全部帧做，然后每条 clip 用
  自己的相机 + 胶囊各渲染一份 81 帧 ids ⇒「同一张地图、不同的按键 / 走位」。每条的帧 0 视角不同 ⇒ 首帧各自生成
  （f0gen qwen21 默认配方，提示词 / 场景 / 种子一致）。
* start（加做）：同一 run、同一相机档位（camera_variant）的几条 clip，按帧 0 玩家位置**平移**到同一个起点（相机矩阵右乘平移：
  行向量约定 M' = T(−t)·M；胶囊 / actor_pos / ground_z 同步平移），Boss 去掉 ⇒ 帧 0 的相机完全重合、语义只差玩家姿势 ⇒
  **一张首帧**（锚 clip 的首帧）+ 不同的按键 / 轨迹 ⇒ 同一个开局、不同的未来。世界同样按平移后足迹的并集搭。

每条 clip 的地图目录 $S/maps/<group>/<k>_<start>/ = procmap 标准文件（ids.npz / prompts.json / actions.npy / source.json /
report.json / map0.png / map0_shaded.png / depth0.png / contact.jpg），f0gen 地图模式直接能用；组目录另有 plan.png（俯视总图：
所有 clip 的轨迹各一种颜色 + 帧 0 视野框）与 group.json。

子命令
------
build     [--groups g1 g2]（GPU）：搭世界 + 渲染每条 clip 的 ids。
f0jobs    → $S/f0/jobs.jsonl（f0gen.py --backend qwen21 吃；union 组每条 clip 一张；start 组只出锚 clip 一张）。
f0edits   → $S/f0/edits.jsonl（edit_app_f0.py 吃）：start 组锚首帧的整图外观编辑（START_LOOKS = night / snow）
          ⇒ 同一开局 × 外观 × 按键（rolljobs 自动带上：名字后缀 __night / __snow）。
rolljobs  → $S/jobs.jsonl（rollout.py 吃；结果 $S/roll/<group>__<k>_<start>/）。
compose   → $A/showcase/sameworld_<group>.mp4（左 = 俯视总图 + 当前帧各 clip 的位置，右 = 各 clip 的生成视频宫格 + 按键 HUD）
          + $A/showcase/sheets/sameworld_<group>.jpg。

usage（env 同 SPEC §0；在 $M 下）：
  CUDA_VISIBLE_DEVICES=6 $PY scripts/teaser/edit_sameworld.py build
  $PY scripts/teaser/edit_sameworld.py f0jobs
  CUDA_VISIBLE_DEVICES=6 $M/envs/img/bin/python scripts/teaser/f0gen.py --backend qwen21 --jobs outputs/teaser/edit/app/sameworld/f0/jobs.jsonl
  $PY scripts/teaser/edit_sameworld.py rolljobs   # 然后 rollout.py --jobs outputs/teaser/edit/app/sameworld/jobs.jsonl
  $PY scripts/teaser/edit_sameworld.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"
S = f"{A}/sameworld"
F = 81
COLORS = [(255, 64, 200), (40, 220, 255), (255, 200, 40), (120, 255, 90), (255, 110, 60), (170, 120, 255)]   # RGB

GROUPS = {
    # 选组：outputs/teaser/edit/app/sameworld/scan.json（全部 clip 的帧 0 视野中心 / 位移 / Boss 可见率）。第一版用过 R10_s8506 /
    # P9C_default_s1944（帧 0 视野宽只有 9–11 m，画面里几乎只有一种地面，看不出「同一张地图」）⇒ 换成视野 14–27 m 的组。
    # 帧 0 视野重叠的 3 条（同一片地方出发、往不同方向；视野 14 m），有 Boss；v10 透视
    "shrine_R10_s5510": dict(mode="union", run="R10_s5510", scene_src="05_Japanese_Spring_Sanctuary",
                             starts=[3840, 3936, 4128], theme="graveyard", seed=5, f0_scene="greek@photoreal"),
    # 帧 0 视野重叠的 3 条（视野 23 m），Boss 可见率 0.71
    "temple_R10_s6507": dict(mode="union", run="R10_s6507", scene_src="06_Overgrown_Tropical_Water_Temple",
                             starts=[2304, 2688, 2784], theme="maze", seed=7, f0_scene="tibetan@photoreal"),
    # 同一 run 的 4 条，往 4 个方向（−123° / 60° / 148° / −20°；视野 27 m），无 Boss
    "court_R10_s1512": dict(mode="union", run="R10_s1512", scene_src="01_Sunlit_Mediterranean_Courtyard",
                            starts=[1824, 2112, 2208, 2400], theme="orchard", seed=3, f0_scene="farm@photoreal"),
    # 同一组 clip 平移到同一起点 ⇒ 一张首帧、4 种按键
    "court_R10_s1512_start": dict(mode="start", run="R10_s1512", scene_src="01_Sunlit_Mediterranean_Courtyard",
                                  starts=[1824, 2112, 2208, 2400], theme="orchard", seed=3, f0_scene="farm@photoreal"),
    # R10_s6507 的 4 条，往 4 个方向（159° / −16° / −121° / 61°；视野 24 m），平移到同一起点
    "temple_R10_s6507_start": dict(mode="start", run="R10_s6507", scene_src="06_Overgrown_Tropical_Water_Temple",
                                   starts=[2304, 2496, 2880, 2976], theme="maze", seed=7, f0_scene="tibetan@photoreal"),
}


def clip_id(g: dict, start: int) -> str:
    return f"{g['scene_src']}__{g['run']}__{start:06d}"


def translate_meta(meta: dict, t: np.ndarray) -> dict:
    """世界整体平移 t（cm，xyz）：相机矩阵 M' = T(−t)·M（行向量 [X,1]·M），胶囊端点 / actor_pos / ground_z 加 t。"""
    T = np.eye(4)
    T[3, :3] = -t
    out = dict(meta)
    for k in ("world_to_clip", "world_to_view"):
        m = np.asarray(meta[k], np.float64).reshape(-1, 4, 4)
        out[k] = np.einsum("ij,fjk->fik", T, m).reshape(-1, 16)
    caps = np.array(meta["capsules"], np.float64)
    caps[..., 0:3] += t
    caps[..., 3:6] += t
    out["capsules"] = caps.astype(np.float32)
    out["actor_pos"] = (np.asarray(meta["actor_pos"], np.float64) + t).astype(np.float32)
    out["ground_z"] = (np.asarray(meta["ground_z"], np.float64) + t[2]).astype(np.float32)
    return out


def drop_boss(meta: dict) -> dict:
    out = dict(meta)
    caps = np.array(meta["capsules"], np.float32)
    caps[:, 1] = np.nan
    out["capsules"] = caps
    ap = np.array(meta["actor_pos"], np.float32)
    ap[:, 1] = np.nan
    out["actor_pos"] = ap
    return out


def plan_image(w, srcs: list[dict], dev, max_px: int = 1100):
    """俯视总图（procmap.plan 的画法）+ 每条 clip 的轨迹（各一色）+ 帧 0 视野框；返回 (BGR 图, 世界→图像素的函数)。"""
    import cv2
    import procmap as PM
    res = max(w.nx, w.ny) * PM.GRID_CM / max_px
    nx, ny = int(w.nx * PM.GRID_CM / res), int(w.ny * PM.GRID_CM / res)
    yy, xx = np.mgrid[0:ny, 0:nx]
    xy = np.stack([w.x0 + (xx.ravel() + 0.5) * res, w.y0 + (yy.ravel() + 0.5) * res], 1)
    ro = np.c_[xy, np.full(len(xy), w.z0 + 5000.0)]
    rd = np.tile(np.array([3e-4, 2e-4, -1.0]) / np.linalg.norm([3e-4, 2e-4, -1.0]), (len(xy), 1))
    prims = w.prims()
    t, k = PM.Caster(prims, dev).cast(ro, rd)
    cls = w.ground.classify(xy, dev)
    hit = k >= 0
    cls[hit] = [prims[i]["cls"] for i in k[hit]]
    hgt = np.where(hit, 5000.0 - t, 0)
    import geometry as GEO
    img = GEO.PALETTE[cls].astype(np.float32) * (0.75 + 0.25 * np.clip(hgt / 800, 0, 1))[:, None]
    img = img.reshape(ny, nx, 3).clip(0, 255).astype(np.uint8)[..., ::-1].copy()
    sc = PM.GRID_CM / res

    def q(P):
        return ((w.px(np.asarray(P, float)) + 0.5) * sc).astype(np.float32)
    return img, q, res


def build_group(name: str, g: dict, dev: str = "cuda") -> dict:
    import cv2
    import procmap as PM
    import themes_ext  # noqa: F401  注册主题
    import geometry as GEO
    import model as GM
    t0 = time.time()
    gd = f"{S}/maps/{name}"
    os.makedirs(gd, exist_ok=True)
    srcs = [PM.load_source(clip_id(g, s)) for s in g["starts"]]
    if g["mode"] == "start":
        m0 = srcs[0]["meta"]
        p0 = np.r_[m0["actor_pos"][0, 0, :2], m0["ground_z"][0]].astype(np.float64)
        for s in srcs:
            m = s["meta"]
            t = p0 - np.r_[m["actor_pos"][0, 0, :2], m["ground_z"][0]].astype(np.float64)
            s["shift_cm"] = t.tolist()
            s["meta"] = drop_boss(translate_meta(m, t))
        d = np.abs(np.asarray(srcs[1]["meta"]["world_to_clip"][0]) - np.asarray(m0["world_to_clip"][0])).max()
        print(f"[{name}] 平移后帧 0 相机矩阵最大差 {d:.3g}", flush=True)
    keys = ("world_to_clip", "world_to_view", "capsules", "actor_pos", "ground_z")
    merged = dict(srcs[0])
    merged["meta"] = {k: np.concatenate([np.asarray(s["meta"][k]) for s in srcs], 0) for k in keys}
    merged["meta"]["ortho"] = srcs[0]["meta"]["ortho"]
    n_tot = F * len(srcs)
    PM.F = n_tot                                     # World 的足迹 / 视野 + 遮挡控制按全部 clip 的全部帧
    try:
        w = PM.World(merged, g["seed"])
        PM.THEME_FN[g["theme"]](w)
        cache = [PM.char_rays(merged["meta"], f) for f in range(n_tot)]
        occ = PM.control_occlusion(w, dev, cache)
    finally:
        PM.F = F
    print(f"[{name}] 世界 {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)
    plan, q, res = plan_image(w, srcs, dev)
    rows = []
    for k, (s, start) in enumerate(zip(srcs, g["starts"])):
        w.src = s
        ids, keep = PM.render(w, dev)
        assert set(np.unique(ids).tolist()) <= GM.valid_raw_ids("parts")
        od = f"{gd}/{k}_{start:06d}"
        os.makedirs(od, exist_ok=True)
        coll = GM.collapse_parts(ids)
        ppx = [int((coll[f] == GEO.PLAYER_ID).sum()) for f in range(F)]
        np.savez_compressed(f"{od}/ids.npz", ids=ids)
        json.dump(s["prompts"], open(f"{od}/prompts.json", "w"), indent=1)
        np.save(f"{od}/actions.npy", s["blocks"][1:])
        row = s["row"]
        json.dump(dict(clip_id=s["clip_id"], subset=row["subset"], scene=row["scene"], run=row["run"], start=row["start"],
                       theme=g["theme"], theme_desc=PM.THEME_DESC[g["theme"]], seed=g["seed"], ortho=s["meta"]["ortho"],
                       z_ground_cm=w.z0, group=name, group_mode=g["mode"], group_clips=[clip_id(g, x) for x in g["starts"]],
                       shift_cm=s.get("shift_cm"), boss_dropped=g["mode"] == "start", actions_block=s["blocks"].tolist(),
                       world_origin_cm=[float(w.x0), float(w.y0)], generator="scripts/teaser/edit_sameworld.py"),
                  open(f"{od}/source.json", "w"), indent=1)
        rep = dict(case=f"{name}/{k}_{start:06d}", clip_id=s["clip_id"], player_px_f0=ppx[0], player_px_min=min(ppx),
                   boss_px_f0=int((coll[0] == GEO.BOSS_ID).sum()), n_prims=len(w.prims()), occlusion_control=occ,
                   frame0={c: round(float((coll[0] == i).mean()), 4) for c, i in
                           (("terrain", 1), ("paving", 2), ("architecture", 3), ("rock", 4), ("vegetation", 5),
                            ("prop", 6), ("sky", 0))}, checks_ok=bool(ppx[0] >= 500), 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{f}") for f in (0, 20, 40, 60, 80)]
        cv2.imwrite(f"{od}/contact.jpg", np.concatenate(tiles, 1)[..., ::-1], [cv2.IMWRITE_JPEG_QUALITY, 90])
        P = s["meta"]["actor_pos"][:F, 0, :2]
        P = P[np.isfinite(P).all(1)]
        foot0 = PM.frame_ground_poly(s["meta"]["world_to_clip"][0], s["meta"]["world_to_view"][0], w.z0)
        rows.append(dict(k=k, start=start, clip_id=s["clip_id"], dir=od, path_px=q(P).tolist(), foot0_px=q(foot0).tolist(),
                         disp_cm=(P[-1] - P[0]).tolist(), player_px_f0=ppx[0], shift_cm=s.get("shift_cm")))
        print(f"[{name}] clip {k} {s['clip_id']}: 帧 0 玩家 {ppx[0]} px · 最少 {min(ppx)} px", flush=True)
    for r in rows:                                    # plan：帧 0 视野框 + 轨迹
        col = COLORS[r["k"] % len(COLORS)][::-1]
        cv2.polylines(plan, [np.asarray(r["foot0_px"], np.float32).round().astype(np.int32)], True, col, 1, cv2.LINE_AA)
    for r in rows:
        col = COLORS[r["k"] % len(COLORS)][::-1]
        pp = np.asarray(r["path_px"], np.float32).round().astype(np.int32)
        cv2.polylines(plan, [pp], False, (0, 0, 0), 5, cv2.LINE_AA)
        cv2.polylines(plan, [pp], False, col, 3, cv2.LINE_AA)
        cv2.circle(plan, tuple(pp[0]), 6, col, -1, cv2.LINE_AA)
        cv2.circle(plan, tuple(pp[0]), 6, (0, 0, 0), 1, cv2.LINE_AA)
    cv2.imwrite(f"{gd}/plan.png", plan)
    info = dict(name=name, **g, plan=f"{gd}/plan.png", plan_res_cm_per_px=res, clips=rows,
                world_extent_m=[round(w.nx * PM.GRID_CM / 100, 1), round(w.ny * PM.GRID_CM / 100, 1)],
                n_prims=len(w.prims()), occlusion_control=occ, sec=round(time.time() - t0, 1))
    json.dump(info, open(f"{gd}/group.json", "w"), indent=1)
    return info


def cmd_build(a) -> None:
    for name in a.groups or list(GROUPS):
        build_group(name, GROUPS[name])


def cmd_f0jobs(a) -> None:
    os.makedirs(f"{S}/f0", exist_ok=True)
    rows = []
    for name, g in GROUPS.items():
        gj = f"{S}/maps/{name}/group.json"
        if not os.path.exists(gj):
            continue
        info = json.load(open(gj))
        clips = info["clips"] if g["mode"] == "union" else info["clips"][:1]
        for c in clips:
            for seed in a.seeds:
                rows.append({"case": f"{name}__{c['k']}", "map_dir": c["dir"], "scene": g["f0_scene"], "seed": seed,
                             "out": f"{S}/f0/{name}/{c['k']}_{c['start']:06d}/qwen21_s{seed}.png"})
    with open(f"{S}/f0/jobs.jsonl", "w") as fh:
        for r in rows:
            fh.write(json.dumps(r) + "\n")
    print(f"{len(rows)} 个首帧任务 → {S}/f0/jobs.jsonl")


START_LOOKS = ("night", "snow")      # start 组的锚首帧再做外观编辑（edit_app_f0 整图编辑）⇒ 同一开局 × 外观 × 按键


def cmd_f0edits(a) -> None:
    """start 组锚 clip 的首帧 → edit_app_f0 的整图编辑任务（night / snow），写 $S/f0/edits.jsonl。"""
    import edit_app as E
    rows = []
    for name, g in GROUPS.items():
        gj = f"{S}/maps/{name}/group.json"
        if g["mode"] != "start" or not os.path.exists(gj):
            continue
        c = json.load(open(gj))["clips"][0]
        src = f"{S}/f0/{name}/{c['k']}_{c['start']:06d}/qwen21_s{a.seed}.png"
        if not os.path.exists(src):
            continue
        for v in START_LOOKS:
            rows.append({"out": src[:-4] + f"__{v}.png", "src": src, "ids": f"{c['dir']}/ids.npz", "kind": "global",
                         "prompt": E.GLOBAL[v], "seed": E.SEED[v], "variant": v, "group": name})
    with open(f"{S}/f0/edits.jsonl", "w") as fh:
        for r in rows:
            fh.write(json.dumps(r) + "\n")
    print(f"{len(rows)} 个锚首帧外观编辑任务 → {S}/f0/edits.jsonl")


def cmd_rolljobs(a) -> None:
    rows = []
    for name, g in GROUPS.items():
        gj = f"{S}/maps/{name}/group.json"
        if not os.path.exists(gj):
            continue
        info = json.load(open(gj))
        anchor = info["clips"][0]
        looks = [None] + (list(START_LOOKS) if g["mode"] == "start" else [])
        for c, look in [(c, lk) for lk in looks for c in info["clips"]]:
            fc = c if g["mode"] == "union" else anchor
            first = f"{S}/f0/{name}/{fc['k']}_{fc['start']:06d}/qwen21_s{a.seed}" + (f"__{look}" if look else "") + ".png"
            if not os.path.exists(first):
                continue
            ed = {"type": "sameworld", "variant": g["mode"], "group": name, "group_mode": g["mode"], "clip_index": c["k"],
                  "source_clip": c["clip_id"], "group_clips": [x["clip_id"] for x in info["clips"]], "theme": g["theme"],
                  "world_seed": g["seed"], "f0_scene": g["f0_scene"], "f0_seed": a.seed, "f0_backend": "qwen21 shp_afix_rs",
                  "shift_cm": c.get("shift_cm"), "boss_dropped": g["mode"] == "start", "edit_frame": None,
                  "map_dir": c["dir"], "actions": f"{c['dir']}/actions.npy", "plan": info["plan"],
                  "look_edit": look, "label": f"Play {c['k'] + 1}" + (f" ({look})" if look else ""),
                  "scene": g["f0_scene"], "map": g["theme"],
                  "desc": ("same procmap world (union of the clips' footprints), each clip's own camera + capsules + keys"
                           if g["mode"] == "union" else
                           "clips translated to one common start (boss removed): same first frame, different keys")}
            if look:
                ed["first_edit"] = {"editor": "Qwen-Image-2.1-Turbo", "kind": "global", "variant": look,
                                    "src": first.replace(f"__{look}.png", ".png")}
            rows.append({"name": f"{name}__{c['k']}_{c['start']:06d}" + (f"__{look}" if look else ""), "first": first, "ids": f"{c['dir']}/ids.npz",
                         "prompts": f"{c['dir']}/prompts.json", "out": f"{S}/roll", "seed": 0, "prio": 1,
                         "scene": g["f0_scene"], "map": g["theme"], "label": ed["label"], "edit": ed})
    with open(f"{S}/jobs.jsonl", "w") as fh:
        for r in rows:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f"{len(rows)} 个 rollout 任务 → {S}/jobs.jsonl")


def cmd_compose(a) -> None:
    import edit_compose as EC
    for name in a.groups or list(GROUPS):
        if os.path.exists(f"{S}/maps/{name}/group.json"):
            EC.compose_sameworld(name, a)
            if GROUPS[name]["mode"] == "start":
                EC.compose_sameworld_looks(name, START_LOOKS)


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("--groups", nargs="*", default=None)
    f = sp.add_parser("f0jobs")
    f.add_argument("--seeds", type=int, nargs="+", default=[0])
    e = sp.add_parser("f0edits")
    e.add_argument("--seed", type=int, default=0)
    r = sp.add_parser("rolljobs")
    r.add_argument("--seed", type=int, default=0, help="用哪个首帧种子")
    c = sp.add_parser("compose")
    c.add_argument("--groups", nargs="*", default=None)
    a = ap.parse_args(argv)
    {"build": cmd_build, "f0jobs": cmd_f0jobs, "f0edits": cmd_f0edits, "rolljobs": cmd_rolljobs,
     "compose": cmd_compose}[a.cmd](a)


if __name__ == "__main__":
    main()