"""DynamicVLA (hzxie/dynamic-vla-DOM) action-chunk demo."""

from __future__ import annotations

import importlib.util
import json
import os
import subprocess
import sys
import time
from pathlib import Path

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

import spaces  # noqa: E402  — must precede any CUDA-touching import

if importlib.util.find_spec("lerobot") is None:
    subprocess.check_call(
        [sys.executable, "-m", "pip", "install", "--no-deps", "lerobot==0.3.3"]
    )

import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402
import plotly.graph_objects as go  # noqa: E402
import torch  # noqa: E402
from huggingface_hub import snapshot_download  # noqa: E402
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature  # noqa: E402
from PIL import Image  # noqa: E402

from policies.dynamicvla.configuration_dynamicvla import DynamicVLAConfig  # noqa: E402
from policies.dynamicvla.modeling_dynamicvla import (  # noqa: E402
    DynamicVLAPolicy,
    load_dynamicvla,
)

import gradio as gr  # noqa: E402

MODEL_ID = "hzxie/dynamic-vla-DOM"
IMG_H, IMG_W = 360, 480
N_OBS = 2
ACTION_COLS = ["x", "y", "z", "roll", "pitch", "yaw", "gripper"]
ROOT = Path(__file__).resolve().parent
EXAMPLES = ROOT / "examples"

FEATURE_TYPES = {
    "STATE": FeatureType.STATE,
    "VISUAL": FeatureType.VISUAL,
    "ACTION": FeatureType.ACTION,
}


def _features(spec: dict) -> dict:
    return {
        key: PolicyFeature(type=FEATURE_TYPES[ft["type"]], shape=tuple(ft["shape"]))
        for key, ft in spec.items()
    }


def _identity_stats(cfg: DynamicVLAConfig) -> dict[str, dict[str, torch.Tensor]]:
    """MEAN_STD with mean=0 / std=1 is a passthrough. Checkpoint strips norm buffers."""
    stats: dict[str, dict[str, torch.Tensor]] = {}
    for name, feat in {**cfg.input_features, **cfg.output_features}.items():
        if feat.type in (FeatureType.STATE, FeatureType.ACTION):
            shape = tuple(feat.shape)
            stats[name] = {
                "mean": torch.zeros(shape, dtype=torch.float32),
                "std": torch.ones(shape, dtype=torch.float32),
                "min": torch.full(shape, -1.0, dtype=torch.float32),
                "max": torch.ones(shape, dtype=torch.float32),
            }
    return stats


def _build_config(ckpt_dir: str) -> DynamicVLAConfig:
    with open(os.path.join(ckpt_dir, "config.json"), encoding="utf-8") as fh:
        raw = json.load(fh)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    cfg = DynamicVLAConfig(
        input_features=_features(raw["input_features"]),
        output_features=_features(raw["output_features"]),
        device=device,
    )
    skip = {"type", "device", "input_features", "output_features"}
    for key, value in raw.items():
        attr = key.lower()
        if attr in skip or value is None or not hasattr(cfg, attr):
            continue
        if attr == "normalization_mapping" and isinstance(value, dict):
            value = {
                k: (NormalizationMode[v] if isinstance(v, str) else v)
                for k, v in value.items()
            }
        setattr(cfg, attr, value)
    cfg.enable_streaming = False
    return cfg


def load_policy() -> DynamicVLAPolicy:
    ckpt_dir = snapshot_download(MODEL_ID)
    cfg = _build_config(ckpt_dir)
    policy = DynamicVLAPolicy(cfg, dataset_stats=_identity_stats(cfg))
    load_dynamicvla(
        policy,
        os.path.join(ckpt_dir, "model.safetensors"),
        device="cpu",
        checkpoint_keys_mapping="model._orig_mod.//model.",
    )
    policy.eval()
    return policy.to("cuda")


POLICY = load_policy()


def _as_image(img) -> Image.Image | None:
    if img is None:
        return None
    if isinstance(img, Image.Image):
        return img.convert("RGB")
    if isinstance(img, np.ndarray):
        if img.ndim == 3 and img.shape[-1] == 4:
            img = img[..., :3]
        if img.dtype != np.uint8:
            img = np.clip(img, 0, 255).astype(np.uint8) if img.max() > 1.5 else (
                np.clip(img * 255.0, 0, 255).astype(np.uint8)
            )
        return Image.fromarray(img).convert("RGB")
    return Image.open(img).convert("RGB")


def _to_nchw(img: Image.Image) -> torch.Tensor:
    resized = img.resize((IMG_W, IMG_H), Image.BILINEAR)
    arr = np.asarray(resized, dtype=np.float32) / 255.0
    return torch.from_numpy(arr).permute(2, 0, 1)


def _stack_obs(current: Image.Image, previous: Image.Image | None) -> torch.Tensor:
    cur = _to_nchw(current)
    prev = _to_nchw(previous) if previous is not None else cur
    return torch.stack([prev, cur], dim=0)  # (n_obs, C, H, W)


