"""DM0.5 Vision-Language-Action demo.

Loads the Dexmal/DM05 open-world VLA foundation model and predicts a chunk of
future bimanual robot actions from three camera views + a natural-language
instruction. The predicted 50-step x 14-dim action trajectory is rendered as
per-joint / per-gripper plots.

Inference logic mirrors OpenDM's `DM05InferenceConfig` (github.com/dexmal/opendm):
the same PixelTransform -> Normalize(state) -> ChatTokenization pipeline feeds
`model.inference_action`, and the flow-matching output is denormalized and
converted from relative deltas back to absolute joint targets.
"""

import os

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

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

import numpy as np
import torch
import gradio as gr
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from huggingface_hub import snapshot_download
from transformers import AutoProcessor

from opendm.constants.robot import ROBOT_STATE_DESCS, RobotType, RobotStateDesc
from opendm.data.augmentations import NoAugmentationPipeline
from opendm.data.transforms import (
    ActionAbsolute,
    ChatTokenization,
    Denormalize,
    Normalize,
    PixelTransform,
    Pipeline,
    ToDevice,
)
from opendm.model.dm05.dm05_arch import DM05Config, DM05ForConditionalGeneration

MODEL_ID = "Dexmal/DM05"
CHUNK_SIZE = 50           # action horizon predicted by the model
OUTPUT_ACTION_DIM = 14    # bimanual: 6 joints + gripper per arm
N_BINS = 256
MODEL_MAX_LENGTH = 1024
IMAGE_KEYS = ["images_1", "images_2", "images_3"]  # head, left wrist, right wrist
ROBOT_TYPE = "DOS W1"
STATE_DESC = ROBOT_STATE_DESCS[RobotType.DOS_W1]
# 14-dim resting state used by OpenDM's demo (grippers open = 1.0).
DEFAULT_STATE = "[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]"

DEVICE = "cuda"

# ---------------------------------------------------------------------------
# Model + transforms — loaded once at module scope (ZeroGPU packs to disk).
# ---------------------------------------------------------------------------
print(f"Downloading {MODEL_ID} ...")
CKPT_DIR = snapshot_download(MODEL_ID)
NORM_STATS_PATH = os.path.join(CKPT_DIR, "norm_stats.json")

print("Building DM05 model ...")
_config = DM05Config.from_pretrained(CKPT_DIR)
_config.chunk_size = CHUNK_SIZE


def _set_attn(cfg, value):
    for attr in ("_attn_implementation", "attn_implementation",
                 "_attn_implementation_internal"):
        try:
            object.__setattr__(cfg, attr, value)
        except Exception:
            try:
                setattr(cfg, attr, value)
            except Exception:
                pass


def _pin_attn_before_init(cfg):
    """Pin attention backends on the config tree before model construction.

    The saved DM05 config bakes ``flash_attention_2`` into the VLM/vision
    configs, which transformers validates at __init__ time (before we can call
    ``set_attention_implementation``). flash_attn isn't installed on the
    ZeroGPU Blackwell runtime. The top-level DM05 wrapper only declares
    ``eager`` support, while the Gemma3 sub-modules accept ``sdpa`` — so pin
    the wrapper to eager and the VLM sub-configs to sdpa.
    """
    _set_attn(cfg, "eager")
    for sub in ("vlm_config", "vision_config", "text_config"):
        child = getattr(cfg, sub, None)
        if child is not None and hasattr(child, "__dict__"):
            _set_attn(child, "sdpa")
            for sub2 in ("vision_config", "text_config"):
                gchild = getattr(child, sub2, None)
                if gchild is not None and hasattr(gchild, "__dict__"):
                    _set_attn(gchild, "sdpa")


_pin_attn_before_init(_config)

model = DM05ForConditionalGeneration.from_pretrained(
    CKPT_DIR,
    config=_config,
    torch_dtype=torch.bfloat16,
)
# Inference-time attention backends. OpenDM forces LLM to eager for inference;
# vision/action use SDPA (torch-native, works on ZeroGPU Blackwell sm_120 —
# avoids the flash_attn dependency entirely).
model.set_attention_implementation(
    llm_attn_implementation="eager",
    vision_attn_implementation="sdpa",
    action_attn_implementation="sdpa",
    bf16=True,
)
model.to(DEVICE)
model.eval()

