"""SyncWorld — visual calibration turns a video world model into a zero-shot robot simulator.

Gradio / ZeroGPU demo for `yyuncong/SyncWorld` (paper 2609.09155, UMass Embodied AGI).

The inference path is ported 1:1 from the authors' reference script
`examples/eval_gripperhead_fdm_rollout.py` in
https://github.com/UMass-Embodied-AGI/SyncWorld — same config surgery, same
multi-item `[calib x 6] + [history] + [current+future]` batch, same sampler
call, same conditioning horizon (25 sparse history frames @ stride 3, 16
predicted frames, 512px, conditioning fps 15, 7-D `[dpos_cm, deuler_deg,
gripper]` actions in the `backward_framewise` convention).
"""
import os
import sys

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("COSMOS_DEVICE", "cuda")

import spaces  # noqa: F401  # MUST precede torch / any CUDA-touching import

from cosmos_framework.inference.common.init import init_script

init_script()  # sets grad off, seeds, wires the cosmos logger (matches the reference script)
sys.excepthook = sys.__excepthook__  # undo init_script's distributed excepthook (Gradio needs the default)

import functools
import json
import math
import pickle
import shutil
import tempfile
import time
import types
from pathlib import Path

import gradio as gr
import imageio
import numpy as np
import torch
import torch.nn.functional as F
from huggingface_hub import hf_hub_download, snapshot_download

from cosmos_framework.configs.base.defaults.compile import CompileConfig
from cosmos_framework.data.vfm.action.calib_segments import build_calib_segment_indices
from cosmos_framework.data.vfm.action.datasets.gripperhead_fdm_dataset import NEUTRAL_CAPTION_SEED
from cosmos_framework.data.vfm.action.domain_utils import get_domain_id
from cosmos_framework.data.vfm.action.pose_utils import pose_abs_to_rel
from cosmos_framework.data.vfm.action.transforms import build_sequence_plan_from_mode
from cosmos_framework.inference.model import Cosmos3OmniConfig, Cosmos3OmniModel
from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import tokenize_caption

HERE = Path(__file__).parent
DEVICE = "cuda"
OUT_DIR = Path(tempfile.gettempdir()) / "syncworld_out"
OUT_DIR.mkdir(parents=True, exist_ok=True)

CKPT_REPO = "yyuncong/SyncWorld"
VAE_REPO = "Wan-AI/Wan2.2-TI2V-5B"
VAE_FILE = "Wan2.2_VAE.pth"
QWEN_REPO = "nvidia/Cosmos3-Nano"
# Only the processor / tokenizer artifacts at the repo root — NOT the 35GB of weights.
QWEN_FILES = [
    "chat_template.json",
    "merges.txt",
    "preprocessor_config.json",
    "tokenizer.json",
    "tokenizer_config.json",
    "video_preprocessor_config.json",
    "vocab.json",
]

# ---------------------------------------------------------------- shipped recipe
# These are the trained checkpoint's conditioning + sampling defaults, i.e. the
# argparse defaults of the authors' eval script. They MUST match the checkpoint.
def _recipe(**over):
    a = types.SimpleNamespace(
        num_history_frames=25,
        num_pred_frames=16,
        history_frame_stride=3,
        resolution=512,
        fps=15.0,
        use_calibration=True,
        calib_null=False,
        calib_segments=6,
        calib_seg_len=5,
        calib_frame_interval=3,
        calib_positive_actions=False,
        action_convention="backward_framewise",
        action_rot_format="euler_xyz",
        action_trans_scale=100.0,
        action_rot_scale=57.2958,
        action_cfg_scale=1.0,
        num_steps=20,
        seed=0,
    )
    for k, v in over.items():
        setattr(a, k, v)
    return a


# ================================================================ model loading
# (ported from the reference script: _REWRITES / _config_from_dir /
#  _patch_offline_vlm_processor / load_model)
_REWRITES = [
    ("cosmos3._src.vfm.configs.base.", "cosmos_framework.configs.base."),
    ("cosmos3._src.vfm.models.", "cosmos_framework.model.vfm."),
    ("cosmos3._src.vfm.tokenizers.", "cosmos_framework.model.vfm.tokenizers."),
    ("cosmos3._src.imaginaire.", "cosmos_framework."),
]


