"""ForgeWM — few-step, action-conditioned Minecraft world model.

Gradio demo for `ForgeWM/ForgeWM` (paper: "ForgeWM: Progressive Causal Training
for Few-Step Action-Conditioned Video World Models").

Faithful to the repo's own `inference.py` / `pipeline/causal_inference.py`:
same 352x640 resolution, 3-latent-frame causal blocks, sliding window of 6
latent frames, the released `warp_denoising_step` schedules, and First-Frame
Enhancement for the 1-/2-step students.
"""

import os

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

import spaces  # noqa: E402  (must precede torch)

import gc  # noqa: E402
import shutil  # noqa: E402
import tempfile  # noqa: E402
import time  # noqa: E402
from typing import List, Tuple  # noqa: E402

import gradio as gr  # noqa: E402
import imageio.v2 as imageio  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
import torch.nn.functional as F  # noqa: E402
from huggingface_hub import hf_hub_download  # noqa: E402
from omegaconf import OmegaConf  # noqa: E402
from PIL import Image  # noqa: E402

from pipeline import CausalInferencePipeline
from utils.wan_wrapper import WanVAEWrapper

# ────────────────────────────── constants ──────────────────────────────
HEIGHT, WIDTH = 352, 640
FPS = 12
VAE_TCR = 4                 # VAE temporal compression ratio
FRAMES_PER_BLOCK = 3        # latent frames per causal block
CAM_VALUE = 0.10            # camera delta magnitude, from inference.py

MINECRAFT_ACTIONS = [
    "forward", "back", "left", "right",
    "turn_right", "turn_left", "look_up", "look_down",
    "forward_turn_right", "random", "no_action",
]

BASE_REPO = "Skywork/Matrix-Game-2.0"
FORGEWM_REPO = "ForgeWM/ForgeWM"

VARIANTS = {
    "ForgeWM-4 · 4 steps": ("configs/stage3_dmd.yaml", "stage3/model.pt"),
    "ForgeWM-2 · 2 steps": ("configs/stage3_dmd_2step.yaml", "2step/model.pt"),
    "ForgeWM-1 · 1 step": ("configs/stage3_dmd_1step.yaml", "1step/model.pt"),
}
DEFAULT_VARIANT = "ForgeWM-4 · 4 steps"

# Measured ballpark on the ZeroGPU Blackwell card; used only to size the
# @spaces.GPU reservation.
SECONDS_PER_CHUNK = {
    "ForgeWM-4 · 4 steps": 1.5,
    "ForgeWM-2 · 2 steps": 1.0,
    "ForgeWM-1 · 1 step": 0.8,
}
FIXED_OVERHEAD = 22.0       # CLIP + VAE encode/decode + mp4 mux

CKPT_DIR = os.path.join(os.getcwd(), "ckpts", "MG2-base")


# ─────────────────────────── weight preparation ───────────────────────────
def _link(src: str, dst: str) -> None:
    if os.path.lexists(dst):
        os.remove(dst)
    os.symlink(src, dst)


def _purge(path: str) -> None:
    """Drop a hub blob from local disk once its tensors are in memory."""
    try:
        real = os.path.realpath(path)
        if os.path.isfile(real):
            os.remove(real)
        if os.path.islink(path):
            os.remove(path)
    except OSError as exc:  # pragma: no cover
        print(f"[disk] could not purge {path}: {exc}")


def _disk() -> str:
    total, used, free = shutil.disk_usage("/")
    return f"disk {used / 2**30:.1f}G used / {free / 2**30:.1f}G free"


os.makedirs(os.path.join(CKPT_DIR, "xlm-roberta-large"), exist_ok=True)

print(f"[setup] fetching Matrix-Game-2.0 base weights … ({_disk()})", flush=True)
_dit = hf_hub_download(BASE_REPO, "base_model/diffusion_pytorch_model.safetensors")
_link(_dit, os.path.join(CKPT_DIR, "diffusion_pytorch_model.safetensors"))
_link(hf_hub_download(BASE_REPO, "base_model/base_config.json"),
      os.path.join(CKPT_DIR, "base_config.json"))
_vae_path = hf_hub_download(BASE_REPO, "Wan2.1_VAE.pth")
_clip_path = hf_hub_download(
    BASE_REPO, "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth")
