#!/usr/bin/env python3
"""推理：首帧 + 场景从句 + 动作脚本 ⇒ 81 帧视频，并对生成结果实测背景位移。

动作在 ActionRoPE 里不走文本，走 RoPE 的世界坐标，所以推理时的"指令"就是一条
逐帧相机偏移 frame_offset_px[81,2]（玩家屏幕投影位移，与 sidecar 同一约定），
再按 cell 取均值得到 offset_px[21,2] 交给 arope_forward。两种来源：

    --actions "right:1.0:21"               方向 × 倍率 × cell 数，速度取 sidecar meta 的 walk_speed_px_s
    --actions "right:1.0:10,left:1.0:11"   分段脚本，cell 数合计必须是 21
    --actions "71,0:21"                    直接给 (dx, dy) px/s
    --replay <clip_id>                     sidecar 真值 + 该 clip 的首帧 + meta 里的原 prompt，可与 GT 比 PSNR

方向名指的是**玩家在屏幕上的真实运动方向**（right ⇒ dx>0 ⇒ 背景左移）。数据集的动作
标签是镜像的（label 7 "moving right" 的 clip 玩家实际向屏幕左走，2026-09-04 人工看视频确认），
plain 臂训练默认已把动作词翻回真实方向（train.py --mirror_labels，查 text_table_eybx_mirror.pt），
所以推理默认 --label_mode literal 按字面写真实方向，--replay 的原 prompt 也同样翻转；
--label_mode mirror 只给 --no-mirror_labels 训出来的旧约定模型用（文本 = 数据集原标签、原表）。

采样照抄 WanVideoPipeline.__call__ 的 TI2V 路径：FlowMatchScheduler("Wan") 50 步 shift 5，
帧 0 的 latent 每步都写回 VAE 编码的首帧（WanVideoUnit_ImageEmbedderFused 的做法），
DiT 前向用 arope_forward（plain 臂 offset_px=None，就是原模型）。

baseline 臂（linear / xattn / prompt / adaln，baseline/SPEC.md）：--arm 可不给，按 ckpt 的 metadata / `arm.` 键集
自动识别并 strict 装上（baseline.detect_arm_from_ckpt）；动作脚本 → action_inputs（offset_px 同 ARoPE，
offset_tok / delta_tok 由它派生，action_idx = velocity_to_label(mirror=False)，replay 时用 meta 标签翻回真实方向），
文本与 arope 臂一样只剩场景（prompt 臂与 plain 一样带动作从句）。

VAE 尺度：数据集 latent 是官方 Wan2_2_VAE.encode 的输出（已做 mean/std 归一化），
DiffSynth 的 WanVideoVAE38.encode 走同一套 scale，non-tiled 时逐 cell 一致；本脚本
首帧编码 / 视频解码都用 tiled=False，与数据集对齐（tests/test_infer.py 实测过误差）。

实测位移：用 actionrope/verify_sidecar.py 的 measure_shift（SIFT+RANSAC 仿射场在玩家
位置取值）量帧 0 → 帧 80 的背景位移；背景位移 = −frame_offset_px，所以与
−frame_offset_px[80] 对比。直接 0→80 测不出时（重叠区小），退化为每 16 帧链式累加。

    CUDA_VISIBLE_DEVICES=4 python -m actionrope.infer --replay clip_Eybx_200000958_000273 \
        --steps 20 --out outputs/samples/smoke_base_replay.mp4
    CUDA_VISIBLE_DEVICES=4 python -m actionrope.infer --first_frame clip_Eybx_200000958_000273 \
        --scene tidal_flats --actions right:1.0:21 --out outputs/samples/right_x1.mp4
"""
from __future__ import annotations

import argparse
import glob
import json
import os
import sys
import time

import numpy as np
import torch

os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True")
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from actionrope import geometry as G  # noqa: E402
from actionrope.arope import arope_forward, install_arope, loss_weight_map, MASK_EMBEDDING_NAME  # noqa: E402
from actionrope.prompts import MIRROR_ACTION, NULL_PROMPT, mirror_action, strip_action  # noqa: E402
from baseline import ARM_NAMES, TEXT_ARMS, build_arm, detect_arm_from_keys, make_action_inputs, parse_arm_kwargs, split_arm_keys  # noqa: E402

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DATA = os.path.join(ROOT, "data")
MODEL_DIR = os.path.join(ROOT, "models", "Wan2.2-TI2V-5B")
DIT_FILES = [os.path.join(MODEL_DIR, f"diffusion_pytorch_model-0000{i}-of-00003.safetensors") for i in (1, 2, 3)]
VAE_PTH = os.path.join(MODEL_DIR, "Wan2.2_VAE.pth")
T5_PTH = os.path.join(MODEL_DIR, "models_t5_umt5-xxl-enc-bf16.pth")
T5_TOK = os.path.join(MODEL_DIR, "google", "umt5-xxl")
SIDECAR = os.path.join(DATA, "latent", "arope_sidecar.pt")
TEXT_TABLE = {"arope": os.path.join(DATA, "latent", "text_table_actionrope.pt"),
              "plain": os.path.join(DATA, "latent", "text_table_eybx_mirror.pt"),   # 含翻回真实方向的全部串
              "plain_raw": os.path.join(DATA, "latent", "text_table_eybx.pt")}      # 旧约定（--label_mode mirror）