def _config_from_dir(d: Path, vae_path: str = "") -> Cosmos3OmniConfig:
    cfg = d / "config.json" if (d / "config.json").exists() else d / "model" / "config.json"
    if not cfg.exists():
        raise FileNotFoundError(f"no config.json under {d}")
    text = cfg.read_text()
    for a, b in _REWRITES:
        text = text.replace(a, b)
    model_cfg = json.loads(text)["model"]
    # The shipped config resolves the Wan2.2 video VAE through an object store
    # (vae_path="pretrained/..." + bucket_name="bucket"). Point it at the local
    # .pth and clear the bucket so loading is self-contained.
    wan = (vae_path or "").strip()
    if wan:
        def _override_vae(o):
            if isinstance(o, dict):
                if "vae_path" in o:  # the VIDEO tokenizer dict ("avae_path" is the audio one)
                    o["vae_path"] = wan
                    o["bucket_name"] = ""
                for v in o.values():
                    _override_vae(v)
            elif isinstance(o, list):
                for v in o:
                    _override_vae(v)
        _override_vae(model_cfg)
    # Sound is DISABLED in the gripperhead FDM recipe, so the checkpoint has no
    # sound expert; leaving sound_gen=True would make the build fetch an audio
    # tokenizer that does not exist for this release.
    if isinstance(model_cfg.get("config"), dict):
        model_cfg["config"]["sound_gen"] = False
        model_cfg["config"]["sound_tokenizer"] = None
    return Cosmos3OmniConfig(model=model_cfg)


def _patch_offline_vlm_processor(qwen_assets: str):
    """Serve the VLM processor from a local dir instead of shelling out to `uvx hf download`
    (which would pull the whole 35GB nvidia/Cosmos3-Nano repo)."""
    import cosmos_framework.utils.checkpoint_db as _ckdb

    _orig = _ckdb._hf_download

    def _patched(cmd_args):
        repo = str(cmd_args[0]) if cmd_args else ""
        if "Cosmos3-Nano" in repo:
            print(f"[boot] offline VLM processor: _hf_download({repo}) -> {qwen_assets}", flush=True)
            return qwen_assets
        return _orig(cmd_args)

    _ckdb._hf_download = _patched


def _load_model(checkpoint: str, vae_path: str, qwen_assets: str):
    _patch_offline_vlm_processor(qwen_assets)
    ckpt = Path(checkpoint)
    wrapper = Cosmos3OmniModel.from_pretrained_dcp(
        ckpt,
        config=_config_from_dir(ckpt, vae_path),
        compile_config=CompileConfig(enabled=False),  # torch.compile is off: ZeroGPU forks a fresh worker
    )
    m = wrapper.model
    # Single-process inference: null parallel_dims so every collective site takes
    # its local no-op branch (the reference does the same for single-GPU eval).
    if getattr(m, "parallel_dims", None) is not None:
        m.parallel_dims = None
    m.eval()
    return m


# ================================================================ episode reading
# (ported: _read_video / _center_square_crop_np / _load_pose / _build_action /
#  read_episode / _to_u8 / build_fdm_batch_multiitem / build_calib_eval_items /
#  _gen_window)
def _read_video(path: str) -> np.ndarray:
    import imageio.v3 as iio

    try:
        return iio.imread(path, plugin="pyav")  # (T,H,W,C) uint8
    except Exception:
        return iio.imread(path)


@functools.lru_cache(maxsize=4)
def _read_video_cached(path: str) -> np.ndarray:
    return _read_video(path)


def _center_square_crop_np(a: np.ndarray) -> np.ndarray:
    h, w = a.shape[1], a.shape[2]
    if h == w:
        return a
    m = min(h, w)
    top, left = (h - m) // 2, (w - m) // 2
    return a[:, top:top + m, left:left + m]


def _load_pose(leaf: str, thr: float = 0.6):
    with open(os.path.join(leaf, "pose.pkl"), "rb") as f:
        d = pickle.load(f)
    mats = np.asarray(d["gripper_matrix"], dtype=np.float32)  # (T,4,4)
    gopen = np.asarray(d.get("gripper_open", np.ones(len(mats))), dtype=np.float32)
    return mats, (gopen > thr).astype(np.float32)


