#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "gradio>=5.49,<6",
#     "spaces",
#     "kernels",
#     "trimesh>=4.4",
#     "torch>=2.5.1",
#     "torchvision>=0.20.1",
#     "dvlt @ https://huggingface.co/spaces/blanchon/dvlt/resolve/main/packages/dvlt/wheels/dvlt-0.0.1-py3-none-any.whl",
# ]
# ///
"""Unofficial Gradio / ZeroGPU demo for Déjà View (DVLT).

Déjà View (DVLT) — Looping Transformers for Multi-View 3D Reconstruction.
DVLT loops a shared transformer block K times; we decode an intermediate point
cloud every N steps so the viewer shows it converge.

Run standalone with `uv run app.py` (deps in the PEP 723 header above), or as a
Hugging Face ZeroGPU Space (deps in requirements.txt; torch from the image).
"""

from __future__ import annotations

import os
import sys
import tempfile
import time
from dataclasses import dataclass, field

import cv2
import gradio as gr
import numpy as np
import spaces
import torch
from accelerate import Accelerator
from PIL import Image

# Flash-Attention 3, prebuilt for ZeroGPU's Blackwell GPUs and exposed under the
# module name DVLT imports (`from flash_attn_interface import flash_attn_func`).
try:
    from kernels import get_kernel

    sys.modules["flash_attn_interface"] = get_kernel("kernels-community/flash-attn3", version=1)
except Exception as exc:  # local / no compatible kernel -> DVLT falls back to SDPA
    print(f"[dvlt] FA3 kernel unavailable ({exc}); using default attention.")

from dvlt.common.constants import DataField, PredictionField
from dvlt.common.geometry import depth_to_world_coords_points
from dvlt.common.pose import to4x4
from dvlt.model.dvlt.model import DVLT, _slice_expand_flatten
from dvlt.model_components import set_attn_backend
from dvlt.util.preprocess import preprocess_images
from dvlt.viz.depth import overlay_depth_map
from dvlt.viz.glb import pointcloud_to_glb
from dvlt.viz.pointcloud import zero_depths_on_pad


# --- config (matched to the released nvidia/dvlt stage-2 checkpoint) ---------
CHECKPOINT = os.environ.get("DVLT_CHECKPOINT", "nvidia/dvlt")
IMG_SIZE = 504
PATCH_SIZE = 14
DEFAULT_STEPS = 12  # K — the model's default inference step count
MAX_STEPS = 24
MAX_FRAMES = 16  # cap views so global attention stays within the ZeroGPU budget

VIDEO_FPS_DEFAULT = 2.0
DECODE_EVERY_DEFAULT = 3
FRUSTUM_DEFAULT = 2.0  # % of scene radius
CONF_DEFAULT = 50.0
MAX_POINTS_DEFAULT = 400_000

EXAMPLES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples")
IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp", ".bmp", ".JPG", ".JPEG", ".PNG")


# --- model (placed on CUDA at module level per the ZeroGPU guidance) ---------
# bf16 on GPU (the repo's inference precision); fp32 on CPU (bf16 breaks the
# CPU pose solver). Autocast handles the cast; weights stay fp32.
_ACCEL = Accelerator(mixed_precision="bf16" if torch.cuda.is_available() else "no")
_MODEL: DVLT | None = None


def load_model() -> DVLT:
    global _MODEL
    if _MODEL is None:
        try:
            set_attn_backend("fa3")
        except Exception as exc:  # noqa: BLE001
            print(f"[dvlt] fa3 unavailable ({exc}); using default attention.")
        model = DVLT(img_size=IMG_SIZE, depth_head_type="conv")
        model.load_pretrained(CHECKPOINT, strict=True)
        model.setup_test(_ACCEL)
        if torch.cuda.is_available():
            model.model.to("cuda")
        _MODEL = model
    return _MODEL


