"""ERVLA — Embodied Chain-of-Thought VLA policy: predict a robot action chunk from camera views.

Loads `ERVLA/liberoplus` (the LIBERO-Plus-tuned MiBoT checkpoint: Qwen3-VL-4B backbone + a
DiT flow-matching action head) and replicates the reference deployment path from
`src/server/deploy_hf_liberoplus.py` in the ERVLA repo 1:1 — same prompt template, same raw
8-dim proprioceptive state zero-padded to the model's 60-dim state window, same manual
action mask (translation+rotation+gripper only), same 5-step Euler flow integration, and
the same LIBERO `libero_all` normalization stats for decoding actions.

Inputs are two LIBERO-style camera views (base + left wrist) plus the instruction; the
output is the predicted end-effector action chunk (Dxyz translation, Drotation, gripper
command) rendered as a chart + table. No simulator is bundled — LIBERO-Plus / VLABench
cannot run inside a Space, so the action chunk itself is the demo output.
"""

import os

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

import spaces  # noqa: E402  (must precede torch)
import time  # noqa: E402
from typing import Any, Dict, List, Tuple  # noqa: E402

import gradio as gr  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
from PIL import Image  # noqa: E402
from transformers import AutoModel, AutoProcessor  # noqa: E402

MODEL_ID = "ERVLA/liberoplus"
NUM_FLOW_STEPS = 5  # reference server default --num-steps
STATE_LENGTH, STATE_DIM, ACTION_DIM = 1, 60, 60  # from config.json

# libero_all normalization stats, configs/data/sources/libero.yaml in the ERVLA repo.
ACTION_MEAN = np.array([0.0628, 0.0868, -0.0904, 0.0005, 0.0056, -0.0052, -0.0496], dtype=np.float32)
ACTION_STD = np.array([0.3355, 0.3784, 0.4447, 0.0392, 0.0634, 0.0780, 0.9988], dtype=np.float32)

LIBERO_ROBOT_TYPES = ["libero_spatial", "libero_object", "libero_goal", "libero_10"]

# Dataviz reference palette: validated categorical slots 1-3, light + dark steps.
SERIES = {
    "light": ["#2a78d6", "#eb6834", "#1baf7a"],
    "dark": ["#3987e5", "#d95926", "#199e70"],
}
INK = {"light": "#0b0b0b", "dark": "#ffffff"}
GRID = {"light": "#e1e0d9", "dark": "#2c2c2a"}
SURFACE = {"light": "#fcfcfb", "dark": "#1a1a19"}

processor: Any = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True, use_fast=False)


def _load_model() -> Any:
    """Load with eager attention.

    The custom `MiBoTForActionGeneration` class declares `_supports_flash_attn` but not
    `_supports_sdpa`, so under transformers 4.57.1 the only universally available
    backends are flash_attention_2 and eager. The prebuilt Blackwell flash-attn wheel
    (pt212) has a torch-ABI mismatch against the runtime's preinstalled torch 2.13
    (`undefined symbol: ...materialize_cow_storage`), so eager it is — transformers'
    own suggested fallback. The workload is a single forward pass (~700 vision+text
    tokens, /no_cot, no autoregressive decoding) + 5 small DiT flow steps, where eager
    attention costs only a few hundred ms over FA2.
    """
    m: Any = AutoModel.from_pretrained(
        MODEL_ID,
        trust_remote_code=True,
        dtype=torch.bfloat16,
        attn_implementation="eager",
    )
    print("[load] attention backend: eager", flush=True)
    return m.to("cuda")


model = _load_model()
model.eval()
MODEL_DTYPE = next(model.parameters()).dtype
print(f"[load] {MODEL_ID} on cuda, dtype={MODEL_DTYPE}", flush=True)


def _messages(base_image: Image.Image, wrist_image: Image.Image, instruction: str) -> List[Dict[str, Any]]:
    """Build the chat messages exactly as the reference deployment server does."""
    content = [
        {"type": "text", "text": "The following observations are captured from multiple views.\n# Base View\n"},
        {"type": "image", "image": base_image},
        {"type": "text", "text": "\n# Left-Wrist View\n"},
        {"type": "image", "image": wrist_image},
        {"type": "text", "text": f"\nGenerate robot actions for the task:\n{instruction} /no_cot"},
    ]
    return [
        {"role": "user", "content": content},
        {"role": "assistant", "content": [{"type": "text", "text": "<cot></cot>"}]},
    ]