def _build_action(mats, gopen, convention="backward_framewise", rot_format="euler_xyz",
                  trans_scale=100.0, rot_scale=57.2958) -> torch.Tensor:
    """abs poses -> per-step delta actions [dpos_cm(3), deuler_deg(3), gripper(1)]."""
    poses_rel = pose_abs_to_rel(mats, rotation_format=rot_format, pose_convention=convention,
                               translation_scale=trans_scale, rotation_scale=rot_scale)
    grip = gopen[1:].reshape(-1, 1)
    return torch.from_numpy(np.concatenate([poses_rel, grip], axis=-1).astype(np.float32))


def _frames_to_pm1(frames: np.ndarray, res: int) -> torch.Tensor:
    v = torch.from_numpy(_center_square_crop_np(frames)).float().permute(0, 3, 1, 2) / 255.0
    _, _, hh, ww = v.shape
    th = res
    tw = max(16, int(round((res * ww / hh) / 16) * 16))  # aspect preserving, /16 for the VAE
    if (hh, ww) != (th, tw):
        v = F.interpolate(v, size=(th, tw), mode="bilinear", align_corners=False)
    return v * 2 - 1


def read_episode(leaf: str, view: str, res: int):
    frames = _read_video_cached(os.path.join(leaf, view, "video.mp4"))
    v = _frames_to_pm1(frames, res)
    mats, gopen = _load_pose(leaf)
    n = min(len(v), len(mats))
    # FDM training always emits NEUTRAL_CAPTION_SEED, so anything else is out of distribution.
    return v[:n], mats[:n], gopen[:n], NEUTRAL_CAPTION_SEED


def _to_u8(x: torch.Tensor) -> torch.Tensor:
    return ((x.clamp(-1, 1) + 1) / 2 * 255.0).round().clamp(0, 255).to(torch.uint8)


def build_fdm_batch_multiitem(model, video_u8_list, action_list, caption, device, fps=15.0,
                              mode: str = "forward_dynamics"):
    """N-item FDM batch: `[calib x K] + [history] + [current+future]`. Every item but the
    last is fully conditioning; the last conditions on its current-frame latent and
    generates the P future latents."""
    maxD = model.config.max_action_dim

    def pad(a):
        p = torch.zeros(a.shape[0], maxD, device=device)
        p[:, : a.shape[1]] = a.to(device)
        return p

    vids = [v.to(device) for v in video_u8_list]
    acts = [pad(a) for a in action_list]
    dims = [torch.tensor(int(a.shape[1]), dtype=torch.long, device=device) for a in action_list]
    sizes = [torch.tensor([[v.shape[-2], v.shape[-1], v.shape[-2], v.shape[-1]]],
                          dtype=torch.float32, device=device) for v in vids]
    cf = vids[-1]
    sp = build_sequence_plan_from_mode(mode, video_length=cf.shape[1],
                                       action_length=action_list[-1].shape[0], has_text=True,
                                       num_condition_latent_frames=1)
    sp.share_vision_temporal_positions = False  # distinct time states per item — MUST match training
    ids = tokenize_caption(caption, model.vlm_tokenizer, is_video=False,
                           use_system_prompt=model.vlm_config.use_system_prompt)
    return {
        model.input_video_key: [vids],
        "action": [acts],
        "raw_action_dim": [dims],
        "image_size": [sizes],
        "mode": [mode],
        model.input_caption_key: [caption],
        "text_token_ids": [torch.tensor(ids, dtype=torch.long, device=device).unsqueeze(0)],
        "fps": torch.tensor([float(fps)], device=device),
        "conditioning_fps": torch.tensor([float(fps)], device=device),
        "domain_id": [torch.tensor(get_domain_id("gripperhead"), dtype=torch.long, device=device)],
        "sequence_plan": [sp],
        # NO is_preprocessed -> the uint8 items get normalized + re-stacked for the VAE
    }