# --- point cloud / GLB / depth helpers ---------------------------------------
def predictions_to_cloud(preds, batch, max_points, conf_threshold):
    """Depth-unproject + filter predictions into ``(points[N,3], colors[N,3] uint8)``."""
    depths = preds[PredictionField.DEPTHS][0].float()
    cameras = preds[PredictionField.CAMERAS][0]
    extrinsics_c2w = to4x4(cameras.camera_to_worlds).float()
    intrinsics = cameras.get_intrinsics_matrices().float()

    world_points, _, valid_mask = depth_to_world_coords_points(depths, extrinsics_c2w, intrinsics)

    images = batch[DataField.IMAGES][0]
    colors = images.detach().float().cpu().permute(0, 2, 3, 1).numpy() * 255.0

    pts = world_points.detach().float().cpu().numpy().reshape(-1, 3)
    cols = colors.reshape(-1, 3)
    mask = valid_mask.detach().cpu().numpy().reshape(-1)

    pad_ok = batch.get("gradio_valid_pixels", None)
    if pad_ok is not None:
        mask = mask & pad_ok[0].detach().cpu().numpy().reshape(-1)

    pts, cols = pts[mask], cols[mask]

    conf = preds.get(PredictionField.DEPTHS_CONF, None)
    if conf is None:
        conf = preds.get(PredictionField.WORLD_POINTS_DIRECT_CONF, None)
    if conf is not None and len(pts) > 0:
        conf_flat = conf[0].detach().float().cpu().numpy().reshape(-1)[mask]
        keep = conf_flat >= np.percentile(conf_flat, float(conf_threshold))
        pts, cols = pts[keep], cols[keep]

    if max_points > 0 and len(pts) > max_points:
        idx = np.random.choice(len(pts), max_points, replace=False)
        pts, cols = pts[idx], cols[idx]

    return pts, cols.astype(np.uint8)


def cameras_to_glb_inputs(preds, batch):
    cameras = preds[PredictionField.CAMERAS][0]
    c2ws = to4x4(cameras.camera_to_worlds).detach().float().cpu().numpy()
    intrinsics = cameras.get_intrinsics_matrices().detach().float().cpu().numpy()
    image_hw = tuple(int(v) for v in batch[DataField.IMAGES].shape[-2:])
    return c2ws, intrinsics, image_hw


def depth_overlays(preds, batch):
    depths = preds[PredictionField.DEPTHS][0].detach().float().cpu().numpy()
    depths = zero_depths_on_pad(depths, batch)
    images = batch[DataField.IMAGES][0].detach().float().cpu().permute(0, 2, 3, 1).numpy()
    return [overlay_depth_map(img, d) for img, d in zip(images, depths)]


def cloud_to_glb(points, colors, c2ws=None, intrinsics=None, image_hw=None, show_cameras=True, frustum_frac=0.02):
    return pointcloud_to_glb(
        points,
        colors,
        cameras_c2w=c2ws if show_cameras else None,
        intrinsics=intrinsics,
        image_hw=image_hw,
        frustum_frac=float(frustum_frac) if show_cameras else 0.0,
    )


# --- streaming looped inference (the GPU function) ---------------------------
@dataclass
class StepResult:
    step: int
    total: int
    points: np.ndarray
    colors: np.ndarray
    c2ws: np.ndarray
    intrinsics: np.ndarray
    image_hw: tuple
    depth_overlays: list = field(default_factory=list)
    timings: str = ""  # phase timings (logged by the main process for visibility)