for _tok in ("sentencepiece.bpe.model", "special_tokens_map.json",
             "tokenizer.json", "tokenizer_config.json"):
    _link(hf_hub_download(BASE_REPO, f"xlm-roberta-large/{_tok}"),
          os.path.join(CKPT_DIR, "xlm-roberta-large", _tok))
print(f"[setup] base weights ready ({_disk()})", flush=True)

# ─────────────────────────────── models ───────────────────────────────
torch.set_grad_enabled(False)
DTYPE = torch.bfloat16

print("[setup] building VAE + CLIP …", flush=True)
VAE = WanVAEWrapper(
    vae_path=_vae_path,
    clip_checkpoint_path=_clip_path,
    clip_tokenizer_path=os.path.join(CKPT_DIR, "xlm-roberta-large"),
).eval()
VAE = VAE.to(device="cuda", dtype=DTYPE)
# `WanVAEWrapper.clip` is a plain object, not a submodule, so `.to()` above does
# not reach it.  Move it explicitly (in fp32, as the repo does) so ZeroGPU packs
# it with everything else instead of paying a 4.8 GB host copy per request.
VAE.clip.model = VAE.clip.model.to("cuda")
gc.collect()
print(f"[setup] VAE + CLIP on GPU ({_disk()})", flush=True)


def _build_variant(name: str) -> CausalInferencePipeline:
    cfg_path, ckpt_file = VARIANTS[name]
    config = OmegaConf.merge(OmegaConf.load("configs/default.yaml"),
                             OmegaConf.load(cfg_path))
    config.model_kwargs.model_name = CKPT_DIR

    pipe = CausalInferencePipeline(config, device=torch.device("cuda"), vae=VAE)

    ckpt = hf_hub_download(FORGEWM_REPO, ckpt_file)
    try:
        state = torch.load(ckpt, map_location="cpu", weights_only=True)
    except Exception:
        state = torch.load(ckpt, map_location="cpu", weights_only=False)
    gen_sd = state.get("generator", state.get("generator_ema", state)) \
        if isinstance(state, dict) else state
    fixed = {k.replace("._fsdp_wrapped_module.", ".")
              .replace("._checkpoint_wrapped_module.", "."): v
             for k, v in gen_sd.items()}
    missing, unexpected = pipe.generator.load_state_dict(fixed, strict=False)
    print(f"[{name}] loaded: missing={len(missing)} unexpected={len(unexpected)}",
          flush=True)
    del state, gen_sd, fixed
    gc.collect()
    _purge(ckpt)

    pipe.generator = pipe.generator.to(device="cuda", dtype=DTYPE).eval()
    pipe.vae = VAE
    gc.collect()
    print(f"[{name}] ready ({_disk()})", flush=True)
    return pipe


PIPELINES = {}
for _name in VARIANTS:
    try:
        PIPELINES[_name] = _build_variant(_name)
    except Exception as exc:  # keep the Space usable if one student fails
        print(f"[setup] FAILED to load {_name}: {exc!r}", flush=True)

if not PIPELINES:
    raise RuntimeError("No ForgeWM student could be loaded.")
if DEFAULT_VARIANT not in PIPELINES:
    DEFAULT_VARIANT = next(iter(PIPELINES))

# The base DiT safetensors were only needed to instantiate the architecture;
# every parameter has since been overwritten by the ForgeWM checkpoints.
_purge(os.path.join(CKPT_DIR, "diffusion_pytorch_model.safetensors"))
print(f"[setup] {len(PIPELINES)} variant(s) ready ({_disk()})", flush=True)