def build_calib_items(calib_dir: str, args, view: str):
    """Load + per-DoF segment the calibration sweep into (K uint8 video items, K action blocks)."""
    frames = _read_video_cached(os.path.join(calib_dir, view, "video.mp4"))
    v = _frames_to_pm1(frames, args.resolution)
    mats, gopen = _load_pose(calib_dir)
    n = min(len(mats), len(v))
    move_order = None
    mrp = os.path.join(calib_dir, "move_range.pkl")
    if os.path.isfile(mrp):
        try:
            with open(mrp, "rb") as f:
                move_order = pickle.load(f).get("movement_order")
        except Exception:
            move_order = None
    sub = list(range(0, n, args.calib_frame_interval)) or [0]
    efficient = (args.calib_segments == 6)  # 6 -> one segment per DoF; 12 -> both signs
    seg_lists = build_calib_segment_indices(mats[sub], args.calib_seg_len, efficient, move_order,
                                            positive_body_actions=bool(args.calib_positive_actions))
    vids, acts = [], []
    for seg in seg_lists:
        real = [sub[i] for i in seg]
        vids.append(_to_u8(v[real].permute(1, 0, 2, 3)))  # (C,seg_len,H,W) uint8
        acts.append(_build_action(mats[real], gopen[real], args.action_convention,
                                  args.action_rot_format, args.action_trans_scale, args.action_rot_scale))
    return vids, acts


def _gen_window(model, ep_frames, gen_frames, mats, gopen, caption, calib_v, calib_a,
                start, args, device, custom_act=None):
    """Generate ONE 1+P frame window at frame `start`. Returns (pred_u8, gt_u8) as (P+1,H,W,C)."""
    T = ep_frames.shape[0]
    H, P, S = args.num_history_frames, args.num_pred_frames, args.history_frame_stride
    clamp = lambda x: max(0, min(int(x), T - 1))  # noqa: E731
    use_history, use_calib = H > 1, args.use_calibration
    cf_idx = [clamp(start)] + [clamp(start + 1 + j) for j in range(P)]
    cf_frames = torch.stack([gen_frames[clamp(start)]]
                            + [ep_frames[clamp(start + 1 + j)] for j in range(P)], dim=1)
    # custom-action mode: drive the window with a synthesized cm/deg action. The HISTORY
    # action block stays real — it is observed context, not the commanded motion.
    cf_act = (custom_act.to(torch.float32) if custom_act is not None else
              _build_action(mats[cf_idx], gopen[cf_idx], args.action_convention,
                            args.action_rot_format, args.action_trans_scale, args.action_rot_scale))
    items_v = list(calib_v) if use_calib else []
    items_a = list(calib_a) if use_calib else []
    if use_history:
        hist_idx = [clamp(start - S * (H - i)) for i in range(H)]
        hist_frames = torch.stack([gen_frames[i] for i in hist_idx], dim=1)
        hist_act = _build_action(mats[hist_idx], gopen[hist_idx], args.action_convention,
                                 args.action_rot_format, args.action_trans_scale, args.action_rot_scale)
        items_v.append(_to_u8(hist_frames))
        items_a.append(hist_act)
    items_v.append(_to_u8(cf_frames))
    items_a.append(cf_act)
    batch = build_fdm_batch_multiitem(model, items_v, items_a, caption, device, fps=args.fps)
    with torch.no_grad():
        outputs = model.generate_samples_from_batch(
            batch, guidance=1.0, action_guidance=args.action_cfg_scale,
            seed=[args.seed], num_steps=args.num_steps)
        dec = model.decode(outputs["vision"][0])[0].clamp(-1, 1)  # (C,P+1,h,w)
    pred_u8 = ((dec.float().permute(1, 2, 3, 0) + 1) / 2 * 255).round().clamp(0, 255).byte().cpu().numpy()
    gt_win = torch.stack([ep_frames[clamp(start + i)] for i in range(P + 1)], dim=0)
    gt_u8 = ((gt_win.float().permute(0, 2, 3, 1) + 1) / 2 * 255).round().clamp(0, 255).byte().cpu().numpy()
    return pred_u8, gt_u8


# ================================================================ boot
_t = time.perf_counter()
print("[boot] fetching processor assets ...", flush=True)
QWEN_DIR = None
for _f in QWEN_FILES:
    QWEN_DIR = os.path.dirname(hf_hub_download(QWEN_REPO, _f))
print(f"[boot] processor assets at {QWEN_DIR}", flush=True)

print("[boot] fetching Wan2.2 video VAE ...", flush=True)
VAE_PATH = hf_hub_download(VAE_REPO, VAE_FILE)