def _to_cuda_dict(values: Dict[str, Any]) -> Dict[str, Any]:
    """Move tensors to CUDA; floating-point tensors take the model dtype (bf16)."""
    moved: Dict[str, Any] = {}
    for key, value in values.items():
        if value is None:
            continue
        if torch.is_tensor(value):
            dtype = MODEL_DTYPE if value.is_floating_point() else None
            moved[key] = value.to("cuda", dtype=dtype)
        else:
            moved[key] = value
    return moved


def build_figure(chunk: np.ndarray):
    """Line chart of the Dx/Dy/Dz translation deltas across the predicted chunk.

    Dataviz rules applied: 3 series use validated categorical slots 1-3; 2px lines;
    >=8px markers with a 2px surface ring; hairline solid gridlines; one axis; a legend
    is present; direct labels only on the final point (never every point); text in ink
    tokens rather than series colors. Rotation + gripper live in the table view (the
    chart's twin), since mixing units on one axis would break the one-axis rule.
    """
    import plotly.graph_objects as go

    mode = "light"
    colors, ink, grid, surface = SERIES[mode], INK[mode], GRID[mode], SURFACE[mode]
    labels = ["Δx (m)", "Δy (m)", "Δz (m)"]
    steps = list(range(1, chunk.shape[0] + 1))
    last = chunk.shape[0] - 1

    fig = go.Figure()
    for i in range(3):
        fig.add_trace(
            go.Scatter(
                x=steps,
                y=chunk[:, i].tolist(),
                mode="lines+markers+text",
                name=labels[i],
                text=[f"{chunk[last, i]:+.3f}" if k == last else "" for k in range(chunk.shape[0])],
                textposition="top right",
                textfont=dict(color=ink, size=10),
                line=dict(width=2, color=colors[i]),
                marker=dict(size=8, color=colors[i], line=dict(width=2, color=surface)),
                hovertemplate="%{y:+.4f} m<extra>" + labels[i] + "</extra>",
            )
        )
    gripper_first = chunk[0, 6]
    gripper_last = chunk[last, 6]
    fig.update_layout(
        title=dict(
            text=(
                f"Predicted {chunk.shape[0]}-step action chunk · gripper cmd "
                f"{gripper_first:+.2f} → {gripper_last:+.2f} (− = close, + = open)"
            ),
            font=dict(color=ink, size=13),
        ),
        xaxis=dict(title="step", showgrid=False, tickfont=dict(color=ink, size=10),
                   linecolor=grid, zeroline=False),
        yaxis=dict(title="Δposition (m)", showgrid=True, gridcolor=grid, gridwidth=1,
                   tickfont=dict(color=ink, size=10), linecolor=grid, zeroline=False),
        plot_bgcolor=surface,
        paper_bgcolor=surface,
        font=dict(color=ink),
        showlegend=True,
        legend=dict(orientation="h", y=1.12, font=dict(color=ink, size=10)),
        margin=dict(l=40, r=20, t=60, b=40),
        height=340,
        uirevision="chunk",
    )
    return fig


def chunk_table(chunk: np.ndarray) -> List[List[float]]:
    """Full predicted chunk as the chart's table twin (all 7 dims per step)."""
    rounded = np.round(chunk.astype(np.float64), 4)
    return [[i + 1, *row.tolist()] for i, row in enumerate(rounded)]


def format_csv(chunk: np.ndarray) -> str:
    """Serialize the predicted chunk as CSV (step, dx, dy, dz, droll, dpitch, dyaw, gripper)."""
    lines = ["step,dx,dy,dz,droll,dpitch,dyaw,gripper"]
    for i, row in enumerate(chunk):
        lines.append(f"{i + 1}," + ",".join(f"{v:+.4f}" for v in row.tolist()))
    return "\n".join(lines)


