#!/usr/bin/env python3
"""3D HAMSTER — 3D Trajectory Prediction demo (ZeroGPU).

Depth-aware VLM planner: predicts metric 3D end-effector trajectories (and 2D
trajectories / pointing / bbox / VQA) from a single RGB image + metric depth map
+ a language instruction.

Adapted from the reference Gradio script in the official repo
(scripts/trajectory_prediction_gradio.py) to run on ZeroGPU: the model is loaded
once at module scope onto CUDA and inference is wrapped in @spaces.GPU.
"""

import os

# DINOv2 geometry encoder uses xformers when available; force the pure-torch
# fallback so we don't need an xformers CUDA build on Blackwell.
os.environ.setdefault("XFORMERS_DISABLED", "1")

import spaces  # noqa: E402  MUST come before torch / transformers

import json  # noqa: E402
import re  # noqa: E402
import tempfile  # noqa: E402
from pathlib import Path  # noqa: E402
from typing import Optional  # noqa: E402

import cv2  # noqa: E402
import gradio as gr  # noqa: E402
import numpy as np  # noqa: E402
import plotly.graph_objects as go  # noqa: E402
import torch  # noqa: E402
from PIL import Image  # noqa: E402
from huggingface_hub import snapshot_download  # noqa: E402

# ── Constants (mirror the reference script) ──────────────────────────────────

MODEL_ID = "DAVIAN-Robotics/3D_HAMSTER"
HERE = Path(__file__).parent.resolve()
EXAMPLES_DIR = str(HERE / "examples")
TARGET_SIZE = 640  # training resolution: longest edge = 640

V5_SYSTEM_PROMPT = ""
VQA_STYLE = "General VQA"
BBOX_STYLE = "2D Bounding Box"

V5_PROMPT_SUFFIXES = {
    "3D Trajectory": (
        "Predict the full manipulation trajectory as point_3d waypoints "
        "with depth and gripper state in JSON."
    ),
    "2D Trajectory": (
        "Predict the full manipulation trajectory as point_2d waypoints "
        "with gripper state in JSON."
    ),
    "3D Pointing": "Report the point_3d location in JSON.",
    "2D Pointing": "Report point_2d locations in JSON.",
    BBOX_STYLE: None,
    VQA_STYLE: None,
}
V5_PROMPT_STYLES = list(V5_PROMPT_SUFFIXES.keys())
V5_DEFAULT_STYLE = "3D Trajectory"
POINTING_STYLES = {"2D Pointing", "3D Pointing"}


def build_v5_human_message(instruction: str, prompt_style: str) -> str:
    instr = instruction.strip()
    if prompt_style == BBOX_STYLE:
        return (
            f"I'm looking for {instr} in this image. Can you locate it? "
            "Report bbox coordinates in JSON format."
        )
    if prompt_style not in V5_PROMPT_SUFFIXES:
        prompt_style = V5_DEFAULT_STYLE
    suffix = V5_PROMPT_SUFFIXES[prompt_style]
    if suffix is None:
        return instr
    return f"{instr}\n{suffix}"


# ── Preprocessing (matches training pipeline) ────────────────────────────────


def resize_to_target(image, target_size=TARGET_SIZE, interp=cv2.INTER_LINEAR):
    h, w = image.shape[:2]
    scale = target_size / max(h, w)
    new_w, new_h = int(round(w * scale)), int(round(h * scale))
    resized = cv2.resize(image, (new_w, new_h), interpolation=interp)
    return resized, scale


def prepare_inputs_from_arrays(rgb, depth, tmp_dir):
    rgb_resized, scale = resize_to_target(rgb, TARGET_SIZE, cv2.INTER_LINEAR)
    depth_resized, _ = resize_to_target(depth, TARGET_SIZE, cv2.INTER_NEAREST)
    new_h, new_w = rgb_resized.shape[:2]
    mask = ((depth_resized > 0.01) & (depth_resized < 10.0)).astype(np.float32)
    pcd = np.zeros((new_h, new_w, 4), dtype=np.float32)
    pcd[:, :, 2] = depth_resized
    pcd[:, :, 3] = mask
    img_path = os.path.join(tmp_dir, "frame_0_640.png")
    npz_path = os.path.join(tmp_dir, "frame_0_640.npz")
    cv2.imwrite(img_path, cv2.cvtColor(rgb_resized, cv2.COLOR_RGB2BGR))
    np.savez_compressed(npz_path, pcd=pcd.astype(np.float16))
    return img_path, npz_path, rgb_resized, depth_resized, scale