processor = AutoProcessor.from_pretrained(CKPT_DIR, trust_remote_code=True)

input_transform = Pipeline(
    [
        PixelTransform(
            transform_pipeline=NoAugmentationPipeline(),
            image_keys=IMAGE_KEYS,
        ),
        Normalize(
            norm_stats_path=NORM_STATS_PATH,
            norm_keys=["state"],
            use_quantiles=True,
        ),
        ChatTokenization(
            processor=processor,
            n_bins=N_BINS,
            max_length=MODEL_MAX_LENGTH,
            image_keys=IMAGE_KEYS,
            add_state=True,
        ),
        ToDevice(device=DEVICE),
    ]
)
# DM05 is trained in RELATIVE action mode -> convert deltas back to absolute.
output_transform = Pipeline(
    [
        Denormalize(
            norm_stats_path=NORM_STATS_PATH,
            norm_keys=["action"],
            use_quantiles=True,
        ),
        ActionAbsolute(),
    ]
)
print("Model ready.")

JOINT_LABELS = [f"J{i + 1}" for i in range(6)]


def _plot_trajectory(actions: np.ndarray) -> str:
    """Render the predicted 14-dim action chunk to a PNG and return its path.

    actions: (chunk_size, 14) array — [L arm joints x6, L gripper,
             R arm joints x6, R gripper].
    """
    steps = np.arange(actions.shape[0])
    fig, axes = plt.subplots(2, 2, figsize=(12, 8))
    fig.suptitle("DM0.5 predicted action trajectory", fontsize=15, fontweight="bold")

    specs = [
        (axes[0, 0], "Left arm — joint targets", actions[:, 0:6], JOINT_LABELS),
        (axes[0, 1], "Right arm — joint targets", actions[:, 7:13], JOINT_LABELS),
        (axes[1, 0], "Left gripper", actions[:, 6:7], ["gripper"]),
        (axes[1, 1], "Right gripper", actions[:, 13:14], ["gripper"]),
    ]
    for ax, title, data, labels in specs:
        for j in range(data.shape[1]):
            ax.plot(steps, data[:, j], marker="", linewidth=1.8, label=labels[j])
        ax.set_title(title)
        ax.set_xlabel("prediction step")
        ax.set_ylabel("value")
        ax.grid(True, alpha=0.3)
        ax.legend(loc="best", fontsize=8, ncol=2)

    fig.tight_layout(rect=[0, 0, 1, 0.96])
    out = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    fig.savefig(out.name, dpi=110, bbox_inches="tight")
    plt.close(fig)
    return out.name


def _parse_state(state_text: str) -> np.ndarray:
    try:
        arr = np.array(json.loads(state_text), dtype=np.float32)
    except Exception as e:
        raise gr.Error(f"Could not parse robot state as a JSON array: {e}")
    if arr.shape[0] != OUTPUT_ACTION_DIM:
        raise gr.Error(
            f"Robot state must have {OUTPUT_ACTION_DIM} values "
            f"(6 joints + gripper per arm); got {arr.shape[0]}."
        )
    return arr