@spaces.GPU(duration=10)
def predict(
    base_view: Image.Image,
    wrist_view: Image.Image,
    instruction: str,
    robot_type: str = "libero_10",
    state_ee_x: float = -0.19947,
    state_ee_y: float = -0.00631,
    state_ee_z: float = 1.18217,
    state_aa1: float = 3.13143,
    state_aa2: float = 0.00487,
    state_aa3: float = -0.00294,
    state_finger_l: float = 0.03622,
    state_finger_r: float = -0.03623,
    seed: int = 42,
    chunk_length: int = 10,
) -> Tuple[Any, List[List[float]], str, str]:
    """Predict a LIBERO end-effector action chunk from two camera views + an instruction.

    Args:
        base_view: base-view camera image (LIBERO agentview orientation).
        wrist_view: left-wrist camera image.
        instruction: the task, e.g. "Turn on the stove and put the moka pot on it."
        robot_type: LIBERO suite the scene comes from (libero_spatial / libero_object /
            libero_goal / libero_10).
        state_ee_x, state_ee_y, state_ee_z: end-effector position (m).
        state_aa1, state_aa2, state_aa3: end-effector orientation, axis-angle (rad).
        state_finger_l, state_finger_r: gripper finger joint positions (rad).
        seed: RNG seed for the flow-matching sampler.
        chunk_length: number of predicted action steps (10 = LIBERO chunk).

    Returns:
        (plotly figure of the Δx/Δy/Δz deltas, chunk table, CSV text, status line).
    """
    t0 = time.time()
    instruction = str(instruction or "").strip()
    if not instruction:
        raise gr.Error("Provide a task instruction.")
    if not isinstance(base_view, Image.Image) or not isinstance(wrist_view, Image.Image):
        raise gr.Error("Provide both camera views (base + left wrist).")
    if robot_type not in LIBERO_ROBOT_TYPES:
        raise gr.Error(f"Unsupported robot type {robot_type!r}")

    # LIBERO interface feeds the instruction capitalized.
    instruction = instruction.capitalize()

    raw_state = np.array(
        [state_ee_x, state_ee_y, state_ee_z,
         state_aa1, state_aa2, state_aa3,
         state_finger_l, state_finger_r],
        dtype=np.float32,
    )
    padded = np.zeros((STATE_LENGTH, STATE_DIM), dtype=np.float32)
    padded[-1, :8] = raw_state
    state = torch.from_numpy(padded)[None]

    chunk_len = int(chunk_length)
    if chunk_len < 1 or chunk_len > 20:
        chunk_len = 10
    # Manual mask: dims 0-6 (translation/rotation/gripper) — the reference server
    # overrides the processor's stats-derived mask, which is all zeros for LIBERO.
    mask = torch.zeros((1, chunk_len, ACTION_DIM), dtype=torch.float32)
    mask[:, :, :7] = 1.0

    model_inputs = processor.apply_chat_template(
        _messages(base_view, wrist_view, instruction),
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
        state=state,
        robot_type=robot_type,
    )
    model_inputs = dict(model_inputs)
    model_inputs["state"] = state
    model_inputs["action_mask"] = mask
    model_inputs["seed"] = int(seed)
    model_inputs["num_steps"] = NUM_FLOW_STEPS
    model_inputs = _to_cuda_dict(model_inputs)

    with torch.no_grad():
        output = model(**model_inputs)
    actions = getattr(output, "actions", None)
    if actions is None or actions.ndim != 3 or actions.shape[0] != 1:
        raise gr.Error(f"Unexpected model output shape {tuple(actions.shape) if actions is not None else None}")

    mean = torch.as_tensor(ACTION_MEAN, device=actions.device).view(1, 1, 7)
    std = torch.as_tensor(ACTION_STD, device=actions.device).view(1, 1, 7)
    decoded = actions[:, :, :7].float() * std + mean
    chunk = decoded[0].cpu().numpy().astype(np.float32)
    if not np.isfinite(chunk).all():
        raise gr.Error("Decoded action chunk contains NaN or inf.")

    fig = build_figure(chunk)
    table = chunk_table(chunk)
    csv_text = format_csv(chunk)
    dt = time.time() - t0
    dt_gpu = dt  # measured end-to-end inside the GPU function
    status = (
        f"Predicted {chunk.shape[0]} actions in {dt_gpu:.1f}s · suite={robot_type} · seed={int(seed)} "
        f"· gripper ends {chunk[-1, 6]:+.2f}"
    )
    return fig, table, csv_text, status