# ──────────────────────────── action scripting ────────────────────────────
def make_action(action_type: str, num_raw_frames: int) -> Tuple[torch.Tensor, torch.Tensor]:
    """Minecraft action palette (mouse_dim_in=2), verbatim from inference.py."""
    mouse = torch.zeros(1, num_raw_frames, 2)
    keyboard = torch.zeros(1, num_raw_frames, 6)

    if action_type == "forward":
        keyboard[:, :, 0] = 1.0
    elif action_type == "back":
        keyboard[:, :, 1] = 1.0
    elif action_type == "left":
        keyboard[:, :, 2] = 1.0
    elif action_type == "right":
        keyboard[:, :, 3] = 1.0
    elif action_type == "turn_right":
        mouse[:, :, 1] = CAM_VALUE
    elif action_type == "turn_left":
        mouse[:, :, 1] = -CAM_VALUE
    elif action_type == "look_up":
        mouse[:, :, 0] = CAM_VALUE
    elif action_type == "look_down":
        mouse[:, :, 0] = -CAM_VALUE
    elif action_type == "forward_turn_right":
        keyboard[:, :, 0] = 1.0
        mouse[:, :, 1] = CAM_VALUE
    elif action_type == "random":
        torch.manual_seed(42)
        mouse = (torch.rand(1, num_raw_frames, 2) - 0.5) * (2 * CAM_VALUE)
        keyboard[:, :, :4] = (torch.rand(1, num_raw_frames, 4) > 0.5).float()
    elif action_type == "no_action":
        pass
    else:
        raise gr.Error(f"Unknown action '{action_type}'.")
    return mouse, keyboard


def _chunk_start(chunk_index: int) -> int:
    """First raw (pixel) frame owned by causal block `chunk_index`.

    Latent frame 0 maps to raw frame 0; latent frame f>=1 maps to raw frames
    4f-3 … 4f.  A block spans 3 latent frames, so block c covers raw frames
    [12c-3, 12c+9) — and [0, 9) for c == 0.
    """
    return 0 if chunk_index == 0 else 12 * chunk_index - 3


def build_action_track(segments: List[Tuple[str, int]]) -> Tuple[torch.Tensor, torch.Tensor]:
    total_chunks = sum(n for _, n in segments)
    num_raw = _chunk_start(total_chunks)
    mouse = torch.zeros(1, num_raw, 2)
    keyboard = torch.zeros(1, num_raw, 6)
    chunk = 0
    for action, count in segments:
        for _ in range(count):
            lo, hi = _chunk_start(chunk), _chunk_start(chunk + 1)
            m, k = make_action(action, hi - lo)
            mouse[:, lo:hi] = m
            keyboard[:, lo:hi] = k
            chunk += 1
    return mouse, keyboard


# ───────────────────────────── preprocessing ─────────────────────────────
def load_reference_frame(image_path: str) -> torch.Tensor:
    """Aspect-preserving resize + centre crop to 352x640, scaled to [-1, 1]."""
    image = Image.open(image_path).convert("RGB")
    arr = torch.from_numpy(np.asarray(image)).permute(2, 0, 1)[None].float() / 255.0
    _, _, h, w = arr.shape
    if h / w > HEIGHT / WIDTH:
        new_w, new_h = WIDTH, max(HEIGHT, int(round(h * WIDTH / w)))
    else:
        new_h, new_w = HEIGHT, max(WIDTH, int(round(w * HEIGHT / h)))
    arr = F.interpolate(arr, size=(new_h, new_w), mode="bilinear", align_corners=False)
    top, left = (new_h - HEIGHT) // 2, (new_w - WIDTH) // 2
    arr = arr[:, :, top:top + HEIGHT, left:left + WIDTH]
    arr = (arr - 0.5) / 0.5
    return arr.unsqueeze(0).to(device="cuda", dtype=DTYPE)  # [1, 1, 3, H, W]


def build_conditional_dict(pipe, pixel, num_frames, mouse_cond, keyboard_cond):
    """MG2-style conditioning: CLIP context + first-frame latent + mask."""
    num_pixel_frames = (num_frames - 1) * VAE_TCR + 1
    visual_context = pipe.vae.encode_visual_context_from_pixels(pixel).to(DTYPE)

    first_frame = pixel[:, 0:1]
    pad = torch.zeros(1, num_pixel_frames - 1, 3, pixel.shape[3], pixel.shape[4],
                      device=pixel.device, dtype=DTYPE)
    padded = torch.cat([first_frame, pad], dim=1).permute(0, 2, 1, 3, 4)
    img_cond = pipe.vae.encode_to_latent(padded).to(DTYPE)

    _, _, _, h_lat, w_lat = img_cond.shape
    mask = torch.zeros(1, num_frames, 4, h_lat, w_lat,
                       device=pixel.device, dtype=DTYPE)
    mask[:, 0:1] = 1
    return {
        "visual_context": visual_context,
        "cond_concat": torch.cat([mask, img_cond], dim=2),
        "mouse_condition": mouse_cond.to(device="cuda", dtype=DTYPE),
        "keyboard_condition": keyboard_cond.to(device="cuda", dtype=DTYPE),
    }