def prepare_inputs(rgb_pil, depth_npy_path, tmp_dir):
    rgb = np.array(rgb_pil.convert("RGB"))
    depth = np.load(depth_npy_path).astype(np.float32)
    img_path, npz_path, rgb_resized, depth_resized, _ = prepare_inputs_from_arrays(
        rgb, depth, tmp_dir
    )
    return img_path, npz_path, rgb_resized, depth_resized


# ── Parsing ──────────────────────────────────────────────────────────────────


def parse_trajectory(output):
    """Legacy <ans>...</ans> parser."""
    waypoints, actions = [], []
    ans_match = re.search(r"<ans>(.*?)</ans>", output, re.DOTALL)
    if not ans_match:
        ans_match = re.search(r"\[\[.*?\]\]", output, re.DOTALL)
        if not ans_match:
            return [], []
        content = ans_match.group(0)
    else:
        content = ans_match.group(1)
    coord_pat = r"\[(\d+(?:\.\d+)?),\s*(\d+(?:\.\d+)?),\s*(\d+(?:\.\d+)?)\]"
    action_pat = r"<action>(.*?)</action>"
    parts = re.split(action_pat, content)
    for i, part in enumerate(parts):
        if i % 2 == 0:
            for c in re.findall(coord_pat, part):
                waypoints.append([float(c[0]), float(c[1]), float(c[2])])
                actions.append(None)
        else:
            if actions:
                actions[-1] = part.strip()
    return waypoints, actions


def parse_gt_structured_json(gpt_value):
    m = re.search(r"```json\s*(.*?)\s*```", gpt_value, re.DOTALL)
    raw = m.group(1) if m else gpt_value.strip()
    if not raw.lstrip().startswith("["):
        arr = re.search(r"\[.*\]", raw, re.DOTALL)
        if not arr:
            return [], []
        raw = arr.group(0)
    try:
        entries = json.loads(raw)
    except json.JSONDecodeError:
        return [], []
    if not isinstance(entries, list):
        return [], []
    key = "point_3d" if any(
        isinstance(e, dict) and "point_3d" in e for e in entries
    ) else "point_2d"
    waypoints, actions = [], []
    for entry in entries:
        if not isinstance(entry, dict) or key not in entry:
            continue
        pt = entry[key]
        waypoints.append(
            [float(pt[0]), float(pt[1]), float(pt[2]) if len(pt) > 2 else 0.0]
        )
        grip = entry.get("gripper", "none")
        if grip == "close":
            actions.append("Close Gripper")
        elif grip == "open":
            actions.append("Open Gripper")
        else:
            actions.append(None)
    return waypoints, actions


def parse_bbox_2d(output):
    m = re.search(r"```json\s*(.*?)\s*```", output, re.DOTALL)
    raw = m.group(1) if m else output.strip()
    if not raw.lstrip().startswith("["):
        arr = re.search(r"\[.*\]", raw, re.DOTALL)
        if not arr:
            return []
        raw = arr.group(0)
    try:
        entries = json.loads(raw)
    except json.JSONDecodeError:
        return []
    if not isinstance(entries, list):
        return []
    boxes = []
    for e in entries:
        if not isinstance(e, dict) or "bbox_2d" not in e:
            continue
        b = e["bbox_2d"]
        if len(b) < 4:
            continue
        boxes.append(
            (float(b[0]), float(b[1]), float(b[2]), float(b[3]), str(e.get("label", "")))
        )
    return boxes


# ── 2D visualization ─────────────────────────────────────────────────────────

COLOR_WP = (0, 255, 0)
COLOR_GRASP = (255, 0, 0)
COLOR_RELEASE = (0, 0, 255)
COLOR_LINE = (255, 255, 0)


