"""StatePlay — State-Aware Game World Model demo (Street Fighter III).

Wraps the official `stateplay` inference pipeline (Wan2.2-TI2V-5B visual expert
+ 0.75B state expert with joint attention) on ZeroGPU. Given a first frame, an
initial game state and a 0.5s-per-slot action sequence, the model rolls out a
5-second video AND predicts the internal game state trajectory (timer, both
players' HP, both players' super meters).
"""

import os
import tempfile

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("MPLCONFIGDIR", tempfile.mkdtemp(prefix="mplconfig-"))

import spaces  # noqa: E402  (must precede torch / CUDA-touching imports)

import math  # noqa: E402
import time  # noqa: E402
from pathlib import Path  # noqa: E402

import gradio as gr  # noqa: E402
import matplotlib  # noqa: E402

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402
import torch  # noqa: E402
from huggingface_hub import hf_hub_download, snapshot_download  # noqa: E402

from stateplay import NEG_PROMPT, SF3_BUTTON_COLS, StatePlayPipeline  # noqa: E402
from stateplay.utils.image import save_mp4  # noqa: E402
from stateplay.utils.state import STATE_COLUMNS, denormalize_state  # noqa: E402

# --------------------------------------------------------------------------- #
# Constants
# --------------------------------------------------------------------------- #

CKPT_REPO = "onepiece1999/StatePlay"
WAN_REPO = "Wan-AI/Wan2.2-TI2V-5B"
FPS = 20
SLOT_FRAMES = 10  # each action slot is held for 10 video frames = 0.5 s @ 20 fps
HEIGHT, WIDTH = 480, 832  # native StatePlay training resolution
STATE_MAX = {"timer": 99, "hp1": 160, "hp2": 160, "meter1": 104, "meter2": 96}
NO_INPUT_TOKENS = {"", "-", ".", "_", "none", "noop", "idle", "neutral", "0"}

# --------------------------------------------------------------------------- #
# Weights: assemble the layout `StatePlayPipeline.from_pretrained` expects
# --------------------------------------------------------------------------- #

print("[startup] downloading weights…", flush=True)
ckpt_path = hf_hub_download(CKPT_REPO, "StatePlay.safetensors")
vae_path = hf_hub_download(WAN_REPO, "Wan2.2_VAE.pth")
t5_path = hf_hub_download(WAN_REPO, "models_t5_umt5-xxl-enc-bf16.pth")
tok_root = snapshot_download(WAN_REPO, allow_patterns=["google/umt5-xxl/*"])

BASE_DIR = Path(tempfile.mkdtemp(prefix="stateplay-base-"))
ti2v_dir = BASE_DIR / "Wan-AI" / "Wan2.2-TI2V-5B"
tok_parent = BASE_DIR / "Wan-AI" / "Wan2.1-T2V-1.3B" / "google"
ti2v_dir.mkdir(parents=True, exist_ok=True)
tok_parent.mkdir(parents=True, exist_ok=True)
os.symlink(vae_path, ti2v_dir / "Wan2.2_VAE.pth")
os.symlink(t5_path, ti2v_dir / "models_t5_umt5-xxl-enc-bf16.pth")
os.symlink(Path(tok_root) / "google" / "umt5-xxl", tok_parent / "umt5-xxl")

print("[startup] building pipeline…", flush=True)
pipe = StatePlayPipeline.from_pretrained(
    base_model_dir=str(BASE_DIR),
    checkpoint_path=ckpt_path,
    torch_dtype=torch.bfloat16,
).to("cuda")
BUTTON_COLS = list(pipe.button_cols)
print(f"[startup] ready. buttons={BUTTON_COLS}", flush=True)


# --------------------------------------------------------------------------- #
# Action-sequence parsing  →  the parquet the pipeline reads
# --------------------------------------------------------------------------- #