print(f"[boot] fetching {CKPT_REPO} (~31 GB) ...", flush=True)
CKPT_DIR = snapshot_download(CKPT_REPO)
print(f"[boot] checkpoint at {CKPT_DIR} ({time.perf_counter() - _t:.0f}s)", flush=True)

print("[boot] building + loading SyncWorld ...", flush=True)
MODEL = _load_model(CKPT_DIR, VAE_PATH, QWEN_DIR)
print(f"[boot] model ready ({time.perf_counter() - _t:.0f}s)", flush=True)

if not os.environ.get("SYNCWORLD_KEEP_CKPT"):
    # The weights now live in (fake-)CUDA tensors that ZeroGPU packs to its own
    # on-disk store; the 31 GB download is dead weight and would blow the Space's
    # ephemeral disk once the pack copy lands.
    try:
        repo_root = Path(CKPT_DIR).parent.parent
        if repo_root.name.startswith("models--"):
            shutil.rmtree(repo_root, ignore_errors=True)
            print(f"[boot] freed {repo_root}", flush=True)
    except Exception as e:  # pragma: no cover
        print(f"[boot] WARN could not free checkpoint dir: {e!r}", flush=True)

# ================================================================ scenes
with open(HERE / "assets" / "scenes" / "scenes.json") as f:
    SCENE_META = json.load(f)
SCENES = {}
for _slug, _m in SCENE_META.items():
    _base = HERE / "assets" / "scenes" / _slug
    SCENES[_m["label"]] = dict(
        slug=_slug,
        leaf=str(_base / "expert"),
        calib=str(_base / "calibration"),
        view=_m["view"],
        num_frames=int(_m["num_frames"]),
        thumb=str(_base / "thumb.jpg"),
        expert_video=str(_base / "expert" / _m["view"] / "video.mp4"),
        calib_video=str(_base / "calibration" / _m["view"] / "video.mp4"),
        suite="ManiSkill" if _m["suite"].endswith("maniskill") else "LIBERO",
        task=_m["task"],
    )
SCENE_NAMES = sorted(SCENES)
MAX_FRAMES = max(v["num_frames"] for v in SCENES.values())

MODE_REPLAY = "Replay the robot's own actions"
MODE_DRIVE = "Drive the arm myself"
GRIP_KEEP, GRIP_OPEN, GRIP_CLOSE = "keep current", "open", "close"


def _write_mp4(frames_u8: np.ndarray, fps: int = 15) -> str:
    d = tempfile.mkdtemp(dir=str(OUT_DIR))
    p = os.path.join(d, "rollout.mp4")
    imageio.mimwrite(p, list(frames_u8), fps=fps, macro_block_size=1, quality=9)
    return p


def _estimate_duration(*a, **k):
    """ZeroGPU reservation, fitted to latency measured on this Space.

    Client-side end-to-end at 20 sampler steps: 1 window 28.6 s, 2 windows 50.6 s,
    3 windows 67.2 s -> ~1.0 s per sampler step per window, plus ~12 s of scene
    prep, VAE encode/decode and mp4 muxing. Keeps ~20% headroom, no more.
    """
    rounds = int(a[10]) if len(a) > 10 else int(k.get("rollout_rounds", 2))
    steps = int(a[11]) if len(a) > 11 else int(k.get("num_steps", 20))
    rounds = max(1, min(rounds, 3))
    steps = max(1, min(steps, 40))
    return int(min(180, 12 + rounds * (4 + steps * 1.0)))