def _plot_path(actions: np.ndarray) -> go.Figure:
    xs, ys, zs = actions[:, 0], actions[:, 1], actions[:, 2]
    fig = go.Figure(
        data=[
            go.Scatter3d(
                x=xs,
                y=ys,
                z=zs,
                mode="lines+markers",
                marker={"size": 4, "color": np.arange(len(xs)), "colorscale": "Viridis"},
                line={"width": 5, "color": "#2ecc71"},
                name="EE path",
            ),
            go.Scatter3d(
                x=[xs[0]],
                y=[ys[0]],
                z=[zs[0]],
                mode="markers",
                marker={"size": 8, "color": "#27ae60"},
                name="start",
            ),
            go.Scatter3d(
                x=[xs[-1]],
                y=[ys[-1]],
                z=[zs[-1]],
                mode="markers",
                marker={"size": 8, "color": "#e74c3c"},
                name="end",
            ),
        ]
    )
    fig.update_layout(
        template="plotly_dark",
        height=420,
        margin={"l": 0, "r": 0, "t": 30, "b": 0},
        scene={
            "xaxis_title": "x (m)",
            "yaxis_title": "y (m)",
            "zaxis_title": "z (m)",
            "aspectmode": "data",
        },
        title="Predicted 20-step end-effector chunk",
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
        legend={"orientation": "h"},
    )
    return fig


def _summarize(actions: np.ndarray, elapsed: float, instruction: str) -> str:
    delta = actions[-1, :3] - actions[0, :3]
    grip = actions[:, -1]
    return (
        f"**Instruction:** {instruction.strip()}\n\n"
        f"**Chunk:** {len(actions)} steps · **{elapsed:.2f}s** GPU\n\n"
        f"- Start xyz: `{actions[0, :3].round(4).tolist()}`\n"
        f"- End xyz: `{actions[-1, :3].round(4).tolist()}`\n"
        f"- Net Δxyz: `{delta.round(4).tolist()}`\n"
        f"- Gripper: min `{grip.min():.3f}` → max `{grip.max():.3f}` "
        f"(last `{grip[-1]:.3f}`)\n"
        f"- Rotation (last rpy): `{actions[-1, 3:6].round(4).tolist()}`"
    )


@spaces.GPU(duration=60)
def predict_action_chunk(
    wrist: np.ndarray | Image.Image | None,
    opposite: np.ndarray | Image.Image | None,
    wrist_prev: np.ndarray | Image.Image | None,
    opposite_prev: np.ndarray | Image.Image | None,
    instruction: str,
    x: float,
    y: float,
    z: float,
    roll: float,
    pitch: float,
    yaw: float,
    apply_delta: bool,
) -> tuple[pd.DataFrame, go.Figure, str]:
    """Predict a 20-step DynamicVLA action chunk from dual-camera frames."""
    if not instruction or not instruction.strip():
        raise gr.Error("Provide a language instruction.")
    wrist_img = _as_image(wrist)
    opp_img = _as_image(opposite)
    if wrist_img is None or opp_img is None:
        raise gr.Error("Upload both wrist and opposite camera frames.")

    device = "cuda"
    wrist_t = _stack_obs(wrist_img, _as_image(wrist_prev)).unsqueeze(0).to(device)
    opp_t = _stack_obs(opp_img, _as_image(opposite_prev)).unsqueeze(0).to(device)
    state = torch.tensor(
        [[[x, y, z, roll, pitch, yaw]] * N_OBS], dtype=torch.float32, device=device
    )

    batch = {
        "observation.images.wrist_cam": wrist_t,
        "observation.images.opst_cam": opp_t,
        "observation.state": state,
        "task": [instruction.strip()],
    }

    POLICY.reset()
    tick = time.perf_counter()
    with torch.inference_mode():
        actions = POLICY.predict_action_chunk(batch)
        if apply_delta and getattr(POLICY.config, "use_delta_action", True):
            actions = actions.clone()
            actions[..., :6] = actions[..., :6] + state[:, -1:, :6]
    elapsed = time.perf_counter() - tick
    acts = actions[0].detach().float().cpu().numpy()

    table = pd.DataFrame(acts, columns=ACTION_COLS)
    table.insert(0, "step", np.arange(len(table)))
    return table, _plot_path(acts), _summarize(acts, elapsed, instruction)


def _example_row(stem: str, instruction: str) -> list:
    img = str(EXAMPLES / f"{stem}.png")
    return [img, img, None, None, instruction, 0.40, 0.00, 0.30, 0.0, 0.0, 0.0, True]