@spaces.GPU(duration=300)
def infer(image_paths, num_steps, conf_threshold, max_points, decode_every):
    """Run the K-step looped inference in one GPU call; return a StepResult per decode.

    Non-generator on purpose: a lazily-consumed @spaces.GPU generator hangs on
    ZeroGPU. We do all GPU work here, return numpy results, and the CPU side
    reveals them progressively.
    """
    model = load_model()
    m = model.model  # DVLTModel
    # Move to the *real* GPU inside the @spaces.GPU function (vggt-omega pattern);
    # relying on module-level placement alone can leave compute on CPU here.
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    m.to(device)

    frames = [Image.open(p).convert("RGB") for p in image_paths]
    batch = preprocess_images(frames, IMG_SIZE, PATCH_SIZE, device)
    images = batch[DataField.IMAGES]
    B, S, _, H, W = images.shape

    K = int(np.clip(int(num_steps), 1, MAX_STEPS))
    every = max(1, int(decode_every))
    results = []

    with torch.no_grad(), _ACCEL.autocast():
        t0 = time.time()
        z_0 = m._encode_images(images)
        register_token = m.register_token.expand(B, S, -1, -1).reshape(B * S, m.num_register_tokens, -1)
        camera_token = _slice_expand_flatten(m.camera_token, B, S)
        x = torch.cat([camera_token, register_token, z_0], dim=1)
        rope_pos = m._get_rope_positions(B * S, H, W, device)

        ts = torch.linspace(0.0, 1.0, K).tolist()
        for i in range(K):
            t_next = ts[i + 1] if i + 1 < K else 1.0
            x = m._interval_step(x, ts[i], t_next, rope_pos, B, S)

            is_last = i == K - 1
            if not (is_last or (i + 1) % every == 0):
                continue

            preds = model._postprocess_predictions(batch, m._decode(x, H, W, B, S, rope_pos))
            pts, rgb = predictions_to_cloud(preds, batch, int(max_points), float(conf_threshold))
            c2ws, intrinsics, image_hw = cameras_to_glb_inputs(preds, batch)
            overlays = depth_overlays(preds, batch) if is_last else []
            timings = (
                f"dev={device} amp={_ACCEL.mixed_precision} xdtype={x.dtype} | "
                f"{S} views K={K} decode@{i + 1} at {time.time() - t0:.1f}s"
            )
            results.append(StepResult(i + 1, K, pts, rgb, c2ws, intrinsics, image_hw, overlays, timings))
            del preds

    return results


# --- inputs: video -> frames -> image batch (the single source of truth) -----
def _cap(paths):
    """Evenly subsample to at most MAX_FRAMES."""
    if len(paths) <= MAX_FRAMES:
        return paths
    idx = np.linspace(0, len(paths) - 1, MAX_FRAMES).round().astype(int)
    return [paths[i] for i in idx]


def video_to_frames(video, video_fps):
    """Sample a video into frames and return their paths (to fill the image batch)."""
    if not video:
        return None
    out_dir = tempfile.mkdtemp(prefix="dvlt_frames_")
    cap = cv2.VideoCapture(video)
    src_fps = cap.get(cv2.CAP_PROP_FPS)
    stride = max(1, int(round((src_fps if src_fps and src_fps > 0 else 24.0) / max(float(video_fps), 0.1))))
    paths, idx, saved = [], 0, 0
    while True:
        ok, frame = cap.read()
        if not ok:
            break
        if idx % stride == 0:
            p = os.path.join(out_dir, f"frame_{saved:04d}.png")
            cv2.imwrite(p, frame)
            paths.append(p)
            saved += 1
        idx += 1
    cap.release()
    return _cap(sorted(paths))


def preview_images(image_paths):
    if not image_paths:
        return None, "Upload a video or images to begin."
    paths = _cap(sorted(image_paths))
    return paths, f"Loaded {len(paths)} frame(s). Click **Reconstruct**."


def _camera_info(n):
    return (
        f"**{n} camera pose(s)** estimated. Each rainbow frustum is one input view "
        "(blue → red = frame order), linked along the camera trajectory."
    )