# ──────────────────────────────── core ────────────────────────────────
@torch.no_grad()
def _rollout(image_path, variant, segments, seed):
    if image_path is None:
        raise gr.Error("Please provide a reference frame.")
    if variant not in PIPELINES:
        variant = DEFAULT_VARIANT
    segments = [(a, int(n)) for a, n in segments if int(n) > 0]
    total_chunks = sum(n for _, n in segments)
    if total_chunks < 1:
        raise gr.Error("Give at least one action segment a non-zero length.")
    if total_chunks > 12:
        raise gr.Error("Keep the total rollout at 12 chunks or fewer.")

    pipe = PIPELINES[variant]
    num_frames = total_chunks * FRAMES_PER_BLOCK

    pixel = load_reference_frame(image_path)
    mouse, keyboard = build_action_track(segments)
    cond = build_conditional_dict(pipe, pixel, num_frames, mouse, keyboard)

    torch.manual_seed(int(seed))
    noise = torch.randn([1, num_frames, 16, HEIGHT // 8, WIDTH // 8],
                        device="cuda", dtype=DTYPE)

    start = time.time()
    video = pipe.inference(noise=noise, conditional_dict=cond, return_latents=False)
    torch.cuda.synchronize()
    elapsed = time.time() - start
    pipe.vae.model.clear_cache()

    frames = (video[0].permute(0, 2, 3, 1).float().cpu().numpy() * 255)
    frames = frames.clip(0, 255).astype(np.uint8)
    out_path = os.path.join(tempfile.mkdtemp(), "forgewm.mp4")
    writer = imageio.get_writer(out_path, fps=FPS, codec="libx264",
                                quality=8, macro_block_size=None)
    for frame in frames:
        writer.append_data(frame)
    writer.close()

    script = " → ".join(f"`{a}` ×{n}" for a, n in segments)
    info = (
        f"**{variant}** · {total_chunks} causal blocks · {num_frames} latent "
        f"frames → {len(frames)} pixel frames ({len(frames) / FPS:.1f}s @ {FPS} fps)\n\n"
        f"Action script: {script}\n\n"
        f"Rollout + VAE decode: **{elapsed:.2f}s** "
        f"({1000 * elapsed / total_chunks:.0f} ms per block, decode included) · seed `{int(seed)}`"
    )
    return out_path, info


def _estimate(variant, chunks) -> int:
    per = SECONDS_PER_CHUNK.get(variant, 1.5)
    return int(FIXED_OVERHEAD + per * max(1, min(int(chunks), 12)) + 8)


def _duration_main(*args):
    variant = args[1] if len(args) > 1 else DEFAULT_VARIANT
    chunks = sum(int(args[i]) for i in (3, 5, 7) if len(args) > i)
    return _estimate(variant, chunks)


def _duration_example(*args):
    chunks = sum(int(args[i]) for i in (2, 4, 6) if len(args) > i)
    return _estimate(DEFAULT_VARIANT, chunks)


@spaces.GPU(duration=_duration_main)
def generate(
    image: str,
    variant: str,
    action_1: str,
    chunks_1: int,
    action_2: str,
    chunks_2: int,
    action_3: str,
    chunks_3: int,
    seed: int = 0,
    progress=gr.Progress(track_tqdm=True),
):
    """Roll out an action-conditioned Minecraft video from one reference frame.

    Args:
        image: Path to the reference frame that anchors the world.
        variant: Which few-step ForgeWM student to sample with.
        action_1: Action held during the first segment of the rollout.
        chunks_1: Length of the first segment, in 3-latent-frame blocks (~1s each).
        action_2: Action held during the second segment.
        chunks_2: Length of the second segment, in blocks.
        action_3: Action held during the third segment.
        chunks_3: Length of the third segment, in blocks.
        seed: Random seed for the initial noise.

    Returns:
        An mp4 of the generated rollout and a markdown summary of the run.
    """
    return _rollout(
        image, variant,
        [(action_1, chunks_1), (action_2, chunks_2), (action_3, chunks_3)],
        seed,
    )


@spaces.GPU(duration=_duration_example)
def generate_example(
    image: str,
    action_1: str,
    chunks_1: int,
    action_2: str,
    chunks_2: int,
    action_3: str,
    chunks_3: int,
    progress=gr.Progress(track_tqdm=True),
):
    """Run a bundled example with the default student and seed 0."""
    return _rollout(
        image, DEFAULT_VARIANT,
        [(action_1, chunks_1), (action_2, chunks_2), (action_3, chunks_3)],
        0,
    )


# ───────────────────────────────── UI ─────────────────────────────────
CSS = """
.gradio-container { max-width: 1200px !important; }
"""

with gr.Blocks(title="ForgeWM") as demo:
    gr.Markdown(
        "# 🎮 ForgeWM — few-step action-conditioned world model\n"
        "Give it **one Minecraft frame** and a short **action script**; the "
        "block-causal diffusion transformer rolls the world forward at "
        "**1, 2 or 4 denoising steps** per 3-frame block.\n\n"
        "[Paper](https://huggingface.co/papers/2608.14022) · "
        "[Model](https://huggingface.co/ForgeWM/ForgeWM) · "
        "[Code](https://github.com/asdfo123/ForgeWM) · "
        "[Project page](https://asdfo123.github.io/ForgeWM/)"
    )

    with gr.Row():
        with gr.Column(scale=1):
            image = gr.Image(label="Reference frame", type="filepath",
                             height=280, sources=["upload", "clipboard"])
            variant = gr.Radio(
                choices=list(PIPELINES.keys()),
                value=DEFAULT_VARIANT,
                label="Student",
                info="Fewer steps = faster. 1-/2-step use First-Frame Enhancement.",
            )
            gr.Markdown("### Action script  \nEach block is 12 frames ≈ 1 second.")
            with gr.Row():
                action_1 = gr.Dropdown(MINECRAFT_ACTIONS, value="forward",
                                       label="Segment 1", scale=2)
                chunks_1 = gr.Slider(0, 8, value=3, step=1, label="blocks", scale=1)
            with gr.Row():
                action_2 = gr.Dropdown(MINECRAFT_ACTIONS, value="turn_right",
                                       label="Segment 2", scale=2)
                chunks_2 = gr.Slider(0, 8, value=2, step=1, label="blocks", scale=1)
            with gr.Row():
                action_3 = gr.Dropdown(MINECRAFT_ACTIONS, value="forward",
                                       label="Segment 3", scale=2)
                chunks_3 = gr.Slider(0, 8, value=2, step=1, label="blocks", scale=1)
            with gr.Accordion("Advanced", open=False):
                seed = gr.Slider(0, 2**31 - 1, value=0, step=1, label="Seed")
            run = gr.Button("Roll out the world", variant="primary")

        with gr.Column(scale=1):
            video = gr.Video(label="Generated rollout", autoplay=True, loop=True)
            info = gr.Markdown()

    gr.Examples(
        examples=[
            ["examples/forest.png", "forward", 3, "turn_right", 2, "forward", 2],
            ["examples/plains.png", "forward", 3, "look_up", 1, "forward_turn_right", 3],
            ["examples/cave.png", "forward", 2, "turn_left", 2, "forward", 3],
        ],
        inputs=[image, action_1, chunks_1, action_2, chunks_2, action_3, chunks_3],
        outputs=[video, info],
        fn=generate_example,
        cache_examples=True,
        cache_mode="lazy",
        label="Examples (reference frames from the ForgeWM repo)",
    )

    gr.Markdown(
        "Rollouts are 352×640 at 12 fps. Attention uses a 6-latent-frame sliding "
        "window with a block-causal KV cache, exactly as the released students "
        "were trained. Base weights: "
        "[Skywork/Matrix-Game-2.0](https://huggingface.co/Skywork/Matrix-Game-2.0)."
    )

    run.click(
        fn=generate,
        inputs=[image, variant, action_1, chunks_1, action_2, chunks_2,
                action_3, chunks_3, seed],
        outputs=[video, info],
        concurrency_limit=1,
        api_name="generate",
    )

demo.queue(max_size=12).launch(
    theme=gr.themes.Citrus(), css=CSS, mcp_server=True, show_error=True)