def parse_action_slots(actions_text: str, n_slots: int) -> list[list[str]]:
    """Parse the slot notation into `n_slots` lists of pressed button names.

    Slots are separated by `|`, `,`, `;` or newlines; buttons inside a slot are
    joined with `+`. `-` (or an empty slot) means "no input". Extra slots are
    dropped, missing slots are padded with no-input.
    """
    raw = str(actions_text or "").replace("\n", "|").replace(",", "|").replace(";", "|")
    chunks = raw.split("|")
    valid = {b.upper() for b in BUTTON_COLS}
    slots: list[list[str]] = []
    for chunk in chunks:
        chunk = chunk.strip()
        if chunk.lower() in NO_INPUT_TOKENS:
            slots.append([])
            continue
        pressed, unknown = [], []
        for tok in chunk.replace(" ", "+").split("+"):
            tok = tok.strip().upper()
            if not tok:
                continue
            if tok in valid:
                if tok not in pressed:
                    pressed.append(tok)
            else:
                unknown.append(tok)
        if unknown:
            raise gr.Error(
                f"Unknown button(s) {unknown} in slot {len(slots) + 1}. "
                f"Valid buttons: {', '.join(BUTTON_COLS)} (or '-' for no input)."
            )
        slots.append(pressed)
    if len(slots) > n_slots:
        slots = slots[:n_slots]
    while len(slots) < n_slots:
        slots.append([])
    return slots


def build_action_parquet(slots, n_rows: int, state: dict, path: str) -> None:
    """Write the action/state parquet the StatePlay pipeline consumes.

    Buttons are written on every 10th row only (matching the dataset's 10 Hz
    controller log); the pipeline's hold-window densification fills the rest.
    Game-state columns are constant — only the first entry is used as the
    conditioning state, the rest of the trajectory is what the model predicts.
    """
    data = {"frame_id": np.arange(n_rows, dtype=np.int64)}
    for name, value in state.items():
        data[name] = np.full(n_rows, float(value), dtype=np.float32)
    for col in BUTTON_COLS:
        data[col] = np.zeros(n_rows, dtype=np.float32)
    df = pd.DataFrame(data)
    for i, pressed in enumerate(slots):
        row = i * SLOT_FRAMES
        if row >= n_rows:
            break
        for col in pressed:
            df.loc[row, col] = 1.0
    df.to_parquet(path, index=False)


def slot_layout(num_frames: int) -> tuple[int, int]:
    """(number of controller rows, number of 0.5s action slots) for a clip."""
    n_rows = max(SLOT_FRAMES, int(num_frames) - 1)
    return n_rows, math.ceil(n_rows / SLOT_FRAMES)


# --------------------------------------------------------------------------- #
# State-trajectory visualisation
# --------------------------------------------------------------------------- #


def state_table(state: torch.Tensor) -> list[list[float]]:
    """[1, F, 5] normalized state → rows of [time_s, timer, hp1, hp2, meter1, meter2]."""
    values = denormalize_state(state.detach().float().cpu()).squeeze(0).numpy()
    rows = []
    for i, row in enumerate(values):
        rows.append([round(i * 4.0 / FPS, 2)] + [round(float(v), 1) for v in row])
    return rows


def state_chart(rows: list[list[float]]) -> str:
    """Render the predicted state trajectory to a PNG and return its path."""
    arr = np.asarray(rows, dtype=np.float32)
    t = arr[:, 0]
    fig, axes = plt.subplots(3, 1, figsize=(7.2, 6.4), sharex=True)
    fig.suptitle("Predicted game-state trajectory", fontsize=13)

    axes[0].plot(t, arr[:, 2], color="#2b7bba", lw=2, marker="o", ms=3, label="P1 HP")
    axes[0].plot(t, arr[:, 3], color="#d1495b", lw=2, marker="o", ms=3, label="P2 HP")
    axes[0].set_ylim(-5, STATE_MAX["hp1"] + 5)
    axes[0].set_ylabel("health")

    axes[1].plot(t, arr[:, 4], color="#3f8f5b", lw=2, marker="o", ms=3, label="P1 meter")
    axes[1].plot(t, arr[:, 5], color="#e0a419", lw=2, marker="o", ms=3, label="P2 meter")
    axes[1].set_ylim(-5, max(STATE_MAX["meter1"], STATE_MAX["meter2"]) + 5)
    axes[1].set_ylabel("super meter")

    axes[2].plot(t, arr[:, 1], color="#6a4c93", lw=2, marker="o", ms=3, label="timer")
    axes[2].set_ylim(-2, STATE_MAX["timer"] + 2)
    axes[2].set_ylabel("round timer")
    axes[2].set_xlabel("time (s)")

    for ax in axes:
        ax.grid(alpha=0.25)
        ax.legend(loc="upper right", fontsize=8)
    fig.tight_layout()
    out = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    fig.savefig(out.name, dpi=110)
    plt.close(fig)
    return out.name


# --------------------------------------------------------------------------- #
# Inference
# --------------------------------------------------------------------------- #