GALLERY_MD = """
## DynamicVLA on DOM

**DynamicVLA** (0.4B, SmolLM2-360M + FastViT) is a VLA for *moving* objects.
It adds **Continuous Inference** and **Latent-aware Action Streaming** so the
policy does not freeze between action chunks.

This Space runs the official [`hzxie/dynamic-vla-DOM`](https://huggingface.co/hzxie/dynamic-vla-DOM)
checkpoint and predicts a **20-step** 7-DoF end-effector chunk
(`[x, y, z, roll, pitch, yaw, gripper]`).

Closed-loop Isaac Lab eval is *not* hosted here — use
[hzxie/DynamicVLA](https://github.com/hzxie/DynamicVLA) for that.

| | |
|---|---|
| Paper | [arXiv:2601.22153](https://arxiv.org/abs/2601.22153) |
| Dataset | [`hzxie/DOM`](https://huggingface.co/datasets/hzxie/DOM) — 200K episodes, 2.8K scenes, 206 objects |
| Weights | [`hzxie/dynamic-vla-DOM`](https://huggingface.co/hzxie/dynamic-vla-DOM) |
| Project | [infinitescript.com/project/dynamic-vla](https://www.infinitescript.com/project/dynamic-vla/) |
| Spotlight | [YouTube](https://youtu.be/NmJnHcI04_Q) |

<iframe width="100%" height="360" src="https://www.youtube.com/embed/NmJnHcI04_Q"
title="DynamicVLA spotlight" frameborder="0"
allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture"
allowfullscreen></iframe>
"""


def build_ui() -> gr.Blocks:
    with gr.Blocks(title="DynamicVLA · DOM") as demo:
        gr.Markdown(
            "# DynamicVLA — DOM action-chunk demo\n"
            "0.4B VLA for dynamic object manipulation. "
            "Upload wrist + scene cameras, write an instruction, get a 20-step EE chunk."
        )
        with gr.Tabs():
            with gr.Tab("Predict"):
                with gr.Row():
                    with gr.Column():
                        wrist = gr.Image(label="Wrist camera (current)", type="numpy")
                        opposite = gr.Image(
                            label="Opposite / scene camera (current)", type="numpy"
                        )
                        with gr.Accordion("Previous frames (optional, n_obs=2)", open=False):
                            wrist_prev = gr.Image(
                                label="Wrist camera (t−1)", type="numpy"
                            )
                            opposite_prev = gr.Image(
                                label="Opposite camera (t−1)", type="numpy"
                            )
                        instruction = gr.Textbox(
                            label="Language instruction",
                            placeholder="Pick up the rolling cylinder and place it onto the wooden block.",
                            lines=2,
                        )
                        with gr.Accordion("Current EE state (meters / rad)", open=False):
                            with gr.Row():
                                x = gr.Number(value=0.40, label="x")
                                y = gr.Number(value=0.00, label="y")
                                z = gr.Number(value=0.30, label="z")
                            with gr.Row():
                                roll = gr.Number(value=0.0, label="roll")
                                pitch = gr.Number(value=0.0, label="pitch")
                                yaw = gr.Number(value=0.0, label="yaw")
                            apply_delta = gr.Checkbox(
                                value=True,
                                label="Add delta actions to current EE state",
                            )
                        run = gr.Button("Predict action chunk", variant="primary")
                    with gr.Column():
                        summary = gr.Markdown("Upload both views and run.")
                        path = gr.Plot(label="EE trajectory")
                        table = gr.Dataframe(label="Action chunk (20 × 7)")

                inputs = [
                    wrist,
                    opposite,
                    wrist_prev,
                    opposite_prev,
                    instruction,
                    x,
                    y,
                    z,
                    roll,
                    pitch,
                    yaw,
                    apply_delta,
                ]
                run.click(
                    fn=predict_action_chunk,
                    inputs=inputs,
                    outputs=[table, path, summary],
                )
                gr.Examples(
                    examples=[
                        _example_row(
                            "franka-coffee",
                            "Pick up the rolling cylinder and place it onto the wooden block.",
                        ),
                        _example_row(
                            "piper-sesame",
                            "Grasp the rolling roasted sesame container and place it onto the blue frisbee.",
                        ),
                        _example_row(
                            "franka-tennis",
                            "Get hold of the moving tennis ball and position it into the paper bowl.",
                        ),
                    ],
                    inputs=inputs,
                    outputs=[table, path, summary],
                    fn=predict_action_chunk,
                    label="DOM-style prompts (official comparison stills as both views)",
                    cache_examples=True,
                    cache_mode="lazy",
                )

            with gr.Tab("About"):
                gr.Markdown(GALLERY_MD)
                if (EXAMPLES / "teaser.webp").exists():
                    gr.Image(
                        value=str(EXAMPLES / "teaser.webp"),
                        label="Official teaser",
                        interactive=False,
                    )
                gr.Markdown(
                    "```bibtex\n"
                    "@article{xie2026dynamicvla,\n"
                    "  title   = {DynamicVLA: A Vision-Language-Action Model for Dynamic Object Manipulation},\n"
                    "  author  = {Xie, Haozhe and Wen, Beichen and Zheng, Jiarui and Chen, Zhaoxi\n"
                    "             and Hong, Fangzhou and Diao, Haiwen and Liu, Ziwei},\n"
                    "  journal = {arXiv preprint arXiv:2601.22153},\n"
                    "  year    = {2026}\n"
                    "}\n"
                    "```"
                )
    return demo


demo = build_ui()

if __name__ == "__main__":
    demo.launch(
        mcp_server=True,
        theme=gr.themes.Soft(primary_hue="green", neutral_hue="zinc"),
    )