sys.path.insert(0, os.path.join(DATA, "code"))
from scenes import ACTIONS, SCENE_NL, PROMPT_TMPL  # noqa: E402

# 转场 trigger 从句：与 data/code/burn_corpus.py 的 TRIG_* 同构。目的地要去掉前导的 "in "，
# 否则拼出 "burning away into in a dark cave"。
TRIG_TMPL = {"burn": "{a} burning away into {b}", "cut": "{a} cutting to {b}"}
# 场景 dropout 的裸串（data/code/encode_latents.py 的 BARE_TMPL，.pt 里的 prompts_bossdrop 就是它）：
# plain 臂训练时 dropout 用的是这 9 个串而不是 NULL_PROMPT（NULL_PROMPT 不在 eybx 表里），
# 所以 plain 臂 CFG 的负条件也用它——负条件必须是模型训练中见过的"无场景"串。
BARE_TMPL = "In a 2.5D top down view, the player is {action}."

# 动作脚本里的方向名 → scenes.ACTIONS 下标（方向向量取 geometry.ACTION_DIRS）
DIR_NAMES = {
    "idle": 0, "up": 1, "up-left": 2, "left": 3, "down-left": 4,
    "down": 5, "down-right": 6, "right": 7, "up-right": 8,
    "ul": 2, "dl": 4, "dr": 6, "ur": 8, "u": 1, "d": 5, "l": 3, "r": 7,
}