EXAMPLES = [
    ["ex_base.png", "ex_wrist.png",
     "Turn on the stove and put the moka pot on it.",
     "libero_10", -0.19947, -0.00631, 1.18217, 3.13143, 0.00487, -0.00294, 0.03622, -0.03623, 42, 10],
    ["ex_bowl_base.png", "ex_bowl_wrist.png",
     "Put the black bowl in the bottom drawer of the cabinet and close it.",
     "libero_10", -0.2201, -0.0014, 1.1623, 3.1392, -0.0067, -0.1071, 0.0362, -0.0362, 42, 10],
]

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

with gr.Blocks(title="ERVLA — robot action prediction") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # ERVLA: Embodied Chain-of-Thought for Generalizable Robot Manipulation

            Predict a **robot action chunk** from two camera views + a task instruction, with the
            LIBERO-Plus-tuned ERVLA checkpoint (`ERVLA/liberoplus` — Qwen3-VL-4B backbone + a DiT
            flow-matching action head; 86.9% success on LIBERO-Plus). Give the robot's eye views,
            tell it what to do, and read the predicted end-effector deltas (Δxyz, Δrotation,
            gripper). No simulator runs here — the action chunk is the output.

            [[paper](https://huggingface.co/papers/2606.03784)] ·
            [code](https://github.com/Oscar-Yang11/ERVLA)] ·
            [model](https://huggingface.co/ERVLA/liberoplus)
            """
        )
        with gr.Row():
            with gr.Column(scale=1):
                base_view = gr.Image(type="pil", label="Base view (LIBERO agentview)")
                wrist_view = gr.Image(type="pil", label="Left-wrist view")
                instruction = gr.Textbox(
                    label="Instruction",
                    placeholder="e.g. Turn on the stove and put the moka pot on it.",
                    lines=1,
                )
                run_btn = gr.Button("Predict actions", variant="primary")
            with gr.Column(scale=1):
                plot = gr.Plot(label="Predicted Δposition per step")
                table_out = gr.Dataframe(
                    headers=["step", "Δx", "Δy", "Δz", "Δroll", "Δpitch", "Δyaw", "gripper"],
                    label="Full predicted action chunk",
                    datatype=["number"] * 8,
                )
                status_out = gr.Textbox(label="Status", interactive=False, buttons=["copy"])
                csv_out = gr.Textbox(label="Chunk as CSV", lines=8, interactive=False,
                                     buttons=["copy"])
        with gr.Accordion("Advanced", open=False):
            robot_type = gr.Dropdown(
                choices=LIBERO_ROBOT_TYPES, value="libero_10", label="LIBERO suite (robot type)"
            )
            gr.Markdown(
                "End-effector state (LIBERO proprioception, metric units). "
                "Defaults come from the first example scene."
            )
            state_ee_x = gr.Number(label="EE position x (m)", value=-0.19947, precision=5)
            state_ee_y = gr.Number(label="EE position y (m)", value=-0.00631, precision=5)
            state_ee_z = gr.Number(label="EE position z (m)", value=1.18217, precision=5)
            state_aa1 = gr.Number(label="Orientation axis-angle 1 (rad)", value=3.13143, precision=5)
            state_aa2 = gr.Number(label="Orientation axis-angle 2 (rad)", value=0.00487, precision=5)
            state_aa3 = gr.Number(label="Orientation axis-angle 3 (rad)", value=-0.00294, precision=5)
            state_finger_l = gr.Number(label="Gripper finger left (rad)", value=0.03622, precision=5)
            state_finger_r = gr.Number(label="Gripper finger right (rad)", value=-0.03623, precision=5)
            seed = gr.Number(label="Seed", value=42, precision=0)
            chunk_length = gr.Number(label="Action chunk length", value=10, precision=0)

        all_inputs = [
            base_view, wrist_view, instruction, robot_type,
            state_ee_x, state_ee_y, state_ee_z,
            state_aa1, state_aa2, state_aa3,
            state_finger_l, state_finger_r,
            seed, chunk_length,
        ]
        run_btn.click(
            predict,
            inputs=all_inputs,
            outputs=[plot, table_out, csv_out, status_out],
            api_name="predict",
        )

        gr.Examples(
            examples=EXAMPLES,
            inputs=all_inputs,
            outputs=[plot, table_out, csv_out, status_out],
            fn=predict,
            cache_examples=True,
            cache_mode="lazy",
        )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)