def _estimate_duration(
    image=None,
    actions_text: str = "",
    prompt: str = "",
    timer: int = 45,
    hp1: int = 160,
    hp2: int = 160,
    meter1: int = 0,
    meter2: int = 0,
    num_frames: int = 101,
    num_inference_steps: int = 30,
    cfg_scale: float = 5.0,
    action_cfg_scale: float = 1.0,
    state_cfg_scale: float = 1.0,
    seed: int = 2,
    *args,
    **kwargs,
):
    """ZeroGPU duration estimate: fixed overhead + per-DiT-pass cost x number of passes.

    Calibrated on zero-a10g (Blackwell, bf16, SDPA) against two measured runs:
      41 frames / 10 steps / cfg 5.0  -> 14.4 s   (latent_t 11, 2 passes)
      101 frames / 30 steps / cfg 5.0 -> 94.6 s   (latent_t 26, 2 passes)
    Both fit t = 1.2 + 0.0599 * latent_t * passes * steps within 4%, i.e. cost is
    essentially linear in latent length. The coefficients below add ~25-30% margin.
    """
    latent_t = (int(num_frames) - 1) // 4 + 1
    passes = 1
    if float(cfg_scale) != 1.0:
        passes += 1
    if float(action_cfg_scale) != 1.0:
        passes += 1
    if float(state_cfg_scale) != 1.0 and float(cfg_scale) == 1.0:
        passes += 1
    total = 8.0 + 0.075 * latent_t * passes * int(num_inference_steps)
    return int(min(260, math.ceil(total)))


@spaces.GPU(duration=_estimate_duration)
def generate(
    image,
    actions_text: str,
    prompt: str,
    timer: int = 45,
    hp1: int = 160,
    hp2: int = 160,
    meter1: int = 0,
    meter2: int = 0,
    num_frames: int = 101,
    num_inference_steps: int = 30,
    cfg_scale: float = 5.0,
    action_cfg_scale: float = 1.0,
    state_cfg_scale: float = 1.0,
    seed: int = 2,
    progress=gr.Progress(track_tqdm=True),
):
    """Roll out a Street Fighter III clip and its internal game state with StatePlay.

    Args:
        image: first frame of the rollout (any aspect; cropped to 832x480).
        actions_text: controller inputs, one 0.5s slot per `|`-separated entry,
            buttons joined with `+` (e.g. "D | UP+RIGHT | - | A"). Valid buttons:
            UP, DOWN, LEFT, RIGHT, Y, X, Z, A, B, C, D.
        prompt: NPC behaviour / strategy description used as text conditioning.
        timer: initial round timer (0-99).
        hp1: initial player-1 health (0-160).
        hp2: initial player-2 health (0-160).
        meter1: initial player-1 super meter (0-104).
        meter2: initial player-2 super meter (0-96).
        num_frames: rollout length in frames at 20 fps.
        num_inference_steps: flow-matching denoising steps.
        cfg_scale: text classifier-free-guidance scale.
        action_cfg_scale: action classifier-free-guidance scale (1.0 = off).
        state_cfg_scale: state classifier-free-guidance scale (1.0 = off).
        seed: RNG seed.

    Returns:
        The generated MP4, a chart of the predicted state trajectory, a table of
        the predicted state values, and a run-info string.
    """
    if image is None:
        raise gr.Error("Please provide a first frame.")

    num_frames = int(num_frames)
    n_rows, n_slots = slot_layout(num_frames)
    slots = parse_action_slots(actions_text, n_slots)

    state = {
        "timer": max(0, min(int(timer), STATE_MAX["timer"])),
        "hp1": max(0, min(int(hp1), STATE_MAX["hp1"])),
        "hp2": max(0, min(int(hp2), STATE_MAX["hp2"])),
        "meter1": max(0, min(int(meter1), STATE_MAX["meter1"])),
        "meter2": max(0, min(int(meter2), STATE_MAX["meter2"])),
    }

    parquet_path = tempfile.NamedTemporaryFile(suffix=".parquet", delete=False).name
    build_action_parquet(slots, n_rows, state, parquet_path)

    started = time.perf_counter()
    out = pipe(
        image=image,
        actions_parquet=parquet_path,
        state_parquet=parquet_path,
        state_sampling="end",
        prompt=(prompt or "SF3 Game.").strip(),
        negative_prompt=NEG_PROMPT,
        num_frames=num_frames,
        num_inference_steps=int(num_inference_steps),
        cfg_scale=float(cfg_scale),
        state_cfg_scale=float(state_cfg_scale),
        action_cfg_scale=float(action_cfg_scale),
        height=HEIGHT,
        width=WIDTH,
        seed=int(seed),
    )
    elapsed = time.perf_counter() - started

    frames = out.frames[0]
    video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    save_mp4(frames, video_path, fps=FPS)

    rows = state_table(out.state)
    chart_path = state_chart(rows)

    pressed = " | ".join("+".join(s) if s else "-" for s in slots)
    info = (
        f"{len(frames)} frames @ {FPS} fps ({len(frames) / FPS:.2f}s), {WIDTH}x{HEIGHT}, "
        f"{int(num_inference_steps)} steps, seed {int(seed)} — inference {elapsed:.1f}s\n"
        f"actions used ({n_slots} x 0.5s slots): {pressed}"
    )
    try:
        os.remove(parquet_path)
    except OSError:
        pass
    return video_path, chart_path, rows, info


