"""ActionRoPE: action as coordinates for spatially consistent control in planar-camera game world models.

Gradio demo of the paper's Wan2.2-TI2V-5B backbone: a first frame (832x480) plus an
action string produce an 81-frame (5 s, 16 fps) video of the game world, with the
camera displacement injected through ActionRoPE (world-coordinate rotary positions)
and Coordinate-Conditioned Noise Scheduling (CNS).
"""
import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  must come before torch
import tempfile  # noqa: E402
import time  # noqa: E402

import gradio as gr  # noqa: E402

from actionrope.geometry import offsets_from_actions  # noqa: E402
from actionrope.media import load_first_frame  # noqa: E402
from actionrope.media import write_video  # noqa: E402
from actionrope.prompts import NULL_PROMPT  # noqa: E402
from arope_loader import Generator, download_space_weights  # noqa: E402

CKPT, NULL_PROMPT_PT, VAE_PATH = download_space_weights()
print("[app] weights cached; building model...")
GEN = Generator(CKPT, VAE_PATH, NULL_PROMPT_PT, device="cuda", torch_dtype=__import__("torch").bfloat16)
print("[app] model ready")


@spaces.GPU(duration=75)
def play(first_frame, actions: str, seed: int, steps: int, use_cns: bool = True):
    """Render an 81-frame game-world video from a first frame and an action plan.

    Args:
        first_frame: initial observation frame (any size; centre-cropped to 26:15 and resized to 832x480).
        actions: comma-separated move plan, each segment "<direction>[*speed]:<frames>",
            e.g. "right:40,left:40". Directions: idle, up, down, left, right, and the four
            diagonals (aliases u/d/l/r/ul/ur/dl/dr). Speed 1.0 is the game's walking speed
            (69.2 px/s horizontally, 46.0 px/s vertically at 832x480).
        seed: RNG seed for the initial noise.
        steps: number of Euler flow-matching steps (paper setting: 50).
        use_cns: use Coordinate-Conditioned Noise Scheduling (the paper's method).

    Returns:
        The rendered 832x480 16 fps mp4 of the 5-second rollout.
    """
    t0 = time.perf_counter()
    try:
        offsets = offsets_from_actions(actions)  # validates the segment spec
    except ValueError as e:
        raise gr.Error(str(e))
    frames, fps = GEN(load_first_frame(first_frame), offsets, prompt=NULL_PROMPT,
                     seed=int(seed), steps=int(steps), cns=bool(use_cns))
    out = os.path.join(tempfile.gettempdir(), f"arope_{int(time.time() * 1000) % 100000}.mp4")
    write_video(out, frames, fps=fps, quality=9)
    print(f"[app] generated {len(frames)} frames in {time.perf_counter() - t0:.1f}s -> {out}")
    return out


def build_ui():
    with gr.Blocks(title="ActionRoPE") as demo:
        gr.Markdown(
            """
            # 🎮 ActionRoPE — action as coordinates for game world models
            Drive a video-diffusion world model with plain **actions** instead of text.
            The first frame pins the scene; the action plan moves the camera through the world.
            Backbones: Wan2.2-TI2V-5B (this demo) — from the paper
            [ActionRoPE: action as coordinates for spatially consistent control in planar-camera game world models](https://huggingface.co/papers/2610.25861).
            """
        )
        with gr.Row():
            with gr.Column():
                first_frame = gr.Image(label="First frame (832x480; other sizes are centre-cropped)", type="filepath")
                actions = gr.Textbox(
                    label="Actions",
                    value="right:40,left:40",
                    info="comma list of <direction>[*speed]:<frames> over 80 frames; "
                         "speed 1.0 = walking speed (69.2 px/s horizontal / 46.0 px/s vertical)",
                    placeholder="right:40,left:40",
                )
                btn = gr.Button("Generate rollout", variant="primary")
                with gr.Accordion("Advanced", open=False):
                    seed = gr.Number(0, label="Seed", precision=0)
                    steps = gr.Slider(10, 50, value=50, step=5, label="Sampling steps", info="paper setting = 50")
                    use_cns = gr.Checkbox(True, label="CNS (coordinate-conditioned noise scheduling)")
            with gr.Column():
                video = gr.Video(label="5-second rollout (81 frames, 16 fps)", autoplay=True)

        gr.Examples(
            examples=[
                ["examples/back_01.png", "right:40,left:40"],
                ["examples/back_264.png", "right:60,left:20"],
                ["examples/bamboo_L.png", "up-right*1.5:32,left:40,down-left:8"],
                ["examples/watercolor_L.png", "right:80"],
            ],
            inputs=[first_frame, actions],
            fn=play,
            outputs=[video],
            cache_examples=True,
            cache_mode="lazy",
            label="Example first frames and action plans",
        )

        btn.click(play, inputs=[first_frame, actions, seed, steps, use_cns], outputs=[video])
    return demo


if __name__ == "__main__":
    build_ui().launch(mcp_server=True, theme=gr.themes.Citrus())