"""
Gradio Space for TinyWorlds — Autoregressive World Model demo.
Deploy at: https://huggingface.co/spaces
"""
import re
import random
import tempfile
from pathlib import Path

import cv2
import gradio as gr
import numpy as np
import spaces
import torch
from einops import repeat
from huggingface_hub import hf_hub_download, list_repo_files

BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = BASE_DIR / "data"
MODEL_CACHE_DIR = BASE_DIR / "model_cache"
DATA_DIR.mkdir(exist_ok=True)
MODEL_CACHE_DIR.mkdir(exist_ok=True)

MODELS_REPO_ID = "AlmondGod/tinyworlds-models"
DATASET_REPO_ID = "AlmondGod/tinyworlds"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
FALLBACK_SUITE = "sonic"
CONTEXT_WINDOW = 4  # match training context_length=4 for in-distribution inference
UPSCALE = 4  # render 128-px frames at 4× for crisp pixel art

# Per-game config: dataset key → (suite name on HF, h5 filename, display label)
GAME_CONFIG: dict[str, dict] = {
    "SONIC":         {"suite": "sonic",        "h5": "sonic_frames.h5",        "label": "Sonic the Hedgehog"},
    "PONG":          {"suite": "pong",         "h5": "pong_frames.h5",         "label": "Pong"},
    "ZELDA":         {"suite": "zelda",        "h5": "zelda_frames.h5",        "label": "Zelda OoT 2D"},
    "POLE_POSITION": {"suite": "pole_position","h5": "pole_position_frames.h5","label": "Pole Position"},
    "PICODOOM":      {"suite": "picodoom",     "h5": "picodoom_frames.h5",     "label": "PicoDoom"},
}

# ---------------------------------------------------------------------------
# Asset download helpers
# ---------------------------------------------------------------------------

def _download_suite_checkpoints(suite: str) -> Path:
    """Download checkpoint files for a suite from HF. Returns suite_dir."""
    suite_dir = MODEL_CACHE_DIR / suite
    suite_dir.mkdir(exist_ok=True)

    all_files = list(list_repo_files(MODELS_REPO_ID, repo_type="model"))
    print(f"[HF] {len(all_files)} files in {MODELS_REPO_ID}:")
    for f in all_files[:30]:
        print(f"  {f}")

    suite_files = [
        f for f in all_files
        if f.startswith(f"{suite}/") and (f.endswith(".pt") or f.endswith(".pth"))
    ]

    if not suite_files:
        raise FileNotFoundError(
            f"No checkpoints for suite '{suite}' in {MODELS_REPO_ID}. "
            f"Available files: {all_files[:10]}"
        )

    for fname in suite_files:
        local = MODEL_CACHE_DIR / fname
        if not local.exists():
            hf_hub_download(
                repo_id=MODELS_REPO_ID,
                filename=fname,
                repo_type="model",
                local_dir=str(MODEL_CACHE_DIR),
                local_dir_use_symlinks=False,
            )
    return suite_dir


def _find_ckpt_pth(suite_dir: Path, model_name: str) -> Path:
    """Find the .pth file with the highest step for a model in suite_dir."""
    aliases = {
        "video_tokenizer": ["video_tokenizer"],
        "latent_actions":  ["latent_actions", "lam", "actions", "action_tokenizer"],
        "dynamics":        ["dynamics"],
    }
    names = aliases.get(model_name, [model_name])
    candidates = [
        f for f in suite_dir.rglob("*.pth")
        if any(n in f.name for n in names)
    ]
    if not candidates:
        raise FileNotFoundError(f"No .pth checkpoint for '{model_name}' in {suite_dir}")
    def _step(p: Path) -> int:
        m = re.search(r"_step_(\d+)", str(p))
        return int(m.group(1)) if m else 0
    return max(candidates, key=_step)


