"""
PPMem: Persistent Planar Memory for Video World Models — interactive demo.

Single-image, camera-posed video generation with the authors' released checkpoint
(Robbyant-Research/PPMem): a starting image, a camera trajectory (or a synthesized
gentle dolly), optional geometry (depth + intrinsics) — and out comes a 25-frame
1024x576 RGB video plus the planar geometry the model predicts along that camera:
depth maps, normal maps and a world-space planar mesh (PLY).

This is a 1:1 port of the reference rollout (ppmem.rollout.step0 + VideoGenerator.
generate_video) restricted to one AR chunk. The one hardware-forced deviation: the
reference encodes the prompt at build time and frees the text encoder; on ZeroGPU no
compute is allowed at module scope and the prompt arrives per request, so tokenizer +
UMT5 text encoder stay resident and prompts are encoded inside the GPU call.
"""

import os

# Cut VRAM fragmentation on the large per-request allocations (DiT + VAE + TE).
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

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

import gc
import tempfile
import time
from pathlib import Path

import cv2
import gradio as gr
import numpy as np
import torch
from einops import rearrange
from PIL import Image

from huggingface_hub import snapshot_download

from ppmem._logging import logger
from ppmem.config import MODEL, JobConfig
from ppmem.engine import depth_latent
from ppmem.engine import depth_latent as _depth_latent
from ppmem.engine.determinism import set_determinism
from ppmem.model.dit import WanTransformerPlana3DModel
from ppmem.model.loaders import (
    load_text_encoder,
    load_tokenizer,
    load_vae,
    read_checkpoint,
)
from ppmem.model.pipeline import PPMemImageToVideoPipeline, t5_prompt_embeds
from ppmem.model.planes import (
    planes_to_mesh,
    save_depth_maps,
    save_normal_maps,
    save_planes_as_ply,
    transform_planes_to_world,
)
from ppmem.model.scheduler import FlowUniPCMultistepScheduler
from ppmem.pool.init_geometry import load_init_geometry
from ppmem.pool.segment_math import chunk_seed
from ppmem.rollout.latents import LatentEncoder
from ppmem.rollout.moge import MoGeEstimator
from ppmem.rollout.step0 import step0

try:
    from diffusers.utils import export_to_video
except ImportError:  # diffuters moved it around between versions
    from diffusers.video_utils import export_to_video

SEGMENT_SIZE = 25           # frames per AR chunk (the released rollout's chunk size)
VIDEO_H, VIDEO_W = 576, 1024
FPS = 15
NUM_CONTEXT = 5
DEFAULT_STEPS = 50
DEFAULT_GUIDANCE = 5.0
DEFAULT_SEED = 42

DTYPE = torch.bfloat16
DEVICE = "cuda"

_job = JobConfig()  # run defaults: steps=50, guidance=5.0, flow_shift=5.0

# ----------------------------------------------------------------------
# Module-scope model load (eager; ZeroGPU intercepts .to("cuda") and streams
# the weights into VRAM on the first @spaces.GPU entry).
# ----------------------------------------------------------------------

ckpt_dir = snapshot_download(repo_id="Robbyant-Research/PPMem")
logger.info(f"[demo] checkpoint at {ckpt_dir}")

set_determinism(_job.run.deterministic)

tokenizer = load_tokenizer(ckpt_dir, subfolder="tokenizer")
text_encoder = load_text_encoder(
    ckpt_dir, subfolder="text_encoder", torch_dtype=DTYPE, torch_device=DEVICE,
)
text_encoder.requires_grad_(False)

with torch.device("meta"):
    transformer = WanTransformerPlana3DModel(MODEL)
transformer.cast_tensors(DTYPE)
# The reference builder materializes via to_empty(device); on ZeroGPU the plain
# .to("cuda") path is the one the hijack intercepts (to_empty allocates tensors
# directly with torch.empty(device=...), which may bypass the monkey-patch), so
# materialize on CPU, load the checkpoint, then hand the result to .to("cuda")
# like a normal load.
transformer.to_empty(device="cpu")
with torch.no_grad():
    transformer.init_weights()

_state_dict, _read_as = read_checkpoint(ckpt_dir)
_missing, _unexpected = transformer.load_state_dict(_state_dict, strict=False)
del _state_dict
gc.collect()
_missing = [k for k in _missing if not k.startswith("vae.")]
if _missing or _unexpected:
    raise RuntimeError(
        f"checkpoint mismatch: missing={_missing[:5]} unexpected={_unexpected[:5]}"
    )