# --- reconstruction (streaming) + CPU-only re-render -------------------------
def reconstruct(
    image_paths, num_steps, conf_thres, max_points, decode_every, show_cam, frustum_scale,
    progress=gr.Progress(),
):
    if not image_paths:
        raise gr.Error("Upload a video or images first.")
    image_paths = _cap(sorted(image_paths))
    frustum_frac = float(frustum_scale) / 100.0
    print(f"[recon] {len(image_paths)} frames, K={num_steps} -> GPU", flush=True)

    results = infer(image_paths, int(num_steps), float(conf_thres), int(max_points), int(decode_every))
    print(f"[recon] GPU returned {len(results)} step(s)", flush=True)

    for res in results:
        print(f"[recon] reveal step {res.step}/{res.total} ({res.timings})", flush=True)
        glb = cloud_to_glb(
            res.points, res.colors, c2ws=res.c2ws, intrinsics=res.intrinsics, image_hw=res.image_hw,
            show_cameras=bool(show_cam), frustum_frac=frustum_frac,
        )
        is_final = res.step >= res.total
        progress(res.step / max(res.total, 1), desc=f"Refinement step {res.step}/{res.total}")
        status = (
            f"✅ Done — {res.total} steps, {len(res.points):,} points, {len(res.c2ws)} cameras."
            if is_final
            else f"🔄 Step {res.step}/{res.total} — {len(res.points):,} points so far…"
        )
        cloud = {
            "points": res.points, "colors": res.colors, "c2ws": res.c2ws,
            "intrinsics": res.intrinsics, "image_hw": res.image_hw,
        }
        depth_gallery = res.depth_overlays if (is_final and res.depth_overlays) else gr.update()
        yield glb, status, depth_gallery, _camera_info(len(res.c2ws)), cloud


def update_view(cloud, max_points, show_cam, frustum_scale):
    """Re-render the cached cloud on CPU (density / cameras only)."""
    if not cloud:
        return None, "Nothing to update — run **Reconstruct** first."
    pts, cols = cloud["points"], cloud["colors"]
    if int(max_points) > 0 and len(pts) > int(max_points):
        idx = np.random.choice(len(pts), int(max_points), replace=False)
        pts, cols = pts[idx], cols[idx]
    glb = cloud_to_glb(
        pts, cols, c2ws=cloud["c2ws"], intrinsics=cloud["intrinsics"], image_hw=cloud["image_hw"],
        show_cameras=bool(show_cam), frustum_frac=float(frustum_scale) / 100.0,
    )
    return glb, f"View updated — {len(pts):,} points. (Change confidence → re-run Reconstruct.)"


# --- examples (videos only) --------------------------------------------------
def _build_examples():
    rows = []
    for vid, conf in [("conf_20_robot.mp4", 20.0), ("conf50.mp4", 50.0), ("conf_30.mp4", 30.0)]:
        p = os.path.join(EXAMPLES_DIR, vid)
        if os.path.isfile(p):
            rows.append([p, DEFAULT_STEPS, conf, MAX_POINTS_DEFAULT, True])
    return rows


_HEADER = """
<div align="center">
<h1>🔁 Déjà View — DVLT</h1>
<p><b>Looping Transformers for Multi-View 3D Reconstruction</b></p>
<p>
<a href="https://arxiv.org/abs/2605.30215">📄 Paper</a> &nbsp;•&nbsp;
<a href="https://research.nvidia.com/labs/dvl/projects/dvlt/">🌐 Project</a> &nbsp;•&nbsp;
<a href="https://github.com/nv-tlabs/dvlt">🐙 Code</a> &nbsp;•&nbsp;
<a href="https://huggingface.co/nvidia/dvlt">🤗 Model</a>
</p>
<p><sub>⚠️ <b>Unofficial</b> demo, not affiliated with NVIDIA. Model weights are
released under the NVIDIA non-commercial research license.</sub></p>
</div>

<p>Upload a <b>video</b> (auto-sampled into frames) or a set of <b>images</b>, then
hit <b>Reconstruct</b>. DVLT loops a shared transformer block <i>K</i> times to
predict depth, camera poses and a 3D point cloud — the viewer updates after every
step so you can watch it converge.</p>
"""

_CSS = ".dvlt-log * { font-size: 18px !important; font-weight: 600 !important; text-align: center !important; }"