# --------------------------------------------------------------------------- #
# UI
# --------------------------------------------------------------------------- #

EXAMPLES = [
    [
        "examples/01_macro_success.png",
        "D | UP+RIGHT | D | UP+RIGHT | UP+LEFT | A | DOWN | C | UP+LEFT | DOWN",
        "NPC: Active_Behavior(Kazegiri: A rising high kick used primarily to intercept aerial opponents.), Passive_Behavior(Standing Block: Mitigates damage from high and mid-level attacks while standing.; Idle: A neutral stationary stance where no action is taken.), Strategy(Passive Guarding: Remains stationary while utilizing standing or crouching blocks to mitigate incoming damage.)",
        45, 35, 111, 104, 96,
    ],
    [
        "examples/02_macro_success.png",
        "D | D | UP+LEFT | - | A | A | DOWN | UP+RIGHT | Z | LEFT",
        "NPC: Active_Behavior(Walk Left: Moves horizontally to the left along the ground.; Jump Backward: Leaps away from the opponent to create distance.), Passive_Behavior(Take Hit: Sustains damage from an opponent's attack.), Strategy(Spacing Control: Maintains an optimal distance from the opponent through movement.)",
        62, 70, 100, 104, 76,
    ],
    [
        "examples/03_result_win.png",
        "D | D | D | Y | RIGHT | X | D | D | D | D",
        "NPC: Active_Behavior(N/A), Passive_Behavior(Take Hit: Sustains damage from an opponent's attack.; Idle: A neutral stationary stance where no action is taken.), Strategy(Defeated: Health has been depleted and the round is lost.)",
        56, 85, 2, 104, 76,
    ],
    [
        "examples/04_result_win.png",
        "Z | LEFT | A | UP+LEFT | UP | DOWN | D | Z | - | UP",
        "NPC: Active_Behavior(N/A), Passive_Behavior(Take Hit: Sustains damage from an opponent's attack.), Strategy(Defeated: Health has been depleted and the round is lost.)",
        35, 11, 13, 76, 96,
    ],
    [
        "examples/05_result_lose.png",
        "D | D | RIGHT | RIGHT | LEFT | D | Z | - | B | Y",
        "NPC: Active_Behavior(Throw: A close-range grab that bypasses blocking.; Standing Attack: A basic attack performed from a standing position.), Passive_Behavior(Idle: A neutral stationary stance where no action is taken.), Strategy(Aggressive Pressure: Continuously attacks to force the opponent into a defensive state.)",
        35, 6, 88, 31, 96,
    ],
    [
        "examples/06_result_lose.png",
        "A | RIGHT | UP+RIGHT | B | DOWN | Z | UP+RIGHT | - | Y | LEFT",
        "NPC: Active_Behavior(N/A), Passive_Behavior(Idle: A neutral stationary stance where no action is taken.), Strategy(Victorious: The opponent's health has been depleted and the round is won.)",
        42, 12, 75, 14, 96,
    ],
    [
        "examples/07_normal.png",
        "DOWN | C | B | Z | LEFT | LEFT | UP | A | UP | B",
        "NPC: Active_Behavior(Crouch: Lowers stance to the ground to duck under high attacks.; Walk Right: Moves horizontally to the right along the ground.), Passive_Behavior(Idle: A neutral stationary stance where no action is taken.), Strategy(Spacing Control: Maintains an optimal distance from the opponent through movement.)",
        95, 160, 160, 15, 0,
    ],
    [
        "examples/08_normal.png",
        "UP | UP+LEFT | RIGHT | X | C | B | A | DOWN | - | RIGHT",
        "NPC: Active_Behavior(N/A), Passive_Behavior(Idle: A neutral stationary stance where no action is taken.), Strategy(Neutral Game: Observes the opponent while maintaining a safe position.)",
        90, 160, 152, 18, 0,
    ],
]