# ------------------------------------------------------------------ 动作脚本 ----
def parse_actions(spec: str, walk_speed: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """动作脚本 → (frame_offset_px [81,2], cell_velocity [21,2] px/s)。

    每段 "<dir>:<mult>:<n_cells>" 或 "<dx>,<dy>:<n_cells>"（px/s）。速度按轴分别取
    walk_speed_px_s × mult × 方向单位向量：横纵步速本来就不同（透视），单位向量只定方向。
    帧 i 属于 cell(i)，frame_offset[i] = frame_offset[i−1] + v[cell(i)] / 16，帧 0 恒为 0，
    常速时正是 SPEC 的 i/16 × speed × 方向。
    """
    # 段之间用逗号分隔；"dx,dy:cells" 里的逗号会把 dx 切成单独一个 token（没有冒号），
    # 遇到就与下一个 token 合回去。
    toks = [t.strip() for t in spec.split(",") if t.strip()]
    segs, i = [], 0
    while i < len(toks):
        if ":" not in toks[i]:
            if i + 1 >= len(toks):
                raise ValueError(f"动作段不完整: {toks[i]!r}")
            segs.append(toks[i] + ";" + toks[i + 1])
            i += 2
        else:
            segs.append(toks[i])
            i += 1
    vel = []
    for seg in segs:
        parts = seg.split(":")
        if len(parts) == 3:
            name, mult, n = parts[0].strip().lower(), float(parts[1]), int(parts[2])
            if name not in DIR_NAMES:
                raise ValueError(f"未知方向 {name!r}，可选 {sorted(DIR_NAMES)}")
            v = G.ACTION_DIRS[DIR_NAMES[name]] * walk_speed * mult
        elif len(parts) == 2 and ";" in parts[0]:
            dx, dy = (float(t) for t in parts[0].split(";"))
            v, n = np.array([dx, dy], np.float64), int(parts[1])
        else:
            raise ValueError(f"动作段格式应为 dir:mult:cells 或 dx,dy:cells，得到 {seg!r}")
        vel += [v] * n
    if len(vel) != G.CELLS:
        raise ValueError(f"动作脚本的 cell 数合计应为 {G.CELLS}，得到 {len(vel)}")
    vel = np.stack(vel, 0)
    frame_off = np.zeros((G.FRAMES, 2), np.float64)
    cell_of = np.zeros(G.FRAMES, int)
    for k in range(G.CELLS):
        cell_of[G.cell_frames(k)] = k
    for i in range(1, G.FRAMES):
        frame_off[i] = frame_off[i - 1] + vel[cell_of[i]] / G.FPS
    return frame_off, vel


def velocity_to_label(v: np.ndarray, mirror: bool, speed_eps: float = 1.0) -> int:
    """cell 速度 (px/s) → 动作标签下标。mirror=True 时按数据集的镜像约定（标签方向 = −真实屏幕方向）。"""
    if np.hypot(*v) < speed_eps:
        return 0
    d = -v if mirror else v
    cos = G.ACTION_DIRS[1:] @ (d / np.linalg.norm(d))
    return int(np.argmax(cos)) + 1


# ------------------------------------------------------------------ 文本 ----
def scene_clause(scene: str) -> str:
    """--scene 接受 scene key（tidal_flats）或完整从句（in a pale tidal flat ...）。"""
    if scene in SCENE_NL:
        return SCENE_NL[scene]
    if scene.lower().startswith("in a"):
        return scene
    raise ValueError(f"--scene 应为 scene key {sorted(SCENE_NL)} 或 'in a ...' 从句，得到 {scene!r}")


def compose_cell_prompts(scene: str, trigger: str | None, labels: list[int]) -> list[str]:
    """按数据集模板拼 21 条完整 prompt（含动作从句）；arope 臂再 strip_action。

    trigger "burn:<B>:<start>:<n>"：cell [start, start+n) 用 trigger 从句，之后的 cell
    切成纯 B 场景（语料就是"烧完切回纯 B"）。
    """
    clauses = [scene_clause(scene)] * G.CELLS
    if trigger:
        kind, b, start, n = trigger.split(":")
        start, n = int(start), int(n)
        if kind not in TRIG_TMPL:
            raise ValueError(f"trigger 类型应为 burn/cut，得到 {kind!r}")
        if not (0 <= start < G.CELLS and n >= 1):
            raise ValueError(f"trigger 的 start/n 越界: {start}/{n}")
        a, b_nl = clauses[0], scene_clause(b)
        trig = TRIG_TMPL[kind].format(a=a, b=b_nl[3:] if b_nl.startswith("in ") else b_nl)
        for k in range(start, G.CELLS):
            clauses[k] = trig if k < start + n else b_nl
    return [PROMPT_TMPL.format(scene=c, action=ACTIONS[l]) for c, l in zip(clauses, labels)]


class TextEncoder:
    """查文本表；查不到的串回退加载 UMT5（与 build_text_table.py 同一路径）。"""

    def __init__(self, table_path: str, device: str):
        self.table = torch.load(table_path, map_location="cpu", weights_only=False)
        self.device = device
        self.pipe = None
        self.n_fallback = 0

    def _t5(self, prompts: list[str]) -> list[torch.Tensor]:
        if self.pipe is None:
            print(f"   文本表里没有的串，加载 UMT5 …", flush=True)
            from diffsynth.pipelines.wan_video import ModelConfig, WanVideoPipeline
            self.pipe = WanVideoPipeline.from_pretrained(
                torch_dtype=torch.bfloat16, device=self.device,
                model_configs=[ModelConfig(path=T5_PTH)],
                tokenizer_config=ModelConfig(path=T5_TOK), redirect_common_files=False)
            self.pipe.load_models_to_device(["text_encoder"])
        ids, mask = self.pipe.tokenizer(prompts, return_mask=True, add_special_tokens=True)
        ids, mask = ids.to(self.device), mask.to(self.device)
        seq_lens = mask.gt(0).sum(dim=1).long()
        with torch.no_grad():
            emb = self.pipe.text_encoder(ids, mask)
        for j, v in enumerate(seq_lens):
            emb[j, v:] = 0
        return [e.to(torch.bfloat16).cpu().clone() for e in emb]

    def encode(self, prompts: list[str]) -> torch.Tensor:
        """list[str] → bf16 [n, 512, 4096]。"""
        missing = sorted({p for p in prompts if p not in self.table})
        if missing:
            self.n_fallback += len(missing)
            for p, e in zip(missing, self._t5(missing)):
                self.table[p] = e
        return torch.stack([self.table[p] for p in prompts], 0)


def build_context(enc: TextEncoder, prompts: list[str], device: str, per_cell: bool = False):
    """21 条 cell prompt → (context, context_ids)。全同 ⇒ 整条 [1,512,4096]（与逐 cell 等价、更省）；
    per_cell=True（prompt 臂：PromptArm 核对 context 必须逐 cell）时全同也展成 [1,21,512,4096]（ids 全 0，只编码一次）。"""
    uniq = sorted(set(prompts))
    if len(uniq) == 1 and not per_cell:
        return enc.encode(uniq).to(device), None
    emb = enc.encode(uniq)
    idx = torch.tensor([uniq.index(p) for p in prompts])
    return emb[idx].unsqueeze(0).to(device), idx.unsqueeze(0).to(device)


# ------------------------------------------------------------------ 数据 ----
def find_clip_mp4(clip_id: str) -> str:
    hits = glob.glob(os.path.join(DATA, "clips", "*", clip_id + ".mp4"))
    if not hits:
        raise FileNotFoundError(f"找不到 {clip_id} 的 mp4（data/clips/*/）")
    return hits[0]


def read_frames(path: str, n: int | None = None) -> np.ndarray:
    """mp4 → uint8 [F, H, W, 3]（与 encode_latents.py 的 ffmpeg rgb24 解码一致）。"""
    import av
    out = []
    with av.open(path) as c:
        for f in c.decode(video=0):
            out.append(f.to_ndarray(format="rgb24"))
            if n is not None and len(out) >= n:
                break
    return np.stack(out, 0)


def load_first_frame(spec: str) -> np.ndarray:
    """png / mp4 / clip_id → uint8 [H, W, 3]，统一缩放到 832×480。"""
    from PIL import Image
    if spec.endswith(".mp4"):
        img = read_frames(spec, 1)[0]
    elif os.path.isfile(spec):
        img = np.array(Image.open(spec).convert("RGB"))
    else:
        img = read_frames(find_clip_mp4(spec), 1)[0]
    if img.shape[:2] != (G.H, G.W):
        # 先按 26:15 居中裁剪再缩放：bg50 这类 3:2 的图直接拉伸会横向压扁 13%，步速尺度就错了
        h, w = img.shape[:2]
        tw, th = (w, round(w * G.H / G.W)) if w * G.H <= h * G.W else (round(h * G.W / G.H), h)
        x0, y0 = (w - tw) // 2, (h - th) // 2
        img = img[y0:y0 + th, x0:x0 + tw]
        img = np.array(Image.fromarray(img).resize((G.W, G.H), Image.BICUBIC))
    return img


def load_meta_row(clip_id: str) -> dict:
    """--replay 需要 clip 的原 prompt 与动作标签；meta 是文本，比 torch.load .pt 省事。"""
    for f in sorted(glob.glob(os.path.join(DATA, "meta", "*.jsonl"))):
        with open(f, encoding="utf-8") as fh:
            for line in fh:
                if clip_id in line:
                    r = json.loads(line)
                    if r["clip_id"] == clip_id:
                        return r
    raise KeyError(f"meta 里没有 {clip_id}")


# ------------------------------------------------------------------ 模型 ----
def resolve_arm(ckpt: str | None, arm: str | None) -> tuple[str, dict]:
    """(--ckpt, --arm) → (臂名, 臂构造 kwargs)。

    有 ckpt：先看 safetensors metadata 的 arm，再按 `arm.` 键集反推（键集能区分 linear / xattn / adaln 并给出 xattn 的
    window_frames / enable_mouse 等结构参数；无 arm 键的 ckpt 是 arope / plain / prompt，靠 metadata 区分）。
    --arm 给了就必须与 ckpt 一致（plain 与 prompt 前向相同，允许互换）。没 ckpt 且没 --arm ⇒ arope（基座冒烟）。
    只读 header / 形状，不加载张量。"""
    if not ckpt:
        return arm or "arope", {}
    from safetensors import safe_open
    with safe_open(ckpt, "pt") as f:
        meta = f.metadata() or {}
        keys = list(f.keys())
        arm_keys = [k[len("arm."):] for k in keys if k.startswith("arm.")]
        shapes = {k: tuple(f.get_slice("arm." + k).get_shape()) for k in arm_keys}
    name_keys, kw = detect_arm_from_keys(arm_keys, shapes)
    name_meta = meta.get("arm")
    if name_meta and name_keys and name_meta != name_keys:
        raise ValueError(f"ckpt metadata 说 arm={name_meta!r}，键集却是 {name_keys!r}")
    if meta.get("arm_kwargs"):
        try:
            kw = {**parse_arm_kwargs(meta["arm_kwargs"]), **kw}
        except (ValueError, json.JSONDecodeError):
            pass
    name = name_meta or name_keys or "arope"
    if arm and arm != name and {arm, name} != {"plain", "prompt"}:
        raise ValueError(f"--arm {arm} 与 ckpt 的臂 {name!r}（metadata/键集）不符")
    return arm or name, kw


def load_models(ckpt: str | None, mask_channel: bool | None, device: str, need_vae: bool = True,
                arm: str | None = None, arm_kwargs: dict | str | None = None):
    """DiT(bf16) + VAE（+ baseline 动作臂）；--ckpt 时 strict 加载，mask_channel 由 ckpt 键集自动判断，
    臂名不给时按 ckpt 自动识别。返回 (pipe, dit, arm_module)，arope / plain 的 arm_module 为 None。"""
    from diffsynth.pipelines.wan_video import ModelConfig, WanVideoPipeline
    arm_name, kw = resolve_arm(ckpt, arm)
    if arm_kwargs:
        kw = {**kw, **parse_arm_kwargs(arm_kwargs)}
    cfgs = [ModelConfig(path=DIT_FILES)] + ([ModelConfig(path=VAE_PTH)] if need_vae else [])
    pipe = WanVideoPipeline.from_pretrained(torch_dtype=torch.bfloat16, device=device,
                                            model_configs=cfgs, tokenizer_config=None,
                                            redirect_common_files=False)
    dit = pipe.dit.eval().requires_grad_(False)
    if ckpt:
        from safetensors.torch import load_file
        dit_sd, arm_sd = split_arm_keys(load_file(ckpt))
        has_mask = any(k.startswith(MASK_EMBEDDING_NAME + ".") for k in dit_sd)
        if mask_channel is not None and mask_channel != has_mask:
            raise ValueError(f"--mask_channel={mask_channel} 与 ckpt 键集不符（含 {MASK_EMBEDDING_NAME}: {has_mask}）")
        install_arope(dit, mask_channel=has_mask)
        dit.load_state_dict(dit_sd, strict=True)
        arm_module = build_arm(arm_name, dit, kw)
        if arm_module is not None:
            arm_module.load_state_dict(arm_sd, strict=True)  # 键集 / 形状不符在这里炸
            arm_module.eval().requires_grad_(False)
        elif arm_sd:
            raise ValueError(f"ckpt 含 {len(arm_sd)} 个 arm.* 键，但臂 {arm_name} 没有参数")
        print(f"   ckpt {ckpt}: {len(dit_sd)} 个 DiT 键 + {len(arm_sd)} 个 arm 键 · mask_channel={has_mask} · arm={arm_name}"
              + (f"（{type(arm_module).__name__}，{arm_module.n_new_params():,} 参数，kwargs={kw}）" if arm_module is not None else ""))
    else:
        install_arope(dit, mask_channel=bool(mask_channel))
        arm_module = build_arm(arm_name, dit, kw)  # 基座权重 + 零初始化的臂：输出与原模型相同（冒烟用）
        if arm_module is not None:
            arm_module.eval().requires_grad_(False)
    return pipe, dit, arm_module


class _vae_fp32:
    """把 VAE 临时切到 fp32 编码（数据集 latent 是 fp32 编码后存 bf16 的；bf16 编码的首帧与训练时
    看到的帧 0 latent 差 ~1%，fp32 只剩 bf16 存储舍入的 0.3%）。单帧 0.9 s，用完切回 bf16。"""

    def __init__(self, pipe, enabled: bool):
        self.pipe, self.enabled = pipe, enabled

    def __enter__(self):
        if self.enabled:
            self.pipe.vae.model.float()
        return self

    def __exit__(self, *exc):
        if self.enabled:
            self.pipe.vae.model.to(self.pipe.torch_dtype)


def encode_first_frame(pipe, img: np.ndarray, tiled: bool = False, fp32: bool = True) -> torch.Tensor:
    """uint8 [H,W,3] → latent [1, 48, 1, 30, 52]（WanVideoUnit_ImageEmbedderFused 的做法，默认 fp32 编码）。"""
    from PIL import Image
    dtype = torch.float32 if fp32 else pipe.torch_dtype
    x = pipe.preprocess_image(Image.fromarray(img), torch_dtype=dtype).transpose(0, 1)  # [C, 1, H, W]
    with torch.no_grad(), _vae_fp32(pipe, fp32):
        z = pipe.vae.encode([x], device=pipe.device, tiled=tiled, tile_size=(30, 52), tile_stride=(15, 26))
    return z.to(dtype=pipe.torch_dtype, device=pipe.device)


def encode_video(pipe, frames: np.ndarray, tiled: bool = False, fp32: bool = False) -> torch.Tensor:
    """uint8 [F,H,W,3] → latent [1, 48, (F−1)/4+1, 30, 52]，与数据集 encode_latents.py 同一路径（fp32=True 时连精度也一致）。"""
    dtype = torch.float32 if fp32 else pipe.torch_dtype
    x = torch.from_numpy(frames).permute(3, 0, 1, 2).to(dtype).div_(127.5).sub_(1.0)
    with torch.no_grad(), _vae_fp32(pipe, fp32):
        z = pipe.vae.encode([x], device=pipe.device, tiled=tiled, tile_size=(30, 52), tile_stride=(15, 26))
    return z


def decode_video(pipe, latents: torch.Tensor, tiled: bool = False) -> np.ndarray:
    """latent [1,48,21,30,52] → uint8 [81, H, W, 3]。"""
    with torch.no_grad():
        video = pipe.vae.decode(latents, device=pipe.device, tiled=tiled, tile_size=(30, 52), tile_stride=(15, 26))
    frames = pipe.vae_output_to_video(video)
    return np.stack([np.array(f) for f in frames], 0)


@torch.no_grad()
def sample(dit, z0: torch.Tensor, context, context_ids, offset_px, mask_input, *,
           steps: int, shift: float, cfg: float, seed: int, neg_context=None, neg_context_ids=None,
           arm=None, action_inputs=None, log=print) -> torch.Tensor:
    """flow-matching 采样。帧 0 每步写回首帧 latent；CFG 的负条件：arope 臂 NULL_PROMPT，plain 臂逐 cell 裸动作串。
    arm / action_inputs：baseline 动作臂及其输入（正负两支都带动作，CFG 只对文本做）。"""
    from diffsynth.diffusion.flow_match import FlowMatchScheduler
    sched = FlowMatchScheduler("Wan")
    sched.set_timesteps(steps, shift=shift)
    shape = (1, z0.shape[1], G.CELLS, G.LAT_H, G.LAT_W)
    gen = torch.Generator("cpu").manual_seed(seed)
    latents = torch.randn(shape, generator=gen, dtype=torch.float32).to(dtype=z0.dtype, device=z0.device)
    latents[:, :, 0:1] = z0
    kw = dict(offset_px=offset_px, mask_input=mask_input, first_frame_cond=True, arm=arm, action_inputs=action_inputs)
    t0 = time.time()
    for i, t in enumerate(sched.timesteps):
        timestep = t.unsqueeze(0).to(dtype=latents.dtype, device=latents.device)
        pred = arope_forward(dit, latents, timestep, context, context_ids=context_ids, **kw)
        if cfg != 1.0:
            pred_neg = arope_forward(dit, latents, timestep, neg_context, context_ids=neg_context_ids, **kw)
            pred = pred_neg + cfg * (pred - pred_neg)
        latents = sched.step(pred, sched.timesteps[i], latents)
        latents[:, :, 0:1] = z0
        if (i + 1) % 10 == 0 or i + 1 == steps:
            log(f"   step {i + 1}/{steps} · {time.time() - t0:.1f}s")
    return latents


# ------------------------------------------------------------------ 位移实测 ----
def to_gray(frames: np.ndarray) -> np.ndarray:
    import cv2
    return np.stack([cv2.cvtColor(f, cv2.COLOR_RGB2GRAY) for f in frames], 0)


def phase_shift(g0: np.ndarray, gk: np.ndarray) -> tuple[float, float]:
    """全图相位相关（次要参考）。返回背景位移 t：帧 k 位置 s 的内容在帧 0 位于 s − t。"""
    import cv2
    from actionrope.verify_sidecar import _mask
    m = (_mask() > 0).astype(np.float32)
    a = g0.astype(np.float32) * m
    b = gk.astype(np.float32) * m
    win = cv2.createHanningWindow((G.W, G.H), cv2.CV_32F)
    (dx, dy), _ = cv2.phaseCorrelate(a - a.mean(), b - b.mean(), win)
    return float(dx), float(dy)


def _chain_shift(gray: np.ndarray, a: int, b: int, min_step: int) -> np.ndarray | None:
    """a→b 的背景位移；直接测不出就对半分段递归，段短到 min_step 仍测不出 ⇒ None。"""
    from actionrope.verify_sidecar import measure_shift
    m = measure_shift(gray[a], gray[b])
    if m is not None:
        return np.array([m["dx"], m["dy"]])
    if b - a <= min_step:
        return None
    mid = (a + b) // 2
    left, right = _chain_shift(gray, a, mid, min_step), _chain_shift(gray, mid, b, min_step)
    return None if left is None or right is None else left + right


def measure_displacement(gray: np.ndarray, k: int = G.FRAMES - 1, chain: int = 16, min_step: int = 2) -> dict:
    """帧 0 → 帧 k 的累计背景位移，主量 sift（verify_sidecar.measure_shift，玩家位置处的仿射场）。

    直接 0→k 测得出就用它（大基线，没有累计误差）；测不出（重叠区小、生成视频特征少）
    就每 chain 帧一段链式累加，段内再测不出对半细分到 min_step。相位相关只作参考：
    这个透视场景下全图相位相关常锁在静态 HUD / 玩家上给出 0（sidecar 模块的实验结论）。
    """
    out = {"k": k}
    from actionrope.verify_sidecar import measure_shift
    m = measure_shift(gray[0], gray[k])
    out["sift_direct"] = None if m is None else [m["dx"], m["dy"]]
    out["sift_direct_inliers"] = None if m is None else m["n_inl"]
    acc = np.zeros(2)
    for j in range(0, k, chain):
        seg = _chain_shift(gray, j, min(j + chain, k), min_step)
        if seg is None:
            acc = None
            out["sift_chain_failed_at"] = j
            break
        acc += seg
    out["sift_chain"] = None if acc is None else acc.tolist()
    out["sift_chain_step"] = chain
    out["sift_chain_n_segments"] = -(-k // chain)
    out["sift"] = out["sift_direct"] if out["sift_direct"] is not None else out["sift_chain"]
    out["sift_source"] = "direct" if out["sift_direct"] is not None else "chain"
    try:
        out["phase_direct"] = list(phase_shift(gray[0], gray[k]))
        acc = np.zeros(2)
        for j in range(0, k, chain):
            acc += phase_shift(gray[j], gray[min(j + chain, k)])
        out["phase_chain"] = acc.tolist()
        out["phase_note"] = ("全图相位相关在真实画面上恒≈0（静态水印 / 透视锁零），不是第二种实测，"
                             "只在合成平移上验证过符号；以 sift 为准")
    except Exception as e:  # 相位相关只是参考，失败不影响主流程
        out["phase_error"] = str(e)
    return out


def psnr(a: np.ndarray, b: np.ndarray) -> float:
    mse = np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)
    return float("inf") if mse == 0 else float(10 * np.log10(255.0 ** 2 / mse))


def fmt(v) -> str:
    return "None" if v is None else f"({v[0]:+.1f},{v[1]:+.1f})"


# ------------------------------------------------------------------ 主流程 ----
def main(argv=None, models=None):
    """models=(pipe, dit) 可传入已加载的模型（测试里复用，省一次 10 GB 的加载）。"""
    ap = argparse.ArgumentParser(description="ActionRoPE 推理")
    ap.add_argument("--arm", choices=list(ARM_NAMES), default=None,
                    help="不给 ⇒ 按 ckpt 自动识别（无 ckpt 时 arope）；plain 与 prompt 前向相同")
    ap.add_argument("--arm_kwargs", default=None, help="baseline 臂构造参数 JSON（一般不用给，从 ckpt 键集反推）")
    ap.add_argument("--ckpt", default=None, help="step-N.safetensors；不给则用基座权重")
    ap.add_argument("--mask_channel", action=argparse.BooleanOptionalAction, default=None,
                    help="基座权重时是否装 mask 通道（零初始化，不改输出）；有 ckpt 时自动判断")
    ap.add_argument("--first_frame", default=None, help="png / mp4 / clip_id")
    ap.add_argument("--replay", default=None, help="clip_id：sidecar 真值 + 该 clip 首帧 + 原 prompt")
    ap.add_argument("--scene", default=None, help="scene key 或 'in a ...' 从句")
    ap.add_argument("--trigger", default=None, help="burn:<B scene>:<start_cell>:<n_cells>（或 cut:）")
    ap.add_argument("--actions", default=None, help='如 "right:1.0:21" 或 "right:1.0:10,left:1.0:11" 或 "71,0:21"（px/s）')
    ap.add_argument("--label_mode", choices=["literal", "mirror"], default="literal",
                    help="plain / prompt 臂动作从句：literal=真实屏幕方向（配 --mirror_labels 训练，默认）；"
                         "mirror=数据集原镜像标签（配 --no-mirror_labels 训练的旧约定模型）")
    ap.add_argument("--steps", type=int, default=50)
    ap.add_argument("--shift", type=float, default=5.0)
    ap.add_argument("--cfg", type=float, default=1.0)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--tiled", action="store_true", help="VAE 分块（默认不分块，与数据集编码一致）")
    ap.add_argument("--out", required=True, help="输出 mp4 路径（同名 _frame0.png / .json 一并写出）")
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--text_table", default=None)
    args = ap.parse_args(argv)
    t_all = time.time()

    sc = torch.load(SIDECAR, map_location="cpu", weights_only=False)
    walk_speed = np.array(sc["meta"]["walk_speed_px_s"], np.float64)

    # ---- 臂：--arm 或按 ckpt 自动识别（只读 header）；动作走文本的臂 = plain / prompt ----
    args.arm, arm_kwargs = resolve_arm(args.ckpt, args.arm)
    text_arm = args.arm in TEXT_ARMS

    # ---- 指令：frame_offset_px + 21 条 cell prompt ----
    info = dict(arm=args.arm, arm_kwargs=arm_kwargs, ckpt=args.ckpt, steps=args.steps, shift=args.shift, cfg=args.cfg, seed=args.seed)
    if args.replay:
        row = load_meta_row(args.replay)
        c = sc["clips"][args.replay]
        frame_off = c["frame_offset_px"].numpy().astype(np.float64)
        prompts = list(row["prompts"])
        labels = [int(a) for a in row["actions"]]
        # 数据集标签是镜像的：真实屏幕方向 = MIRROR_ACTION[标签]（dataset.py 给 action_idx 用的同一约定）
        true_labels = [MIRROR_ACTION[l] for l in labels]
        if text_arm and args.label_mode == "literal":
            # 数据集原 prompt 是镜像标签；literal 模型学的是真实方向，replay 时翻回来
            prompts = [mirror_action(p) for p in prompts]
            labels = true_labels
        first_spec = args.replay
        info.update(replay=args.replay, sidecar_valid=bool(c["valid"]), kind=c["kind"],
                    actions_gt=row.get("actions"), rt_ratio=row.get("rt_ratio"))
        if not c["valid"]:
            print(f"   注意：{args.replay} 在 sidecar 里 valid=False（{c.get('max_log_dt', float('nan')):.3f}s 日志缺口），偏移仍照用")
    else:
        if not (args.first_frame and args.scene and args.actions):
            ap.error("非 --replay 模式需要 --first_frame、--scene、--actions")
        frame_off, vel = parse_actions(args.actions, walk_speed)
        # 真实屏幕方向的标签（action_idx 用）；文本标签只在旧约定 mirror 模型上取镜像
        true_labels = [velocity_to_label(v, mirror=False) for v in vel]
        labels = [velocity_to_label(v, mirror=args.label_mode == "mirror") for v in vel]
        prompts = compose_cell_prompts(args.scene, args.trigger, labels)
        first_spec = args.first_frame
        info.update(actions=args.actions, scene=args.scene, trigger=args.trigger,
                    labels=labels, label_mode=args.label_mode, walk_speed_px_s=walk_speed.tolist())
    if not text_arm:
        prompts = [strip_action(p) for p in prompts]
    offset_px = G.frames_to_cells(frame_off)
    info["frame_offset_px_80"] = frame_off[80].tolist()
    info["offset_px"] = offset_px.tolist()
    info["action_idx"] = true_labels
    print(f"指令：frame_offset_px[80]=({frame_off[80, 0]:+.1f},{frame_off[80, 1]:+.1f}) px ⇒ 期望背景位移 "
          f"({-frame_off[80, 0]:+.1f},{-frame_off[80, 1]:+.1f})")
    runs, k0 = [], 0
    for k in range(1, G.CELLS + 1):
        if k == G.CELLS or prompts[k] != prompts[k0]:
            runs.append(f"cell {k0}–{k - 1}: {prompts[k0]}")
            k0 = k
    print(f"文本（{args.arm}）：\n   " + "\n   ".join(runs))

    # ---- 模型 ----
    t0 = time.time()
    if models is not None:
        pipe, dit, *rest = models  # 测试里复用已加载的模型；旧的二元组也接受（无臂）
        arm_module = rest[0] if rest else None
    else:
        pipe, dit, arm_module = load_models(args.ckpt, args.mask_channel, args.device, arm=args.arm, arm_kwargs=args.arm_kwargs or arm_kwargs)
    has_mask = hasattr(dit, MASK_EMBEDDING_NAME)
    info["n_new_params"] = arm_module.n_new_params() if arm_module is not None else 0
    print(f"模型加载 {time.time() - t0:.0f}s · arm={args.arm} · mask_channel={has_mask}"
          + (f" · {type(arm_module).__name__} {info['n_new_params']:,} 参数" if arm_module is not None else ""))

    # 文本表：动作走文本的臂查含动作从句的表，其余查纯场景表
    table_key = ("plain_raw" if args.label_mode == "mirror" else "plain") if text_arm else "arope"
    enc = TextEncoder(args.text_table or TEXT_TABLE[table_key], args.device)
    # prompt 臂（PromptArm）要求 context 逐 cell，全同也展开；其它臂全同时用整条（逐位等价、更省）
    per_cell = getattr(arm_module, "per_cell_context", False) and getattr(arm_module, "require_per_cell", False)
    context, context_ids = build_context(enc, prompts, args.device, per_cell=per_cell)
    neg_context = neg_ids = None
    if args.cfg != 1.0:
        neg_prompts = [BARE_TMPL.format(action=ACTIONS[l]) for l in labels] if text_arm else [NULL_PROMPT] * G.CELLS
        neg_context, neg_ids = build_context(enc, neg_prompts, args.device, per_cell=per_cell)
        info["neg_prompts"] = sorted(set(neg_prompts))
    info["text_fallback_t5"] = enc.n_fallback

    img = load_first_frame(first_spec)
    z0 = encode_first_frame(pipe, img, tiled=args.tiled)

    if args.arm == "arope":
        off_t = torch.tensor(offset_px, dtype=torch.float32).unsqueeze(0).to(args.device)
        mask_in = loss_weight_map(off_t).eq(1).float() if has_mask else None
    else:
        off_t, mask_in = None, None
    # baseline 臂的动作输入（baseline/SPEC.md）：与 dataset.py 同一个 make_action_inputs，[1, 21, ·]
    action_inputs = None
    if arm_module is not None:
        action_inputs = {k: v.to(args.device) for k, v in make_action_inputs(
            torch.tensor(offset_px, dtype=torch.float32).unsqueeze(0), torch.tensor([true_labels])).items()}

    # ---- 采样 + 解码 ----
    t0 = time.time()
    latents = sample(dit, z0, context, context_ids, off_t, mask_in, steps=args.steps, shift=args.shift,
                     cfg=args.cfg, seed=args.seed, neg_context=neg_context, neg_context_ids=neg_ids,
                     arm=arm_module, action_inputs=action_inputs)
    torch.cuda.synchronize()
    info["sample_sec"] = time.time() - t0
    t0 = time.time()
    frames = decode_video(pipe, latents, tiled=args.tiled)
    info["decode_sec"] = time.time() - t0
    info["peak_mem_gb"] = torch.cuda.max_memory_allocated() / 1e9
    print(f"采样 {info['sample_sec']:.1f}s（{args.steps} 步）· 解码 {info['decode_sec']:.1f}s · 峰值显存 {info['peak_mem_gb']:.1f} GB")

    from PIL import Image
    from diffsynth.utils.data import save_video
    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
    stem = os.path.splitext(args.out)[0]
    save_video(list(frames), args.out, fps=int(G.FPS), quality=9)
    Image.fromarray(img).save(stem + "_frame0.png")
    info["out"] = args.out
    info["frames"] = int(frames.shape[0])
    info["resolution"] = [int(frames.shape[2]), int(frames.shape[1])]

    # ---- 实测位移 ----
    meas = measure_displacement(to_gray(frames))
    expect = (-frame_off[80]).tolist()
    info["measured"] = meas
    info["expected_bg_shift_80"] = expect
    err = None if meas["sift"] is None else [meas["sift"][0] - expect[0], meas["sift"][1] - expect[1]]
    info["error_px"] = err
    print(f"实测位移 帧0→80：sift={fmt(meas['sift'])}（直接 {fmt(meas['sift_direct'])} / 链式 {fmt(meas['sift_chain'])}，"
          f"{meas['sift_chain_n_segments']} 段×{meas['sift_chain_step']} 帧；采用 {meas['sift_source']}）"
          f" · 指令 −frame_offset_px[80]={fmt(expect)} · 误差 {fmt(err)}")

    if args.replay:
        gt = read_frames(find_clip_mp4(args.replay))
        n = min(len(gt), len(frames))
        per = [psnr(frames[i], gt[i]) for i in range(n)]
        info["psnr_frame0"] = per[0]
        info["psnr_mean_1_80"] = float(np.mean(per[1:]))
        gmeas = measure_displacement(to_gray(gt))
        info["gt_measured"] = gmeas
        print(f"与 GT 比：帧 0 PSNR {per[0]:.1f} dB · 帧 1–80 平均 PSNR {info['psnr_mean_1_80']:.1f} dB · "
              f"GT 实测 sift={fmt(gmeas['sift'])}（{gmeas['sift_source']}）")

    info["total_sec"] = time.time() - t_all
    with open(stem + ".json", "w", encoding="utf-8") as fh:
        json.dump(info, fh, ensure_ascii=False, indent=1, default=float)
    print(f"DONE {args.out} · {info['total_sec']:.0f}s")
    return info


if __name__ == "__main__":
    main()