@spaces.GPU(duration=_estimate_duration)
def simulate(
    scene: str,
    action_mode: str = MODE_REPLAY,
    start_frame: int = 0,
    delta_x_cm: float = 0.0,
    delta_y_cm: float = 0.0,
    delta_z_cm: float = 0.0,
    delta_rot_x_deg: float = 0.0,
    delta_rot_y_deg: float = 0.0,
    delta_rot_z_deg: float = 0.0,
    gripper: str = GRIP_KEEP,
    rollout_rounds: int = 2,
    num_steps: int = 20,
    action_guidance: float = 1.0,
    seed: int = 0,
    progress=gr.Progress(track_tqdm=True),
):
    """Roll out the SyncWorld world model on a visually-calibrated robot episode.

    Args:
        scene: name of a bundled calibrated episode (a ManiSkill or LIBERO tabletop scene).
        action_mode: "Replay the robot's own actions" to re-simulate the expert trajectory,
            or "Drive the arm myself" to command a constant per-step end-effector twist.
        start_frame: index of the episode frame used as the current observation.
        delta_x_cm: commanded per-step end-effector translation along body X, in centimetres.
        delta_y_cm: commanded per-step end-effector translation along body Y, in centimetres.
        delta_z_cm: commanded per-step end-effector translation along body Z, in centimetres.
        delta_rot_x_deg: commanded per-step end-effector rotation about X, in degrees.
        delta_rot_y_deg: commanded per-step end-effector rotation about Y, in degrees.
        delta_rot_z_deg: commanded per-step end-effector rotation about Z, in degrees.
        gripper: commanded gripper state ("keep current", "open" or "close").
        rollout_rounds: number of autoregressive 16-frame windows to chain.
        num_steps: rectified-flow sampler steps per window.
        action_guidance: action classifier-free-guidance scale.
        seed: sampler seed.

    Returns:
        A tuple of (predicted rollout mp4, ground-truth window mp4, markdown run report).
    """
    t0 = time.perf_counter()
    sc = SCENES.get(scene) or SCENES[SCENE_NAMES[0]]
    args = _recipe(num_steps=int(num_steps), action_cfg_scale=float(action_guidance), seed=int(seed))

    ep_frames, mats, gopen, caption = read_episode(sc["leaf"], sc["view"], args.resolution)
    calib_v, calib_a = build_calib_items(sc["calib"], args, sc["view"])
    t_prep = time.perf_counter() - t0

    T = int(ep_frames.shape[0])
    P = args.num_pred_frames
    start = max(0, min(int(start_frame), T - 1))
    rounds = max(1, min(int(rollout_rounds), 3))
    drive = str(action_mode) == MODE_DRIVE

    custom_row = None
    if drive:
        if gripper == GRIP_OPEN:
            g = 1.0
        elif gripper == GRIP_CLOSE:
            g = 0.0
        else:
            g = float(gopen[start])
        custom_row = torch.tensor(
            [float(delta_x_cm), float(delta_y_cm), float(delta_z_cm),
             float(delta_rot_x_deg), float(delta_rot_y_deg), float(delta_rot_z_deg), g],
            dtype=torch.float32)

    gen_frames = ep_frames.clone()
    pred_acc, gt_acc = [], []
    for r in range(rounds):
        s = start + r * P
        cact = custom_row.unsqueeze(0).repeat(P, 1) if custom_row is not None else None
        pred_u8, gt_u8 = _gen_window(MODEL, ep_frames, gen_frames, mats, gopen, caption,
                                     calib_v, calib_a, s, args, DEVICE, custom_act=cact)
        if rounds > 1:  # closed-loop write-back: the next round conditions on generated frames
            for i in range(min(pred_u8.shape[0], gen_frames.shape[0] - s)):
                fr = pred_u8[i].astype(np.float32) / 255.0 * 2.0 - 1.0
                gen_frames[s + i] = torch.from_numpy(fr).permute(2, 0, 1).to(gen_frames)
        skip = 0 if r == 0 else 1  # drop the 1-frame overlap between rounds
        m = min(pred_u8.shape[0], gt_u8.shape[0])
        pred_acc.extend(pred_u8[i] for i in range(skip, m))
        gt_acc.extend(gt_u8[i] for i in range(skip, m))

    pred = np.stack(pred_acc, 0)
    gt = np.stack(gt_acc, 0)
    pred_path = _write_mp4(pred)
    gt_path = _write_mp4(gt)
    dt = time.perf_counter() - t0

    if drive:
        act_desc = (f"commanded twist  `Δpos = ({delta_x_cm:+.2f}, {delta_y_cm:+.2f}, "
                    f"{delta_z_cm:+.2f}) cm/step`,  `Δrot = ({delta_rot_x_deg:+.2f}, "
                    f"{delta_rot_y_deg:+.2f}, {delta_rot_z_deg:+.2f}) deg/step`,  gripper "
                    f"**{gripper}** — held for {rounds * P} steps")
    else:
        act_desc = f"the episode's own recorded actions for frames {start}…{min(start + rounds * P, T - 1)}"
    report = (
        f"**{sc['suite']} · `{sc['task']}`**  ·  camera `{sc['view']}`  ·  "
        f"episode has {T} frames\n\n"
        f"Conditioned on **6 per-DoF calibration segments** + **25 sparse history frames** "
        f"(stride 3, from frame {max(0, start - 3 * 25)}) + current frame **{start}**.\n\n"
        f"Action: {act_desc}.\n\n"
        f"Generated **{len(pred)} frames** in {rounds} autoregressive window(s) × {args.num_steps} "
        f"sampler steps  ·  **{dt:.1f}s** total (scene prep {t_prep:.1f}s)."
    )
    return pred_path, gt_path, report