def build_demo() -> gr.Blocks:
    with gr.Blocks(theme=gr.themes.Ocean(), css=_CSS, title="DVLT 3D Demo") as demo:
        gr.HTML(_HEADER)
        cloud_state = gr.State()  # cached final cloud for Update View

        with gr.Row():
            with gr.Column(scale=2):
                input_video = gr.Video(label="Upload Video (sampled into frames)")
                video_fps = gr.Slider(0.5, 8.0, value=VIDEO_FPS_DEFAULT, step=0.5, label="Video sampling FPS")
                input_images = gr.File(
                    file_count="multiple", file_types=["image"], type="filepath",
                    label="Image batch (videos land here; or upload images directly)",
                )
                gallery = gr.Gallery(label="Input frames", columns=4, height=240, object_fit="contain")

            with gr.Column(scale=4):
                status = gr.Markdown("Upload a video or images to begin.", elem_classes=["dvlt-log"])
                model3d = gr.Model3D(label="Point cloud + camera poses", height=560, clear_color=[0, 0, 0, 0])
                camera_info = gr.Markdown("")

                with gr.Row():
                    run_btn = gr.Button("Reconstruct", variant="primary", scale=2)
                    view_btn = gr.Button("Update View", scale=1)
                    gr.ClearButton([input_video, input_images, gallery, model3d, camera_info, cloud_state], scale=1)

                with gr.Accordion("Parameters", open=True):
                    with gr.Row():
                        num_steps = gr.Slider(
                            1, MAX_STEPS, value=DEFAULT_STEPS, step=1, label="Refinement steps (K)",
                            info="More steps = sharper geometry, more compute.",
                        )
                        decode_every = gr.Slider(
                            1, 6, value=DECODE_EVERY_DEFAULT, step=1, label="Preview every N steps",
                            info="How often to stream an intermediate cloud.",
                        )
                    with gr.Row():
                        conf_thres = gr.Slider(
                            0, 99, value=CONF_DEFAULT, step=1, label="Confidence threshold (percentile)",
                            info="Drops the least-confident points.",
                        )
                        max_points = gr.Slider(
                            50_000, 2_000_000, value=MAX_POINTS_DEFAULT, step=50_000, label="Max points",
                        )
                    with gr.Row():
                        frustum_scale = gr.Slider(0.0, 6.0, value=FRUSTUM_DEFAULT, step=0.1, label="Camera size (%)")
                        show_cam = gr.Checkbox(value=True, label="Show cameras")

                depth_gallery = gr.Gallery(label="Depth maps", columns=4, height=240, object_fit="contain")

        recon_inputs = [input_images, num_steps, conf_thres, max_points, decode_every, show_cam, frustum_scale]
        recon_outputs = [model3d, status, depth_gallery, camera_info, cloud_state]

        examples = _build_examples()
        if examples:
            def run_example(video, k, conf, mp, show):
                paths = video_to_frames(video, VIDEO_FPS_DEFAULT)
                final = (None, "No output.", gr.update(), "", None)
                for out in reconstruct(paths, k, conf, mp, DECODE_EVERY_DEFAULT, show, FRUSTUM_DEFAULT):
                    final = out
                glb, st, depth, cam, cloud = final
                return glb, st, depth, cam, cloud, paths

            gr.Markdown("### Examples &nbsp;<sub>(click a row to load & reconstruct)</sub>")
            gr.Examples(
                examples=examples,
                inputs=[input_video, num_steps, conf_thres, max_points, show_cam],
                outputs=[model3d, status, depth_gallery, camera_info, cloud_state, input_images],
                fn=run_example,
                cache_examples=False,
                examples_per_page=6,
            )

        # video -> frames -> image batch; the image batch drives the preview + reconstruction
        input_video.change(video_to_frames, [input_video, video_fps], [input_images])
        input_images.change(preview_images, [input_images], [gallery, status])

        run_btn.click(lambda: (None, "🔄 Reconstructing…"), None, [model3d, status]).then(
            reconstruct, recon_inputs, recon_outputs
        )
        view_btn.click(update_view, [cloud_state, max_points, show_cam, frustum_scale], [model3d, status])

    return demo


# Place the model on CUDA at module level (ZeroGPU guidance). `infer` is already
# defined above, so ZeroGPU detects the GPU function.
if os.environ.get("DVLT_SKIP_AUTOLOAD") != "1":
    load_model()

demo = build_demo()
demo.queue(max_size=20)

if __name__ == "__main__":
    # ssr_mode=False: the experimental Gradio SSR server can swallow streamed updates.
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True, ssr_mode=False)