@spaces.GPU(duration=120)
def predict(
    image_head,
    image_left_wrist,
    image_right_wrist,
    instruction: str,
    state_text: str = DEFAULT_STATE,
    control_mode: str = "joint",
    speed: str = "normal",
    diffusion_steps: int = 10,
):
    """Predict a bimanual robot action trajectory from cameras + an instruction.

    Args:
        image_head: Head (overhead) camera view (PIL image).
        image_left_wrist: Left-wrist camera view (PIL image).
        image_right_wrist: Right-wrist camera view (PIL image).
        instruction: Natural-language task, e.g. "pick up the roller".
        state_text: JSON array of the current 14-dim robot state.
        control_mode: Text conditioning field required by the base DM05 model.
        speed: Overall motion speed conditioning ("slow"/"normal"/"fast").
        diffusion_steps: Number of flow-matching integration steps.

    Returns:
        A trajectory plot (image) and the raw predicted action array as JSON text.
    """
    if image_head is None or image_left_wrist is None or image_right_wrist is None:
        raise gr.Error("Please provide all three camera images (head, left wrist, right wrist).")
    if not instruction or not instruction.strip():
        raise gr.Error("Please enter a task instruction.")

    state = _parse_state(state_text)

    data = {
        "images_1": image_head.convert("RGB"),
        "images_2": image_left_wrist.convert("RGB"),
        "images_3": image_right_wrist.convert("RGB"),
        "prompt": instruction.strip(),
        "state": state,
        "meta_data": {
            "robot_type": ROBOT_TYPE,
            "speed": speed or None,
            "control_mode": control_mode or None,
            "state_desc": STATE_DESC,
        },
    }

    t0 = time.perf_counter()
    model_input = input_transform(data)
    with torch.no_grad():
        actions = model.inference_action(
            input_ids=model_input["input_ids"],
            attention_mask=model_input["attention_mask"],
            pixel_values=model_input["pixel_values"],
            token_type_ids=model_input["token_type_ids"],
            diffusion_steps=int(diffusion_steps),
        )
    raw_actions = (
        actions.to(torch.float32).detach().cpu().numpy()[0, :, :OUTPUT_ACTION_DIM]
    )
    outputs = output_transform(
        {"action": raw_actions, "state": state, "meta_data": data["meta_data"]}
    )
    final_actions = np.asarray(outputs["action"], dtype=np.float32)
    elapsed = time.perf_counter() - t0

    plot_path = _plot_trajectory(final_actions)
    summary = {
        "instruction": instruction.strip(),
        "chunk_size": int(final_actions.shape[0]),
        "action_dim": int(final_actions.shape[1]),
        "inference_seconds": round(elapsed, 2),
        "actions": np.round(final_actions, 4).tolist(),
    }
    return plot_path, json.dumps(summary, indent=2)


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1150px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

DESCRIPTION = """
# 🤖 DM0.5 — Vision-Language-Action robot control

[**Dexmal/DM05**](https://huggingface.co/Dexmal/DM05) is an open-world Vision-Language-Action (VLA) foundation
model (Gemma3-4B backbone + 680M action expert). Given three robot camera views and a natural-language
instruction, it predicts a chunk of **50 future bimanual actions** (6 joint targets + 1 gripper command per arm)
via flow-matching.

Provide the three camera images and a task instruction, then inspect the predicted joint / gripper trajectories.
Built with [OpenDM](https://github.com/dexmal/opendm). Robot type: **DOS W1** (14-dim state).
"""


with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(DESCRIPTION)

        with gr.Row():
            image_head = gr.Image(label="Head camera (images_1)", type="pil", height=240)
            image_left = gr.Image(label="Left-wrist camera (images_2)", type="pil", height=240)
            image_right = gr.Image(label="Right-wrist camera (images_3)", type="pil", height=240)

        with gr.Row():
            instruction = gr.Textbox(
                label="Task instruction",
                placeholder="e.g. Hold the roller with both arms to pick it up",
                scale=4,
            )
            run = gr.Button("Predict actions", variant="primary", scale=1)

        with gr.Accordion("Advanced settings", open=False):
            state_text = gr.Textbox(
                label="Robot state (14-dim JSON array: 6 joints + gripper per arm)",
                value=DEFAULT_STATE,
            )
            with gr.Row():
                control_mode = gr.Textbox(label="Control mode", value="joint")
                speed = gr.Textbox(label="Overall speed", value="normal")
            diffusion_steps = gr.Slider(
                label="Flow-matching integration steps",
                minimum=1, maximum=20, value=10, step=1,
            )

        with gr.Row():
            output_plot = gr.Image(label="Predicted action trajectory", type="filepath")
            output_json = gr.Code(label="Raw predicted actions (JSON)", language="json")

        gr.Examples(
            examples=[
                [
                    "examples/cam_high.jpg",
                    "examples/cam_left_wrist.jpg",
                    "examples/cam_right_wrist.jpg",
                    "Hold the roller for smoothing materials with both arms to pick it up",
                ],
            ],
            inputs=[image_head, image_left, image_right, instruction],
            outputs=[output_plot, output_json],
            fn=predict,
            cache_examples=True,
            cache_mode="lazy",
        )

    run.click(
        fn=predict,
        inputs=[
            image_head, image_left, image_right, instruction,
            state_text, control_mode, speed, diffusion_steps,
        ],
        outputs=[output_plot, output_json],
        api_name="predict",
    )


if __name__ == "__main__":
    demo.launch(mcp_server=True)