transformer.cast_for_inference(DTYPE)
transformer.to(DEVICE)  # ZeroGPU-hijacked: weights packed to disk, streamed into VRAM per call

vae = load_vae(ckpt_dir, subfolder="vae", torch_dtype=DTYPE, torch_device=DEVICE)
vae.requires_grad_(False)
transformer.vae = vae  # dit.py reads vae.config latents_mean/std during denoising

scheduler = FlowUniPCMultistepScheduler(shift=_job.run.flow_shift)
pipeline = PPMemImageToVideoPipeline(
    tokenizer=None, text_encoder=None, transformer=transformer, vae=vae,
    scheduler=scheduler, pipe_dtype=DTYPE,
)

encoder = LatentEncoder(vae, device=DEVICE, dtype=DTYPE)


# The repo binds InferenceEngine.encode_depth_to_latent AND normalize_depth to the
# plain `self`-taking functions (engine/builder.py:115-116: `InferenceEngine.normalize_depth
# = depth_latent.normalize_depth`); encode_depth_to_latent calls self.normalize_depth
# internally (depth_latent.py:88). Mirror that assignment on a host object exactly:
# plain-function class attributes bind as methods, so `self` resolves correctly.
class _DepthCodec:
    normalize_depth = _depth_latent.normalize_depth
    encode_depth_to_latent = _depth_latent.encode_depth_to_latent


codec = _DepthCodec()
codec.vae = vae
codec.model_dtype = DTYPE
codec.device = DEVICE
encode_depth_to_latent = codec.encode_depth_to_latent

moge = None  # lazy: bundled examples ship geometry.npz and never touch MoGe


# ----------------------------------------------------------------------
# Helpers
# ----------------------------------------------------------------------


def synthesize_dolly(dist: float = 1.5, n: int = SEGMENT_SIZE) -> np.ndarray:
    """A gentle forward dolly (authors' `traj_dolly` shape), (n, 4, 4) c2w float32.

    Frame 0 is the identity pose; the camera translates forward along its optical
    axis, within the per-frame motion guidelines the repo documents (<=0.20 m/frame).
    """
    half = (n - 1) / 2.0
    zs = [dist * (i / half) if i <= half else dist * (2.0 - i / half) for i in range(n)]
    poses = np.stack([np.eye(4, dtype=np.float64) for _ in range(n)])
    for i, z in enumerate(zs):
        poses[i, 2, 3] = z
    return poses.astype(np.float32)


def _load_camera(camera_path) -> np.ndarray:
    """Read a user-supplied camera .npy and validate its shape; or synthesize a dolly."""
    if not camera_path:
        return synthesize_dolly()
    arr = np.load(camera_path)
    if arr.ndim != 3 or arr.shape[1:] != (4, 4):
        raise ValueError(
            f"camera file must hold an (N, 4, 4) array of c2w poses, got {arr.shape}"
        )
    return arr