def _pth_to_ckpt_dir(pth_path: Path, model_name: str) -> Path:
    """
    Convert a single .pth file into the directory format that
    load_*_from_checkpoint expects (model_state_dict.pt + state.pt).
    Returns the created directory.
    """
    ckpt_dir = pth_path.parent / f"{model_name}_extracted"
    model_file = ckpt_dir / "model_state_dict.pt"
    state_file = ckpt_dir / "state.pt"
    if model_file.exists() and state_file.exists():
        return ckpt_dir  # already extracted

    ckpt_dir.mkdir(exist_ok=True)
    data = torch.load(str(pth_path), map_location="cpu", weights_only=False)

    # Handle common single-file checkpoint formats
    if isinstance(data, dict):
        if "model_state_dict" in data:
            model_sd = data["model_state_dict"]
            cfg = data.get("config", {}) or {}
        elif "state_dict" in data:
            model_sd = data["state_dict"]
            cfg = data.get("config", {}) or {}
        elif "model" in data:
            model_sd = data["model"]
            cfg = data.get("config", {}) or {}
        else:
            # Assume the dict IS the state dict directly
            model_sd = data
            cfg = {}
    else:
        raise ValueError(f"Unexpected type in {pth_path}: {type(data)}")

    torch.save(model_sd, str(model_file))
    torch.save({"config": cfg, "step": None, "timestamp": ""}, str(state_file))
    return ckpt_dir


def _ensure_h5(h5_filename: str) -> None:
    dest = DATA_DIR / h5_filename
    if dest.exists():
        return
    print(f"Downloading {h5_filename}…")
    hf_hub_download(
        repo_id=DATASET_REPO_ID,
        filename=h5_filename,
        repo_type="dataset",
        local_dir=str(DATA_DIR),
        local_dir_use_symlinks=False,
    )


# ---------------------------------------------------------------------------
# Model cache
# ---------------------------------------------------------------------------

_model_cache: dict = {}


def get_models(suite: str):
    if suite in _model_cache:
        return _model_cache[suite]

    actual_suite = suite
    try:
        suite_dir = _download_suite_checkpoints(suite)
    except FileNotFoundError:
        print(f"No checkpoints for '{suite}', falling back to '{FALLBACK_SUITE}'")
        actual_suite = FALLBACK_SUITE
        suite_dir = _download_suite_checkpoints(FALLBACK_SUITE)

    vt_pth  = _find_ckpt_pth(suite_dir, "video_tokenizer")
    lam_pth = _find_ckpt_pth(suite_dir, "latent_actions")
    dyn_pth = _find_ckpt_pth(suite_dir, "dynamics")
    vt_dir  = _pth_to_ckpt_dir(vt_pth,  "video_tokenizer")
    lam_dir = _pth_to_ckpt_dir(lam_pth, "latent_actions")
    dyn_dir = _pth_to_ckpt_dir(dyn_pth, "dynamics")

    from utils.inference_utils import load_models  # noqa: PLC0415
    vt, lam, dyn = load_models(str(vt_dir), str(lam_dir), str(dyn_dir), DEVICE, use_actions=True)
    _model_cache[suite] = (vt, lam, dyn)
    print(f"Models for '{actual_suite}' loaded on {DEVICE}.")
    return _model_cache[suite]


# ---------------------------------------------------------------------------
# Dataset cache
# ---------------------------------------------------------------------------

_dataset_cache: dict = {}


def get_val_dataset(game: str, num_frames: int):
    key = (game, num_frames)
    if key in _dataset_cache:
        return _dataset_cache[key]

    h5_file = GAME_CONFIG[game]["h5"]
    _ensure_h5(h5_file)

    from datasets.data_utils import load_data_and_data_loaders  # noqa: PLC0415
    _, val_data, _, _, _ = load_data_and_data_loaders(
        dataset=game, batch_size=1, num_frames=num_frames, preload_ratio=0.15
    )
    _dataset_cache[key] = val_data
    return val_data


# ---------------------------------------------------------------------------
# Rendering helpers
# ---------------------------------------------------------------------------

def _to_rgb_u8(tensor_chw: torch.Tensor) -> np.ndarray:
    """[-1,1] C×H×W tensor → uint8 H×W×C numpy RGB."""
    arr = ((tensor_chw.float().clamp(-1, 1) + 1) / 2 * 255).byte()
    return arr.permute(1, 2, 0).cpu().numpy()


def _nearest(img: np.ndarray, scale: int = UPSCALE) -> np.ndarray:
    """Nearest-neighbour upscale — preserves pixel-art look."""
    H, W = img.shape[:2]
    return cv2.resize(img, (W * scale, H * scale), interpolation=cv2.INTER_NEAREST)


def _border(img: np.ndarray, color: tuple[int, int, int], thickness: int = 6) -> np.ndarray:
    """Draw a solid colored border around a frame to indicate its type."""
    out = img.copy()
    out[:thickness, :]  = color
    out[-thickness:, :] = color
    out[:, :thickness]  = color
    out[:, -thickness:] = color
    return out


_CTX_COLOR = (160, 160, 180)   # muted grey-blue  → context frames
_GEN_COLOR = (80,  200, 120)   # bright green      → generated frames