HEAD = """
# 🎮 StatePlay — State-Aware Game World Model

Roll out **Street Fighter III** gameplay from a single frame *and* read out the
game state the model believes it is producing — health, super meters and the
round timer are predicted jointly with the pixels, so the video stays consistent
with the game's mechanics.

[paper](https://huggingface.co/papers/2607.26754) · [project page](https://jimntu.github.io/stateplay_page/) · [code](https://github.com/Jimntu/StatePlay) · [model](https://huggingface.co/onepiece1999/StatePlay)
"""

ACTION_HELP = f"""
**Action sequence** — one slot per `|`, each slot held for **0.5 s** (10 frames @ 20 fps).
Press several buttons at once with `+`, use `-` for no input.

Buttons: `UP` `DOWN` `LEFT` `RIGHT` (stick) · `{'` `'.join(SF3_BUTTON_COLS[4:])}` (attack / special channels).
A 101-frame rollout uses the first **10** slots.
"""

CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="StatePlay") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(HEAD)

        with gr.Row():
            with gr.Column(scale=1):
                image = gr.Image(label="First frame", type="pil", height=260)
                actions = gr.Textbox(
                    label="Action sequence",
                    value="D | UP+RIGHT | D | UP+RIGHT | UP+LEFT | A | DOWN | C | UP+LEFT | DOWN",
                    lines=2,
                )
                gr.Markdown(ACTION_HELP)
                prompt = gr.Textbox(
                    label="Prompt (NPC behaviour / strategy)",
                    value="SF3 Game.",
                    lines=3,
                )
                gr.Markdown("**Initial game state** (conditioning — the model predicts the rest)")
                with gr.Row():
                    timer = gr.Slider(0, 99, value=45, step=1, label="Timer")
                    hp1 = gr.Slider(0, 160, value=160, step=1, label="P1 HP")
                    hp2 = gr.Slider(0, 160, value=160, step=1, label="P2 HP")
                with gr.Row():
                    meter1 = gr.Slider(0, 104, value=0, step=1, label="P1 meter")
                    meter2 = gr.Slider(0, 96, value=0, step=1, label="P2 meter")

                run = gr.Button("Generate rollout", variant="primary")

                with gr.Accordion("Advanced settings", open=False):
                    num_frames = gr.Dropdown(
                        [41, 61, 81, 101], value=101, label="Frames (20 fps)"
                    )
                    steps = gr.Slider(10, 40, value=30, step=1, label="Denoising steps")
                    cfg = gr.Slider(1.0, 10.0, value=5.0, step=0.1, label="Text CFG")
                    action_cfg = gr.Slider(
                        1.0, 5.0, value=1.0, step=0.1, label="Action CFG (1.0 = off)"
                    )
                    state_cfg = gr.Slider(
                        1.0, 10.0, value=1.0, step=0.1, label="State CFG (1.0 = off)"
                    )
                    seed = gr.Number(value=2, precision=0, label="Seed")

            with gr.Column(scale=1):
                video = gr.Video(label="Generated rollout", autoplay=True, height=300)
                chart = gr.Image(label="Predicted state trajectory", height=420)
                info = gr.Textbox(label="Run info", lines=3)
                with gr.Accordion("Predicted state values", open=False):
                    table = gr.Dataframe(
                        headers=["time_s", *STATE_COLUMNS],
                        datatype=["number"] * 6,
                        label="one row per latent frame (0.2 s)",
                    )

        inputs = [
            image, actions, prompt, timer, hp1, hp2, meter1, meter2,
            num_frames, steps, cfg, action_cfg, state_cfg, seed,
        ]
        outputs = [video, chart, table, info]

        run.click(generate, inputs=inputs, outputs=outputs, api_name="generate")

        gr.Examples(
            examples=EXAMPLES,
            inputs=[image, actions, prompt, timer, hp1, hp2, meter1, meter2],
            outputs=outputs,
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
            label="Held-out StatePlay clips (first frame, real controller log, real initial state)",
        )

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