def _depth_video_file(out_dir: Path, prefix: str = "depth") -> str:
    """Stitch the turbo-colormapped depth PNGs into an mp4 (10 evenly spaced frames)."""
    import glob as _glob

    frames = sorted(_glob.glob(str(out_dir / f"{prefix}_*.png")))
    if not frames:
        return None
    pickup = frames[:: max(1, len(frames) // 10)][:10]
    arr = np.stack([np.asarray(Image.open(f).convert("RGB")) for f in pickup])
    path = out_dir / f"{prefix}_video.mp4"
    try:
        import imageio

        imageio.mimwrite(str(path), arr, fps=5, quality=9, macro_block_size=1)
    except Exception as exc:  # ffmpeg backend problems: fall back to gr.Video-free PNGs
        logger.warning(f"depth video write failed ({exc}); returning first PNG")
        return frames[0]
    return str(path)


def _gpu_duration(
    image=None, camera_file=None, geometry_file=None, prompt=None,
    seed=DEFAULT_SEED, steps=DEFAULT_STEPS, guidance=DEFAULT_GUIDANCE,
    *args, **kwargs,
):
    """Reserve GPU seconds from the measured 50-step run (scene01, geometry supplied:
    ~290 s GPU time end-to-end, cold worker included). Denoise scales ~linearly with
    steps, so the request shrinks with the slider and visitors with partial quota can
    still run at fewer steps. Without a geometry file MoGe-2 depth estimation runs
    first, which adds a fixed cost."""
    t = 60 + float(steps) * 4.8
    if not geometry_file:
        t += 35
    return int(min(t, 360))


@spaces.GPU(duration=_gpu_duration)
def generate(
    image,
    camera_file,
    geometry_file,
    prompt: str,
    seed: int = DEFAULT_SEED,
    steps: int = DEFAULT_STEPS,
    guidance: float = DEFAULT_GUIDANCE,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a 25-frame camera-posed video plus planar geometry from one image.

    Args:
        image: the starting frame (RGB image).
        camera_file: optional .npy of (N, 4, 4) camera-to-world poses. When empty,
            a gentle forward dolly is synthesized.
        geometry_file: optional .npz with depth/fx/fy/cx/cy for frame 0 (as the
            authors' examples ship). When empty, MoGe-2 estimates it from the image.
        prompt: scene description for the text conditioning.
        seed: generation seed (chunk seed = seed + 0, as the reference rollout at step 0).
        steps: number of UniPC inference steps.
        guidance: classifier-free guidance scale.
        progress: Gradio progress.

    Returns:
        (video_path, depth_video or first depth map, [normal map paths], ply_path)
    """
    global moge
    t0 = time.time()
    progress((1, 12), desc="prep")
    work = Path(tempfile.mkdtemp(prefix="ppmem_"))

    if image is None:
        raise gr.Error("Please provide a starting image.")
    src_image = Path(work) / "input.png"
    if isinstance(image, np.ndarray):  # gr.Image numpy (H, W, 3)
        Image.fromarray(image.astype(np.uint8)).save(src_image)
    else:
        Image.open(image).convert("RGB").save(src_image)

    # ---- geometry bootstrap (reference: ppmem.cli --init_geometry) ----
    init_depth = init_K = None
    if geometry_file:
        with Image.open(src_image) as im:
            hw_image = (im.size[1], im.size[0])
        init_depth, init_K = load_init_geometry(geometry_file, (VIDEO_H, VIDEO_W), hw_image)

    extrinsic_c2w = _load_camera(camera_file)

    # ---- step0: projection + conditioning latents (1:1 reference path) ----
    if init_depth is None:
        if moge is None:
            moge = MoGeEstimator(model_type="moge-2-vitl-normal", device=DEVICE)
    progress((2, 12), desc="step0: project + encode")
    meta = step0(
        image_path=str(src_image),
        extrinsic_c2w=extrinsic_c2w,
        prompt=prompt,
        incremental_dir=str(work / "pool"),
        h_tgt=VIDEO_H,
        w_tgt=VIDEO_W,
        device=DEVICE,
        moge=moge,
        encoder=encoder,
        num_context=NUM_CONTEXT,
        segment_size=SEGMENT_SIZE,
        init_depth=init_depth,
        init_K=init_K,
    )
    meta.pop("proj_rgb_frames", None)

    # ---- prompt embeds (the reference pre-encodes; ZeroGPU forces per-call) ----
    progress((3, 12), desc="encode prompt")
    with torch.no_grad():
        prompt_embeds = t5_prompt_embeds(tokenizer, text_encoder, prompt, device=DEVICE)
        negative_prompt_embeds = t5_prompt_embeds(tokenizer, text_encoder, "", device=DEVICE)

    # ---- func_input, exactly as VideoGenerator.generate_video assembles it ----
    latent_h = meta["latent_height"]
    latent_w = meta["latent_width"]
    latent_f = meta["latent_num_frames"]

    generator = torch.manual_seed(chunk_seed(seed, 0))

    func_input = {
        "prompt_embeds": prompt_embeds,
        "negative_prompt_embeds": negative_prompt_embeds,
        "height": meta["video_height"],
        "width": meta["video_width"],
        "num_frames": meta["video_num_frames"],
        "num_inference_steps": steps,
        "guidance_scale": guidance,
        "generator": generator,
        "image": meta["first_frame_pil_images"].resize((meta["video_width"], meta["video_height"])),
    }

    num_ctx = meta.get("context_num_frames", 0)
    if num_ctx > 0 and "latent_context_frames" in meta:
        func_input["latents_context"] = rearrange(
            meta["latent_context_frames"], "(f h w) c -> c f h w",
            f=num_ctx, h=latent_h, w=latent_w,
        ).to(DEVICE)

    vf_data = meta["latent_very_first_frame"]
    n_vf = vf_data.shape[0] // (latent_h * latent_w)
    func_input["latents_very_first"] = rearrange(
        vf_data, "(f h w) c -> c f h w", f=n_vf, h=latent_h, w=latent_w,
    ).to(DEVICE)

    for suffix in ("1", "2"):
        for prefix in ("proj_rgb", "proj_mask"):
            key = f"latent_{prefix}_{suffix}"
            func_input[f"latent_{prefix}_{suffix}"] = rearrange(
                meta[key], "(f h w) c -> c f h w",
                f=latent_f, h=latent_h, w=latent_w,
            ).to(DEVICE)

    depth_norm_method = MODEL.depth_norm_method  # 'log'
    depth_range = (MODEL.depth_norm_range_min, MODEL.depth_norm_range_max)  # (0.2, 99.0)

    next_ar_depth = meta["next_ar_depth"].to(DEVICE)
    func_input["latents_depth"] = encode_depth_to_latent(
        next_ar_depth, norm_method=depth_norm_method, depth_range=depth_range,
    ).squeeze(0)
    # hv = history (none at step 0) + very first (frame 0's depth)
    hv_vf = meta["next_ar_depth"][0].cpu().numpy().astype(np.float32)
    hv_depth = torch.from_numpy(hv_vf[None, ...]).to(DEVICE)
    func_input["latents_hv_depth"] = encode_depth_to_latent(
        hv_depth, norm_method=depth_norm_method, depth_range=depth_range,
    ).squeeze(0)
    func_input["raw_next_ar_depth"] = [next_ar_depth]
    if num_ctx > 0 and meta.get("context_depth", None) is not None and meta["context_depth"].numel() > 0:
        func_input["latents_context_depth"] = encode_depth_to_latent(
            meta["context_depth"].to(DEVICE),
            norm_method=depth_norm_method, depth_range=depth_range,
        ).squeeze(0)

    K_list = []
    for intr in meta["intrinsics"]:
        K_list.append(torch.tensor([
            [intr["fx"], 0.0, intr["cx"]],
            [0.0, intr["fy"], intr["cy"]],
            [0.0, 0.0, 1.0],
        ], dtype=torch.float32))
    intrinsics = torch.stack(K_list).to(DEVICE)

    extrinsics_c2w_full = torch.tensor(meta["extrinsics_c2w"], dtype=torch.float32).to(DEVICE)
    if extrinsics_c2w_full.shape[0] < SEGMENT_SIZE:  # reference pads short sequences
        pad = SEGMENT_SIZE - extrinsics_c2w_full.shape[0]
        extrinsics_c2w_full = torch.cat(
            [extrinsics_c2w_full, extrinsics_c2w_full[-1:].repeat(pad, 1, 1)], dim=0
        )
    intrinsics = intrinsics[:SEGMENT_SIZE]
    while intrinsics.shape[0] < SEGMENT_SIZE:
        intrinsics = torch.cat([intrinsics, intrinsics[-1:]], dim=0)
    func_input["intrinsics"] = intrinsics
    func_input["extrinsics_c2w"] = extrinsics_c2w_full[:SEGMENT_SIZE]

    # This VAE maps frames 1:1; pass the count the pipeline maps back to
    # segment_size latent frames (Wan temporal convention: (n-1)*4+1).
    func_input["num_frames"] = (func_input["num_frames"] - 1) * 4 + 1

    progress((4, 12), desc="denoising (UniPC)")
    with torch.no_grad():
        pipeline_output = pipeline(**func_input)
    videos = pipeline_output.frames
    plane_params = getattr(pipeline_output, "plane_params", None)

    progress((10, 12), desc="export")
    out_dir = work / "out"
    out_dir.mkdir(parents=True, exist_ok=True)
    video_path = out_dir / "generated_video.mp4"
    export_to_video(videos[0], str(video_path), fps=FPS)

    depth_out = None
    normal_paths = []
    ply_path = None
    if plane_params is not None:
        extr = plane_params.get("extrinsics_c2w", None)
        pps = plane_params.get("plane_params_per_video", [])
        if extr is not None and len(pps):
            pp = pps[0]
            extr_t = extr if torch.is_tensor(extr) else torch.tensor(np.asarray(extr), dtype=torch.float32)
            verts_w, faces_w, colors_w, vnormals_w = planes_to_mesh(pp, extr_t.to(pp["plane_center_camera"].device))
            ply_path = out_dir / "planes_world.ply"
            save_planes_as_ply(str(ply_path), verts_w, faces_w, colors_w, vnormals_w)
            ply_path = str(ply_path)

            d_dir = out_dir / "cam_0_depth"
            save_depth_maps(str(d_dir), pp)
            depth_out = _depth_video_file(Path(d_dir))

            n_dir = out_dir / "cam_0_normals"
            save_normal_maps(str(n_dir), pp)
            import glob as _g

            normal_paths = sorted(_g.glob(str(n_dir / "*.png")))[:: max(1, len(sorted(_g.glob(str(n_dir / "*.png")))) // 4)][:4]

    logger.info(f"[demo] generate done in {time.time()-t0:.1f}s -> {video_path}")

    if depth_out is None:
        depth_out = str(video_path)  # keep the slot filled if geometry was skipped
    return str(video_path), depth_out, normal_paths, ply_path


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

examples_dir = Path(__file__).parent / "examples"
_scenes = sorted(p.name for p in examples_dir.glob("scene*") if p.is_dir())
EXAMPLES = []
for scene in _scenes:
    base = examples_dir / scene
    prompt_text = (base / "prompt.txt").read_text().strip() if (base / "prompt.txt").exists() else "A video of an indoor scene."
    EXAMPLES.append([
        str(base / "input.png"),
        str(base / "camera.npy"),
        str(base / "geometry.npz"),
        prompt_text,
    ])

with gr.Blocks(title="PPMem: Persistent Planar Memory for Video World Models") as demo:
    gr.Markdown(
        """
        # PPMem: Persistent Planar Memory for Video World Models
        **Paper** [2610.23795](https://huggingface.co/papers/2610.23795) · **Code** [robbyant-research/PPMem](https://github.com/robbyant-research/PPMem) · **Weights** [Robbyant-Research/PPMem](https://huggingface.co/Robbyant-Research/PPMem)

        An interactive video world model: give it a **starting image** and a **camera trajectory**, and it generates a 25-frame 1024×576 video from the first pose — plus the **3D planar geometry** it predicts along that camera: depth maps, normal maps, and a world-space planar mesh.
        """
    )
    with gr.Row():
        with gr.Column():
            image_in = gr.Image(label="Starting image", type="filepath", height=280)
            camera_in = gr.File(label="Camera trajectory (.npy, (N,4,4) c2w poses) — optional; a gentle dolly is used if empty", file_types=[".npy"])
            geometry_in = gr.File(label="Geometry .npz (depth, fx, fy, cx, cy) — optional; MoGe-2 estimates it if empty", file_types=[".npz"])
            prompt_in = gr.Textbox(label="Prompt", lines=3, placeholder="Describe the scene…")
            with gr.Accordion("Advanced", open=False):
                seed_in = gr.Number(label="Seed", value=DEFAULT_SEED, precision=0)
                steps_in = gr.Slider(label="Inference steps", minimum=10, maximum=50, value=DEFAULT_STEPS, step=1)
                guidance_in = gr.Slider(label="Guidance scale", minimum=1.0, maximum=10.0, value=DEFAULT_GUIDANCE, step=0.5)
            run_btn = gr.Button("Generate video", variant="primary")
        with gr.Column():
            video_out = gr.Video(label="Generated RGB video (1024×576, 25 frames @ 15fps)")
            depth_out = gr.Video(label="Predicted depth (turbo, from planar geometry)")
            normals_out = gr.Gallery(label="Predicted normals (sample frames)", columns=4, height=180)
            planes_out = gr.Model3D(label="World-space planar mesh (PLY)")

    run_btn.click(
        generate,
        inputs=[image_in, camera_in, geometry_in, prompt_in, seed_in, steps_in, guidance_in],
        outputs=[video_out, depth_out, normals_out, planes_out],
        api_name="generate",
    )

    if EXAMPLES:
        gr.Examples(
            examples=EXAMPLES,
            inputs=[image_in, camera_in, geometry_in, prompt_in],
            fn=generate,
            outputs=[video_out, depth_out, normals_out, planes_out],
            cache_examples=True,
            cache_mode="lazy",
        )

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