def visualize_2d(image, waypoints, actions):
    if not waypoints:
        return image
    img = image.copy()
    h, w = img.shape[:2]
    pixels = [(int(wp[0] / 1000 * w), int(wp[1] / 1000 * h)) for wp in waypoints]
    for i in range(len(pixels) - 1):
        cv2.line(img, pixels[i], pixels[i + 1], COLOR_LINE, 2, cv2.LINE_AA)
    for i, (px, py) in enumerate(pixels):
        act = actions[i] if i < len(actions) else None
        if act and "Close" in act:
            color, r = COLOR_GRASP, 12
        elif act and "Open" in act:
            color, r = COLOR_RELEASE, 12
        else:
            color, r = COLOR_WP, 8
        cv2.circle(img, (px, py), r, color, -1)
        cv2.circle(img, (px, py), r, (255, 255, 255), 2)
        cv2.putText(img, str(i), (px - 5, py + 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 2)
        cv2.putText(
            img, f"d={waypoints[i][2]:.2f}m", (px + 15, py),
            cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1,
        )
    y = 30
    for label, color, xo in [
        ("Waypoint", COLOR_WP, 10), ("Grasp", COLOR_GRASP, 110), ("Release", COLOR_RELEASE, 180)
    ]:
        cv2.circle(img, (xo, y - 5), 6, color, -1)
        cv2.putText(img, label, (xo + 10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)
    return img


def visualize_points(image, points):
    if not points:
        return image
    img = image.copy()
    h, w = img.shape[:2]
    for i, p in enumerate(points):
        px, py = int(p[0] / 1000 * w), int(p[1] / 1000 * h)
        cv2.circle(img, (px, py), 8, (0, 255, 0), -1)
        cv2.circle(img, (px, py), 8, (255, 255, 255), 2)
        cv2.putText(
            img, str(i + 1), (px + 11, py + 4),
            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1,
        )
    return img


def visualize_bbox(image, boxes):
    if not boxes:
        return image
    img = image.copy()
    h, w = img.shape[:2]
    palette = [(0, 255, 0), (255, 80, 0), (0, 160, 255), (255, 0, 200), (255, 220, 0)]
    for i, (x1, y1, x2, y2, label) in enumerate(boxes):
        p1 = (int(x1 / 1000 * w), int(y1 / 1000 * h))
        p2 = (int(x2 / 1000 * w), int(y2 / 1000 * h))
        color = palette[i % len(palette)]
        cv2.rectangle(img, p1, p2, color, 2)
        tag = label or f"obj{i}"
        (tw, th), _ = cv2.getTextSize(tag, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        cv2.rectangle(img, (p1[0], p1[1] - th - 6), (p1[0] + tw + 4, p1[1]), color, -1)
        cv2.putText(img, tag, (p1[0] + 2, p1[1] - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
    return img


# ── 3D scene + trajectory visualization ──────────────────────────────────────


def colorize_depth(depth, min_d=0.1, max_d=3.0):
    d = np.clip(depth, min_d, max_d)
    d = ((d - min_d) / (max_d - min_d) * 255).astype(np.uint8)
    colored = cv2.applyColorMap(d, cv2.COLORMAP_TURBO)
    colored = cv2.cvtColor(colored, cv2.COLOR_BGR2RGB)
    colored[depth <= 0] = [0, 0, 0]
    return colored


def _create_sphere_points(center, radius=0.005, n=150):
    phi = np.random.uniform(0, 2 * np.pi, n)
    ct = np.random.uniform(-1, 1, n)
    theta = np.arccos(ct)
    return np.stack([
        center[0] + radius * np.sin(theta) * np.cos(phi),
        center[1] + radius * np.sin(theta) * np.sin(phi),
        center[2] + radius * np.cos(theta),
    ], axis=1)


def _create_tube_points(p1, p2, radius=0.002):
    d = p2 - p1
    L = np.linalg.norm(d)
    if L < 1e-8:
        return np.empty((0, 3))
    d = d / L
    perp1 = np.cross(d, [1, 0, 0]) if abs(d[0]) < 0.9 else np.cross(d, [0, 1, 0])
    perp1 /= np.linalg.norm(perp1)
    pts = []
    for ti in np.linspace(0, 1, max(int(L / 0.002), 10)):
        c = p1 + ti * (p2 - p1)
        for a in np.linspace(0, 2 * np.pi, 6, endpoint=False):
            pts.append(c + radius * (np.cos(a) * perp1 + np.sin(a) * np.cross(d, perp1)))
    return np.array(pts)


def _uvd_to_xyz(coords, intrinsics_3x3, img_w, img_h, uvd_norm=1000.0):
    K = np.array(intrinsics_3x3, dtype=np.float64)
    Kinv = np.linalg.inv(K)
    u_px = (coords[:, 0] / uvd_norm) * img_w
    v_px = (coords[:, 1] / uvd_norm) * img_h
    pixels = np.stack([u_px, v_px, np.ones(len(u_px))], axis=1)
    return coords[:, 2:3] * (pixels @ Kinv.T)


def default_intrinsics(img_h, img_w):
    f = float(max(img_h, img_w))
    return [[f, 0.0, img_w / 2.0], [0.0, f, img_h / 2.0], [0.0, 0.0, 1.0]]


def build_scene_pcd_simple(rgb, depth, intrinsics_3x3, depth_trunc=3.0, stride=2):
    H, W = depth.shape[:2]
    if rgb.shape[:2] != (H, W):
        rgb = cv2.resize(rgb, (W, H))
    K = np.asarray(intrinsics_3x3, dtype=np.float64)
    fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2]
    vs, us = np.mgrid[0:H:stride, 0:W:stride]
    z = depth[vs, us].astype(np.float32)
    valid = (z > 0.1) & (z < depth_trunc)
    z, us, vs = z[valid], us[valid], vs[valid]
    if z.size == 0:
        return None, None
    pts = np.stack([(us - cx) * z / fx, (vs - cy) * z / fy, z], axis=1)
    cols = rgb[vs, us].astype(np.float32) / 255.0
    return pts, cols


def build_3d_scene_figure(scene_pts, scene_cols, traj_dict, title="", subsample=2):
    fig = go.Figure()
    if scene_pts is not None and len(scene_pts) > 0:
        sp = scene_pts[::subsample]
        sr = (scene_cols[::subsample] * 255).astype(np.uint8)
        fig.add_trace(go.Scatter3d(
            x=sp[:, 0], y=sp[:, 1], z=sp[:, 2], mode="markers",
            marker=dict(size=1.5, color=[f"rgb({r},{g},{b})" for r, g, b in sr], opacity=0.6),
            name="Scene", hoverinfo="skip",
        ))
    for label, (xyz, color) in traj_dict.items():
        if xyz is None or len(xyz) < 2:
            continue
        c = np.array(color)
        tpts, tcols = [], []
        sp = _create_sphere_points(xyz[0], radius=0.008, n=400)
        tpts.append(sp)
        tcols.append(np.tile(np.clip(c * 1.3, 0, 1), (len(sp), 1)))
        ep = _create_sphere_points(xyz[-1], radius=0.008, n=400)
        tpts.append(ep)
        tcols.append(np.tile(c * 0.7, (len(ep), 1)))
        for i in range(len(xyz) - 1):
            tube = _create_tube_points(xyz[i], xyz[i + 1], radius=0.003)
            if len(tube) > 0:
                tpts.append(tube)
                tcols.append(np.tile(c, (len(tube), 1)))
        tpts = np.vstack(tpts)
        tcols = np.vstack(tcols)
        tr_rgb = (np.clip(tcols, 0, 1) * 255).astype(np.uint8)
        fig.add_trace(go.Scatter3d(
            x=tpts[:, 0], y=tpts[:, 1], z=tpts[:, 2], mode="markers",
            marker=dict(size=2.5, color=[f"rgb({r},{g},{b})" for r, g, b in tr_rgb], opacity=1.0),
            name=label, hoverinfo="skip",
        ))
    fig.update_layout(
        title=title, height=600,
        scene=dict(
            xaxis_title="X", yaxis_title="Y", zaxis_title="Z",
            aspectmode="data", bgcolor="white",
            camera=dict(eye=dict(x=0, y=0, z=-1.5), up=dict(x=0, y=-1, z=0)),
        ),
        paper_bgcolor="white",
        legend=dict(x=0.01, y=0.99, bgcolor="rgba(255,255,255,0.8)"),
    )
    return fig


def format_conversation(system, user, assistant):
    return (
        "════════ SYSTEM ════════\n"
        f"{system.strip()}\n\n"
        "════════ USER ════════\n"
        f"{user.strip()}\n\n"
        "════════ ASSISTANT ════════\n"
        f"{assistant.strip()}"
    )


# ── Model (loaded once at module scope, eager .to("cuda")) ───────────────────

print(f"Downloading checkpoint {MODEL_ID} …")
CKPT_DIR = snapshot_download(MODEL_ID)
print(f"Checkpoint at {CKPT_DIR}")

# Register the custom Qwen3-VL geometry model class with transformers Auto* .
from hamster3d.model import register_qwen3_vl_geometry  # noqa: E402

try:
    register_qwen3_vl_geometry()
except Exception as e:  # pragma: no cover
    print(f"register_qwen3_vl_geometry: {e!r}")

from transformers import AutoModelForImageTextToText, AutoProcessor  # noqa: E402

print("Loading processor …")
PROCESSOR = AutoProcessor.from_pretrained(CKPT_DIR, trust_remote_code=True)

print("Loading model …")
MODEL = AutoModelForImageTextToText.from_pretrained(
    CKPT_DIR,
    dtype=torch.bfloat16,
    trust_remote_code=True,
).to("cuda")
MODEL.eval()
_PARAM_DTYPE = next(MODEL.parameters()).dtype
print(f"Model loaded ({sum(p.numel() for p in MODEL.parameters()):,} params, dtype={_PARAM_DTYPE})")


def _run_model(image_path, npz_path, query, system_prompt=V5_SYSTEM_PROMPT):
    """Greedy generation matching the reference ModelServer.predict."""
    device = "cuda"
    rgb = np.array(Image.open(image_path).convert("RGB"))
    pcd = np.load(npz_path)["pcd"]  # (H, W, 4) float16
    depth = pcd[:, :, 2].astype(np.float32)

    rgb_tensor = torch.from_numpy(rgb).float().permute(2, 0, 1).unsqueeze(0) / 255.0
    depth_tensor = torch.from_numpy(depth).float().unsqueeze(0)
    geometry_encoder_inputs = [rgb_tensor.to(device=device, dtype=_PARAM_DTYPE)]
    depth_maps = [depth_tensor.to(device=device, dtype=_PARAM_DTYPE)]

    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": [{"type": "text", "text": system_prompt}]})
    messages.append({
        "role": "user",
        "content": [{"type": "image"}, {"type": "text", "text": query}],
    })
    text = PROCESSOR.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    model_inputs = PROCESSOR(
        text=[text],
        images=[Image.open(image_path).convert("RGB")],
        return_tensors="pt",
    ).to(device)
    for _k, _v in list(model_inputs.items()):
        if torch.is_tensor(_v) and torch.is_floating_point(_v):
            model_inputs[_k] = _v.to(_PARAM_DTYPE)
    model_inputs["geometry_encoder_inputs"] = geometry_encoder_inputs
    model_inputs["depth_maps"] = depth_maps

    with torch.inference_mode():
        output_ids = MODEL.generate(
            **model_inputs,
            max_new_tokens=1024,
            do_sample=False,
            temperature=None,
            top_p=None,
        )
    input_len = model_inputs["input_ids"].shape[1]
    generated_ids = output_ids[:, input_len:]
    return PROCESSOR.batch_decode(generated_ids, skip_special_tokens=True)[0]


# ── Inference handler (ZeroGPU) ──────────────────────────────────────────────


@spaces.GPU(duration=120)
def predict(rgb_image, depth_file, instruction, prompt_style):
    """Predict a robot manipulation trajectory / pointing / bbox / VQA answer.

    Args:
        rgb_image: RGB scene image (PIL). Auto-resized to longest edge 640.
        depth_file: metric depth map as a .npy file (float32, meters, aligned to RGB).
        instruction: free-form language instruction (e.g. "Pick up the red cup").
        prompt_style: one of "3D Trajectory", "2D Trajectory", "3D Pointing",
            "2D Pointing", "2D Bounding Box", "General VQA".

    Returns:
        (overlay_image, raw_output_text, conversation_text, plotly_3d_figure)
    """
    if rgb_image is None:
        return None, "Please provide an RGB image.", "", None
    if depth_file is None:
        return None, "Please provide a metric depth .npy file.", "", None
    if not instruction or not instruction.strip():
        return None, "Please enter a task instruction.", "", None

    depth_path = depth_file if isinstance(depth_file, str) else depth_file.name

    tmp_dir = tempfile.mkdtemp(prefix="hamster3d_")
    img_path, npz_path, rgb_resized, depth_resized = prepare_inputs(
        rgb_image, depth_path, tmp_dir
    )
    h, w = rgb_resized.shape[:2]

    human_msg = build_v5_human_message(instruction, prompt_style)
    raw = _run_model(img_path, npz_path, human_msg, system_prompt=V5_SYSTEM_PROMPT)
    conversation = format_conversation(V5_SYSTEM_PROMPT, f"<image>{human_msg}", raw)

    K = default_intrinsics(h, w)
    scene_pts, scene_cols = build_scene_pcd_simple(rgb_resized, depth_resized, K)

    # General VQA → free-form answer, scene with no trajectory.
    if prompt_style == VQA_STYLE:
        fig = build_3d_scene_figure(scene_pts, scene_cols, {}, title=instruction.strip())
        return rgb_resized, raw, conversation, fig

    # 2D bounding box.
    if prompt_style == BBOX_STYLE:
        boxes = parse_bbox_2d(raw)
        viz = visualize_bbox(rgb_resized, boxes)
        fig = build_3d_scene_figure(scene_pts, scene_cols, {}, title=instruction.strip())
        return viz, raw, conversation, fig

    # Pointing → independent numbered points.
    if prompt_style in POINTING_STYLES:
        pts, _ = parse_gt_structured_json(raw)
        viz = visualize_points(rgb_resized, pts) if pts else rgb_resized
        fig = build_3d_scene_figure(scene_pts, scene_cols, {}, title=instruction.strip())
        if pts:
            arr = np.array(pts, dtype=np.float32)
            if prompt_style == "2D Pointing" and depth_resized is not None:
                for i in range(len(arr)):
                    up = int(np.clip(round(arr[i, 0] / 1000 * w), 0, w - 1))
                    vp = int(np.clip(round(arr[i, 1] / 1000 * h), 0, h - 1))
                    arr[i, 2] = float(depth_resized[vp, up])
            xyz = _uvd_to_xyz(arr, K, w, h)
            fig.add_trace(go.Scatter3d(
                x=xyz[:, 0], y=xyz[:, 1], z=xyz[:, 2], mode="markers+text",
                marker=dict(size=8, color="lime", line=dict(width=2, color="white")),
                text=[str(i + 1) for i in range(len(xyz))], textposition="top center",
                name="Points", hoverinfo="skip",
            ))
        return viz, raw, conversation, fig

    # Trajectory (2D / 3D).
    waypoints, actions = parse_gt_structured_json(raw)
    if not waypoints:
        waypoints, actions = parse_trajectory(raw)
    viz_2d = visualize_2d(rgb_resized, waypoints, actions) if waypoints else rgb_resized
    traj_dict = {}
    if waypoints:
        pred_xyz = _uvd_to_xyz(np.array(waypoints, dtype=np.float32), K, w, h)
        traj_dict["Predicted trajectory"] = (pred_xyz, [1.0, 0.2, 0.0])
    fig = build_3d_scene_figure(scene_pts, scene_cols, traj_dict, title=instruction.strip())
    return viz_2d, raw, conversation, fig


# ── Examples browser helpers ─────────────────────────────────────────────────


def _load_example_assets(idx):
    """Return (rgb_pil, depth_npy_path, instruction) for bundled example idx."""
    prefix = f"sample_{int(idx)}"
    rgb_path = os.path.join(EXAMPLES_DIR, f"{prefix}_rgb.png")
    depth_path = os.path.join(EXAMPLES_DIR, f"{prefix}_depth.npy")
    instr_path = os.path.join(EXAMPLES_DIR, f"{prefix}_instruction.txt")
    instruction = ""
    if os.path.isfile(instr_path):
        instruction = open(instr_path).read().strip()
    return rgb_path, depth_path, instruction


def _example_rows():
    rows = []
    for i in range(6):
        rgb_path, depth_path, instruction = _load_example_assets(i)
        if os.path.isfile(rgb_path) and os.path.isfile(depth_path):
            rows.append([rgb_path, depth_path, instruction, V5_DEFAULT_STYLE])
    return rows


def on_depth_upload(depth_file):
    if depth_file is None:
        return None
    try:
        path = depth_file if isinstance(depth_file, str) else depth_file.name
        depth = np.load(path).astype(np.float32)
        return colorize_depth(depth)
    except Exception:
        return None


# ── UI ───────────────────────────────────────────────────────────────────────

CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
.mono textarea { font-family: monospace; font-size: 13px; }
"""

with gr.Blocks(title="3D HAMSTER") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🐹 3D HAMSTER — 3D Trajectory Prediction\n"
            "Depth-aware **Vision-Language-Action** planner (Qwen3-VL-8B + LingBot-Depth "
            "geometry encoder). From a single **RGB image + metric depth map + language "
            "instruction**, it predicts a metric **3D end-effector trajectory** "
            "(`[u, v, depth]` waypoints + gripper states), plus 2D trajectory / pointing / "
            "bounding-box / VQA modes.\n\n"
            "[Paper](https://huggingface.co/papers/2606.31329) · "
            "[Model](https://huggingface.co/DAVIAN-Robotics/3D_HAMSTER) · "
            "[Code](https://github.com/DAVIAN-Robotics/3D_HAMSTER) · "
            "[Project page](https://davian-robotics.github.io/3D_HAMSTER/)\n\n"
            "> ⚠️ Depth must be **metric (meters)** and aligned to the RGB frame. "
            "Disparity / normalized / millimeter depth will degrade the geometry."
        )

        with gr.Row():
            with gr.Column(scale=1):
                rgb_input = gr.Image(label="RGB Image", type="pil", height=320)
                depth_input = gr.File(
                    label="Metric Depth (.npy, float32, meters)", file_types=[".npy"]
                )
                depth_preview = gr.Image(label="Depth preview", type="numpy", height=180)
                instruction_input = gr.Textbox(
                    label="Task instruction",
                    placeholder="e.g. Pick up the red block and place it on the blue plate.",
                    lines=2,
                )
                prompt_style_input = gr.Radio(
                    label="Prompt style",
                    choices=V5_PROMPT_STYLES,
                    value=V5_DEFAULT_STYLE,
                    info="Trajectory: manipulation waypoints | Pointing: object points | "
                         "2D Bounding Box: object box | General VQA: free-form answer",
                )
                run_btn = gr.Button("Predict trajectory", variant="primary", size="lg")

            with gr.Column(scale=1):
                overlay_output = gr.Image(label="2D overlay", type="numpy", height=340)
                plot_3d_output = gr.Plot(label="3D scene + trajectory (rotate / zoom)")

        with gr.Accordion("Model output", open=False):
            raw_output = gr.Textbox(
                label="Raw model output", lines=4, interactive=False,
                elem_classes=["mono"],
            )
            conversation_output = gr.Textbox(
                label="Full conversation", lines=10, interactive=False,
                elem_classes=["mono"],
            )

        gr.Examples(
            examples=_example_rows(),
            inputs=[rgb_input, depth_input, instruction_input, prompt_style_input],
            outputs=[overlay_output, raw_output, conversation_output, plot_3d_output],
            fn=predict,
            cache_examples=False,
            run_on_click=True,
            label="Bundled examples (RGB + depth + instruction)",
        )

    depth_input.change(fn=on_depth_upload, inputs=[depth_input], outputs=[depth_preview])
    run_btn.click(
        fn=predict,
        inputs=[rgb_input, depth_input, instruction_input, prompt_style_input],
        outputs=[overlay_output, raw_output, conversation_output, plot_3d_output],
        api_name="predict",
    )


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