# ================================================================ UI
def _scene_preview(scene: str, start_frame: int):
    sc = SCENES.get(scene) or SCENES[SCENE_NAMES[0]]
    T = sc["num_frames"]
    s = max(0, min(int(start_frame), T - 1))
    try:
        frame = _read_video_cached(sc["expert_video"])[s]
    except Exception:
        frame = None
    info = (f"`{sc['suite']}` · `{sc['task']}` · camera `{sc['view']}` · **{T} frames** "
            f"· current frame **{s}**")
    return frame, sc["calib_video"], sc["expert_video"], info


THEME = gr.themes.Citrus()

DESC = """\
# 🦾 SyncWorld — a zero-shot robot simulator

[**SyncWorld**](https://huggingface.co/papers/2609.09155) shows that a pretrained video world model
becomes a usable robot simulator once you *visually calibrate* it: prepend a short clip of the arm
sweeping each degree of freedom, and the model infers the unseen camera↔robot mapping on the fly —
no per-scene finetuning.

Pick a calibrated scene, then either **replay the robot's own actions** or **drive the arm yourself**
with a per-step end-effector twist, and watch the model render what happens next.

*Model: [`yyuncong/SyncWorld`](https://huggingface.co/yyuncong/SyncWorld) (16B Cosmos-3 mixture-of-transformers,
Wan2.2 video VAE). Scenes and calibration clips are the authors' own evaluation episodes from
[`yyuncong/SyncWorld-Evaluation`](https://huggingface.co/datasets/yyuncong/SyncWorld-Evaluation) (OpenMDW-1.1).*
"""