# ---------------------------------------------------------------------------
# Auto-generate inference
# ---------------------------------------------------------------------------

@spaces.GPU(duration=120)
def run_inference(game: str, n_steps: int, temperature: float, use_random_actions: bool) -> str:
    prediction_horizon = 1
    frames_to_load = CONTEXT_WINDOW + n_steps * prediction_horizon

    suite = GAME_CONFIG[game]["suite"]
    print(f"[{game}] Loading models (suite={suite})…")
    video_tokenizer, latent_action_model, dynamics_model = get_models(suite)

    print(f"[{game}] Loading dataset sample…")
    val_data = get_val_dataset(game, frames_to_load)

    n = len(val_data)
    idx = random.randint(max(0, n // 5), min(n - 1, 4 * n // 5))
    gt_frames = val_data[idx][0].unsqueeze(0).to(DEVICE)

    context_frames = gt_frames[:, :CONTEXT_WINDOW]
    generated_frames = context_frames.clone()

    n_actions = latent_action_model.quantizer.codebook_size
    inferred_actions: list = []

    with torch.no_grad():
        for i in range(n_steps):
            print(f"[{game}] Generating frame {i + 1}/{n_steps}…")

            ctx = generated_frames[:, -CONTEXT_WINDOW:]
            video_indices = video_tokenizer.tokenize(ctx)
            video_latents = video_tokenizer.quantizer.get_latents_from_indices(video_indices)

            if use_random_actions:
                from utils.inference_utils import sample_random_action  # noqa: PLC0415
                action_idx = sample_random_action(n_actions)
                inferred_actions.append(action_idx)
                recent = inferred_actions[-CONTEXT_WINDOW:]
                recent_t = repeat(torch.tensor(recent, device=DEVICE), "i -> 1 i")
                action_latent = latent_action_model.quantizer.get_latents_from_indices(recent_t)
                if len(recent) < CONTEXT_WINDOW:
                    pad_len = CONTEXT_WINDOW - len(recent) + 1
                    pad_enc = latent_action_model.encode(ctx[:, :pad_len])
                    pad_q   = latent_action_model.quantizer(pad_enc)
                    action_latent = torch.cat([pad_q, action_latent], dim=1)
            else:
                # Infer past actions from context and extrapolate for next step
                act_past = latent_action_model.encode(ctx)          # [1, CONTEXT_WINDOW-1, A]
                action_latent = torch.cat([act_past, act_past[:, -1:]], dim=1)  # [1, CONTEXT_WINDOW, A]

            def idx_to_latents(idx):
                return video_tokenizer.quantizer.get_latents_from_indices(idx, dim=-1)

            next_latents = dynamics_model.forward_inference(
                context_latents=video_latents,
                prediction_horizon=prediction_horizon,
                num_steps=10,
                index_to_latents_fn=idx_to_latents,
                conditioning=action_latent,
                temperature=temperature,
            )
            next_frames = video_tokenizer.detokenize(next_latents)
            generated_frames = torch.cat([generated_frames, next_frames[:, -prediction_horizon:]], dim=1)

    print("Rendering…")
    gen = torch.clamp((generated_frames.detach().cpu() + 1) / 2, 0, 1)
    gt  = torch.clamp((gt_frames[:, :gen.shape[1]].detach().cpu() + 1) / 2, 0, 1)

    B, T, C, H, W = gen.shape
    out_W, out_H = W * 2 * UPSCALE, H * UPSCALE

    tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    out_path = tmp.name
    tmp.close()

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(out_path, fourcc, 4, (out_W, out_H))
    for t in range(T):
        gt_f = (gt [0, t].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
        gn_f = (gen[0, t].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
        combined = np.concatenate([gt_f, gn_f], axis=1)
        combined = cv2.resize(combined, (out_W, out_H), interpolation=cv2.INTER_NEAREST)
        writer.write(cv2.cvtColor(combined, cv2.COLOR_RGB2BGR))
    writer.release()

    print("Done!")
    return out_path


# ---------------------------------------------------------------------------
# Interactive play
# ---------------------------------------------------------------------------

# D-pad / face buttons → action token index (wraps with % n_actions at runtime)
_BUTTON_TO_IDX = {"←": 0, "→": 1, "↑": 2, "↓": 3, "A": 4, "B": 5}


@spaces.GPU(duration=30)
def init_game(game: str):
    suite = GAME_CONFIG[game]["suite"]
    vt, lam, dyn = get_models(suite)

    val_data = get_val_dataset(game, CONTEXT_WINDOW + 1)
    # Sample from the middle 60 % of the dataset to skip title / loading screens
    n = len(val_data)
    lo, hi = max(0, n // 5), min(n - 1, 4 * n // 5)
    idx = random.randint(lo, hi)
    gt = val_data[idx][0].unsqueeze(0).to(DEVICE)
    context = gt[:, :CONTEXT_WINDOW]  # [1, CONTEXT_WINDOW, C, H, W]

    # Quick sanity-check: tokenize → detokenize context and show reconstruction
    with torch.no_grad():
        idx_check  = vt.tokenize(context)
        lat_check  = vt.quantizer.get_latents_from_indices(idx_check)
        recon      = vt.detokenize(lat_check)
        print(f"[init] context range [{context.min():.2f}, {context.max():.2f}]")
        print(f"[init] recon   range [{recon.min():.2f}, {recon.max():.2f}]")

    # Seed gallery: context frames + video-tokenizer reconstruction (sanity check)
    gallery: list[tuple[np.ndarray, str]] = [
        (_border(_nearest(_to_rgb_u8(context[0, i])), _CTX_COLOR), f"Context {i + 1}")
        for i in range(CONTEXT_WINDOW)
    ]
    # Append VT reconstruction of last context frame to show tokenizer quality
    recon_img = _border(_nearest(_to_rgb_u8(recon[0, -1])), (200, 160, 50))  # gold border
    gallery.append((recon_img, "VT recon"))

    state = {
        "game": game,
        "context": context.cpu().numpy(),
        "action_history": [],
        "n_actions": lam.quantizer.codebook_size,
        "gallery": gallery,
        "step": 0,
    }
    label = GAME_CONFIG[game]["label"]
    return gallery, state, f"▶ {label} — Context + VT reconstruction shown. Press Continue to generate."


@spaces.GPU(duration=60)
def step_game(btn_label: str, state: dict, temperature: float, maskgit_steps: int):
    if state is None:
        return None, None, "⚠ Click 'New game' first."

    game      = state["game"]
    n_actions = state["n_actions"]
    suite     = GAME_CONFIG[game]["suite"]
    vt, lam, dyn = get_models(suite)

    context = torch.from_numpy(state["context"].copy()).to(DEVICE)
    hist    = state["action_history"]

    # "▶ Continue" runs the dynamics model unconditioned — most coherent output
    use_action = btn_label != "▶"
    action_idx  = _BUTTON_TO_IDX.get(btn_label, 0) % n_actions if use_action else 0

    with torch.no_grad():
        ctx = context[:, -CONTEXT_WINDOW:]

        # lam.encode: for T frames → [1, T-1, A] inferred actions between consecutive frames
        # This matches training: latent_action_model.encode(x) → [B, T-1, A]
        act_past = lam.encode(ctx)   # [1, CONTEXT_WINDOW-1, A]

        if use_action:
            # Past T-1 inferred actions + user-chosen action for the predicted frame
            cur_t = torch.tensor([[action_idx]], device=DEVICE)      # [1, 1]
            cur_q = lam.quantizer.get_latents_from_indices(cur_t)    # [1, 1, A]
            action_latent = torch.cat([act_past, cur_q], dim=1)      # [1, CONTEXT_WINDOW, A]
        else:
            # Free: repeat the last inferred action (between frame T-1→T) as the next step
            action_latent = torch.cat([act_past, act_past[:, -1:]], dim=1)  # [1, CONTEXT_WINDOW, A]

        video_indices = vt.tokenize(ctx)
        video_latents = vt.quantizer.get_latents_from_indices(video_indices)

        print(f"[play] ctx={ctx.shape} video_latents={video_latents.shape} action_latent={action_latent.shape}")

        def idx_to_latents(idx):
            return vt.quantizer.get_latents_from_indices(idx, dim=-1)

        next_latents = dyn.forward_inference(
            context_latents=video_latents,
            prediction_horizon=1,
            num_steps=maskgit_steps,
            index_to_latents_fn=idx_to_latents,
            conditioning=action_latent,
            temperature=temperature,
        )
        next_frame = vt.detokenize(next_latents)[:, -1:]  # [1, 1, C, H, W]
        print(f"[play] next_frame min={next_frame.min():.3f} max={next_frame.max():.3f} mean={next_frame.mean():.3f}")

    new_context = torch.cat([context, next_frame], dim=1)[:, -CONTEXT_WINDOW:]
    step     = state["step"] + 1
    new_img  = _border(_nearest(_to_rgb_u8(next_frame[0, 0])), _GEN_COLOR)
    caption  = f"Step {step}  {'(free)' if not use_action else btn_label}"
    new_gallery = state["gallery"] + [(new_img, caption)]

    new_state = {
        "game": game,
        "context": new_context.cpu().numpy(),
        "action_history": (hist + [action_idx])[-10:] if use_action else hist,
        "n_actions": n_actions,
        "gallery": new_gallery,
        "step": step,
    }
    status = f"Step {step} — {'unconditioned (free)' if not use_action else f'{btn_label} token {action_idx}/{n_actions}'}"
    return new_gallery, new_state, status


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

GAME_CHOICES = [(cfg["label"], key) for key, cfg in GAME_CONFIG.items()]

_CSS = """
/* ── header ── */
#tw-header { text-align:center; padding:1.2rem 0 0.4rem; }
#tw-header h1 { font-size:2rem; font-weight:700; margin-bottom:0.2rem; }
#tw-header p  { color:var(--body-text-color-subdued); margin:0; font-size:0.95rem; }
/* ── pixel art rendering ── */
#play-frame img { image-rendering:pixelated; image-rendering:crisp-edges; }
/* ── d-pad ── */
.dpad-btn button { min-width:54px !important; min-height:54px !important;
                   font-size:1.4rem !important; font-weight:900 !important;
                   padding:0 !important; border-radius:10px !important; }
.dpad-space { min-width:54px; }
/* ── face buttons ── */
.face-btn button { min-width:60px !important; min-height:42px !important;
                   font-size:1rem !important; font-weight:700 !important;
                   border-radius:50px !important; }
/* ── misc ── */
.label-note { font-size:0.78rem; color:var(--body-text-color-subdued); margin-top:0.3rem; }
#gen-btn, #newgame-btn { margin-top:0.5rem; }
"""

with gr.Blocks(title="TinyWorlds — World Model Demo", theme=gr.themes.Soft(), css=_CSS) as demo:

    # ── Header ──────────────────────────────────────────────────────────────
    gr.Markdown(
        """
<div id="tw-header">
<h1>🎮 TinyWorlds</h1>
<p>Autoregressive World Model &nbsp;·&nbsp;
<a href="https://arxiv.org/pdf/2402.15391" target="_blank">Genie architecture</a> &nbsp;·&nbsp;
<a href="https://github.com/AlmondGod/tinyworlds" target="_blank">GitHub</a></p>
</div>
        """,
        sanitize_html=False,
    )

    # ── Tabs ────────────────────────────────────────────────────────────────
    with gr.Tabs():

        # ── Tab 1: Auto-generate ──────────────────────────────────────────
        with gr.Tab("🎬 Auto-generate"):
            gr.Markdown(
                "The model generates a sequence of frames from a random starting point. "
                "**Left**: ground truth · **Right**: world-model prediction."
            )
            with gr.Row(equal_height=False):
                with gr.Column(scale=1, min_width=260):
                    game_sel = gr.Radio(choices=GAME_CHOICES, value="SONIC", label="Game")
                    n_steps  = gr.Slider(2, 15, value=8, step=1, label="Frames to generate",
                                         info="More frames = longer video, slower")
                    with gr.Accordion("Advanced", open=False):
                        temperature = gr.Slider(0.0, 1.0, value=0.5, step=0.05,
                                                label="Temperature",
                                                info="0 = deterministic · 1 = max diversity")
                        use_actions = gr.Checkbox(value=True, label="Condition on random actions",
                                                  info="Injects random action tokens into the dynamics model")
                    gen_btn = gr.Button("Generate", variant="primary", size="lg", elem_id="gen-btn")
                    gr.HTML('<p class="label-note">⏱ First run downloads models + dataset (~1–2 GB). Subsequent runs use cache.</p>')

                with gr.Column(scale=2):
                    video_out = gr.Video(label="Ground Truth (left)  ·  Predicted (right)",
                                         autoplay=True, show_download_button=True)

            gen_btn.click(fn=run_inference,
                          inputs=[game_sel, n_steps, temperature, use_actions],
                          outputs=video_out, api_name=False)

        # ── Tab 2: Play ───────────────────────────────────────────────────
        with gr.Tab("🕹 Play"):
            gr.Markdown(
                "Pick an action each turn — each button press adds a new frame to the grid. "
                "Actions are **unsupervised tokens**; explore to discover what each button does!"
            )
            with gr.Row(equal_height=False):

                # Left: controls
                with gr.Column(scale=1, min_width=220):
                    game_play   = gr.Radio(choices=GAME_CHOICES, value="SONIC", label="Game")
                    newgame_btn = gr.Button("▶ New game", variant="primary", elem_id="newgame-btn")
                    play_status = gr.Markdown("_Press **New game** to start._")

                    # Primary: unconditioned continuation
                    free_btn = gr.Button("▶  Continue (no action)", variant="primary", size="lg")
                    gr.HTML('<p class="label-note" style="margin-bottom:0.8rem">Most coherent — model predicts the next frame freely from context.</p>')

                    gr.Markdown("#### Action buttons _(experimental)_")
                    gr.HTML('<p class="label-note" style="margin-bottom:0.4rem">Condition on an unsupervised action token.<br>Explore to discover what each one does.</p>')

                    # D-pad
                    with gr.Column():
                        with gr.Row():
                            gr.HTML('<div class="dpad-space"></div>')
                            up_btn    = gr.Button("↑", elem_classes=["dpad-btn"])
                            gr.HTML('<div class="dpad-space"></div>')
                        with gr.Row():
                            left_btn  = gr.Button("←", elem_classes=["dpad-btn"])
                            gr.HTML('<div class="dpad-space"></div>')
                            right_btn = gr.Button("→", elem_classes=["dpad-btn"])
                        with gr.Row():
                            gr.HTML('<div class="dpad-space"></div>')
                            down_btn  = gr.Button("↓", elem_classes=["dpad-btn"])
                            gr.HTML('<div class="dpad-space"></div>')

                    gr.Markdown("&nbsp;")
                    with gr.Row():
                        a_btn = gr.Button("A", elem_classes=["face-btn"])
                        b_btn = gr.Button("B", elem_classes=["face-btn"])

                    with gr.Accordion("Quality settings", open=True):
                        play_temp = gr.Slider(
                            0.0, 1.0, value=0.1, step=0.05,
                            label="Temperature",
                            info="Lower = sharper frames; 0.0–0.2 recommended",
                        )
                        play_steps = gr.Slider(
                            5, 40, value=20, step=5,
                            label="MaskGIT steps",
                            info="More steps = sharper frames, slower",
                        )

                # Right: frame grid
                with gr.Column(scale=3):
                    play_gallery = gr.Gallery(
                        label="World model — frame history",
                        columns=4,
                        rows=2,
                        height="auto",
                        object_fit="contain",
                        elem_id="play-frame",
                        show_download_button=False,
                    )

            play_state = gr.State(None)

            # Wire up buttons
            newgame_btn.click(fn=init_game, inputs=[game_play],
                              outputs=[play_gallery, play_state, play_status], api_name=False)

            for label, btn in [("▶", free_btn),
                                ("←", left_btn), ("→", right_btn),
                                ("↑", up_btn),   ("↓", down_btn),
                                ("A", a_btn),    ("B", b_btn)]:
                btn.click(fn=lambda s, t, st, l=label: step_game(l, s, t, st),
                          inputs=[play_state, play_temp, play_steps],
                          outputs=[play_gallery, play_state, play_status],
                          api_name=False)

        # ── How it works ──────────────────────────────────────────────────
        with gr.Tab("📖 How it works"):
            gr.Markdown("""
## Architecture

TinyWorlds has **three learned components** trained end-to-end on raw gameplay video:

| Component | Role |
|---|---|
| **Video Tokenizer** | FSQ VAE that compresses each frame into a grid of discrete tokens |
| **Action Tokenizer** | FSQ VAE that infers a latent *action* token between consecutive frames — **no action labels needed** |
| **Dynamics Model** | Transformer that predicts next-frame tokens from past frame+action tokens (MaskGIT masked diffusion) |

## Inference loop

1. Tokenize the context window of frames → video tokens
2. (Optional) inject an action token for the next step
3. Dynamics model unmasks next-frame tokens iteratively (10 steps, exponential schedule)
4. Video tokenizer decodes the tokens back to pixels

## Why it's slow

Each button press runs the full dynamics model (10 MaskGIT steps) on GPU. Expect 2–5 s per frame.
Real-time play would require a much smaller / distilled model.

## Action tokens

The action tokenizer learns action representations **unsupervised**. Token 0 might mean "run right",
token 1 might mean "jump" — but this varies by game and training run.
Explore the buttons to discover the mapping!
            """)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0")