with gr.Blocks(theme=THEME, title="SyncWorld — zero-shot robot simulator") as demo:
    gr.Markdown(DESC)

    with gr.Row():
        with gr.Column(scale=5):
            scene = gr.Dropdown(SCENE_NAMES, value=SCENE_NAMES[0], label="Calibrated scene")
            info_md = gr.Markdown()
            with gr.Row():
                cur_img = gr.Image(label="Current frame (the model's last observation)",
                                   height=232, interactive=False)
                calib_vid = gr.Video(label="Calibration sweep (6 DoF)", height=232,
                                     interactive=False, autoplay=True, loop=True)
            start_frame = gr.Slider(0, MAX_FRAMES - 1, value=0, step=1,
                                    label="Current frame index (clamped to the episode length)")
            action_mode = gr.Radio([MODE_REPLAY, MODE_DRIVE], value=MODE_REPLAY, label="Actions")
            with gr.Group():
                gr.Markdown("**Commanded end-effector twist** — used in *Drive the arm myself* mode. "
                            "Held constant for every predicted step (the calibration sweeps cover "
                            "roughly ±3 cm/step and ±3 °/step).")
                with gr.Row():
                    dx = gr.Slider(-3.0, 3.0, value=0.0, step=0.05, label="Δ X (cm/step)")
                    dy = gr.Slider(-3.0, 3.0, value=0.0, step=0.05, label="Δ Y (cm/step)")
                    dz = gr.Slider(-3.0, 3.0, value=0.0, step=0.05, label="Δ Z (cm/step)")
                with gr.Row():
                    rx = gr.Slider(-3.0, 3.0, value=0.0, step=0.05, label="Δ rot X (°/step)")
                    ry = gr.Slider(-3.0, 3.0, value=0.0, step=0.05, label="Δ rot Y (°/step)")
                    rz = gr.Slider(-3.0, 3.0, value=0.0, step=0.05, label="Δ rot Z (°/step)")
                gripper = gr.Radio([GRIP_KEEP, GRIP_OPEN, GRIP_CLOSE], value=GRIP_KEEP,
                                   label="Gripper")
            rollout_rounds = gr.Slider(1, 3, value=2, step=1,
                                       label="Rollout windows (16 generated frames each)")
            with gr.Accordion("Advanced", open=False):
                num_steps = gr.Slider(4, 40, value=20, step=1, label="Sampler steps per window")
                action_guidance = gr.Slider(1.0, 5.0, value=1.0, step=0.1,
                                            label="Action guidance (CFG)")
                seed = gr.Slider(0, 2**31 - 1, value=0, step=1, label="Seed")
            run = gr.Button("Simulate", variant="primary")

        with gr.Column(scale=5):
            out_pred = gr.Video(label="SyncWorld rollout", autoplay=True, loop=True,
                                interactive=False, height=360)
            out_gt = gr.Video(label="Ground truth for the same window (reference)",
                              autoplay=True, loop=True, interactive=False, height=240)
            report_md = gr.Markdown()

    with gr.Accordion("Full expert episode (reference)", open=False):
        expert_vid = gr.Video(label="Expert demonstration", interactive=False, height=360)

    INPUTS = [scene, action_mode, start_frame, dx, dy, dz, rx, ry, rz, gripper,
              rollout_rounds, num_steps, action_guidance, seed]
    OUTPUTS = [out_pred, out_gt, report_md]

    # Every example row supplies the FULL input list. Gradio inserts its Progress
    # object at the positional index of the `progress` parameter, so a short row
    # would silently shift values into the wrong arguments.
    EXAMPLES = [
        # scene, mode, start, dx, dy, dz, rx, ry, rz, gripper, rounds, steps, cfg, seed
        ["ManiSkill · PushCube", MODE_REPLAY, 0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, GRIP_KEEP, 2, 20, 1.0, 0],
        ["LIBERO · put the bowl on the plate", MODE_REPLAY, 16, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, GRIP_KEEP, 2, 20, 1.0, 0],
        ["ManiSkill · StackCube", MODE_DRIVE, 32, 0.0, 1.2, 0.0, 0.0, 0.0, 0.0, GRIP_KEEP, 2, 20, 1.0, 0],
        ["ManiSkill · PushCube", MODE_DRIVE, 0, 1.5, 0.0, 0.0, 0.0, 0.0, 0.0, GRIP_KEEP, 2, 20, 1.0, 0],
        ["LIBERO · open the middle drawer", MODE_DRIVE, 24, 0.0, 0.0, -1.2, 0.0, 0.0, 0.0, GRIP_KEEP, 2, 20, 1.0, 0],
        ["LIBERO · alphabet soup into basket", MODE_DRIVE, 48, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, GRIP_CLOSE, 2, 20, 1.0, 0],
    ]
    EXAMPLE_LABELS = [
        "PushCube · replay the expert",
        "LIBERO bowl→plate · replay the expert",
        "StackCube · push +Y at 1.2 cm/step",
        "PushCube · push +X at 1.5 cm/step",
        "LIBERO drawer · pull −Z at 1.2 cm/step",
        "LIBERO soup · close the gripper, hold still",
    ]

    run.click(fn=simulate, inputs=INPUTS, outputs=OUTPUTS)

    for _ev in (scene.change, start_frame.release):
        _ev(fn=_scene_preview, inputs=[scene, start_frame],
            outputs=[cur_img, calib_vid, expert_vid, info_md], show_progress="minimal")

    gr.Examples(
        examples=EXAMPLES,
        example_labels=EXAMPLE_LABELS,
        inputs=INPUTS,
        outputs=OUTPUTS,
        fn=simulate,
        cache_examples=True,
        cache_mode="lazy",
        label="Examples",
    )

    demo.load(fn=_scene_preview, inputs=[scene, start_frame],
              outputs=[cur_img, calib_vid, expert_vid, info_md])

if __name__ == "__main__":
    demo.queue(max_size=12).launch(mcp_server=True, allowed_paths=[str(OUT_DIR), str(HERE / "assets")])