import os

# Set expandable segments to avoid allocator fragmentation
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / any CUDA-touching import
import sys
import json
import time
import tempfile
import logging
from pathlib import Path

import numpy as np
import torch
from PIL import Image

# Add the Space repo root to sys.path so `src` and `deployment` modules are importable
SPACE_ROOT = Path(__file__).resolve().parent
if str(SPACE_ROOT) not in sys.path:
    sys.path.insert(0, str(SPACE_ROOT))

# Also add the LabVLA repo root for deployment imports
LABVLA_ROOT = str(SPACE_ROOT)
os.environ.setdefault("LABVLA_ROOT", LABVLA_ROOT)

import gradio as gr

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
logger = logging.getLogger(__name__)

MODEL_ID = "zjunlp/LabVLA-5B-Base"

# ---- Download model checkpoint ----
from huggingface_hub import snapshot_download

logger.info(f"Downloading model from {MODEL_ID}...")
_model_dir = snapshot_download(
    repo_id=MODEL_ID,
    repo_type="model",
    local_dir=str(SPACE_ROOT / "model_cache"),
)
# The checkpoint files are at the root of the repo
PRETRAINED_PATH = str(SPACE_ROOT / "model_cache")
logger.info(f"Model downloaded to {PRETRAINED_PATH}")

# ---- Load the LabVLA policy ----
from labsim_transforms import parse_image_to_uint8_hwc
from src.policies.LabVLA.configuration_labvla import LabVLAConfig
from src.policies.LabVLA.modeling_labvla import LabVLAPolicy

# Load config from the checkpoint
config_path = Path(PRETRAINED_PATH) / "config.json"
with open(config_path) as f:
    saved_config = json.load(f)

# Build LabVLAConfig from saved config, with deploy overrides
allowed_fields = set(LabVLAConfig.__dataclass_fields__.keys())
saved_fwd = {k: v for k, v in saved_config.items() if k in allowed_fields}

# Normalize tuple fields
if "image_resolution" in saved_fwd and isinstance(saved_fwd["image_resolution"], list):
    saved_fwd["image_resolution"] = tuple(saved_fwd["image_resolution"])
if "optimizer_betas" in saved_fwd and isinstance(saved_fwd["optimizer_betas"], list):
    saved_fwd["optimizer_betas"] = tuple(saved_fwd["optimizer_betas"])

# Drop str-fallback fields that can't be properly deserialized
for field_name in ("input_features", "output_features"):
    if field_name in saved_fwd and isinstance(saved_fwd[field_name], str):
        saved_fwd.pop(field_name)

# Deploy overrides: no GC, no compile, use the checkpoint's bundled VLM
saved_fwd["vlm_pretrained_path"] = PRETRAINED_PATH
saved_fwd["freeze_vision_encoder"] = False
saved_fwd["gradient_checkpointing"] = False
saved_fwd["compile_model"] = False

# Use SDPA instead of flash_attention_2 for ZeroGPU compatibility
saved_fwd["attn_implementation"] = "sdpa"

# Load to CPU first — ZeroGPU's safetensors load_file(device="cuda") bypasses
# the spaces hijack and fails with "No CUDA GPUs". Load on CPU, then .to("cuda")
# which IS intercepted by the hijack.
saved_fwd["device"] = "cpu"

config = LabVLAConfig(**saved_fwd)

# Load the policy
logger.info("Loading LabVLA policy...")
policy = LabVLAPolicy.from_pretrained(
    pretrained_name_or_path=PRETRAINED_PATH,
    config=config,
    strict=True,
)
policy.to("cuda")
policy.eval()
logger.info("LabVLA policy loaded successfully!")

# Load VLM processor for image/token processing
from transformers import Qwen3VLProcessor
vlm_processor = Qwen3VLProcessor.from_pretrained(PRETRAINED_PATH)
vision_start_token_id = vlm_processor.vision_start_token_id
vision_end_token_id = vlm_processor.vision_end_token_id
image_token_id = vlm_processor.image_token_id

# Image processing constants
SPATIAL_MERGE_SIZE = 2
h, w = config.image_resolution
image_size = (h, w)
dummy_img = torch.zeros(h, w, 3)
dummy_out = vlm_processor.image_processor([dummy_img], do_rescale=False, return_tensors="pt")
_fixed_grid_thw = dummy_out["image_grid_thw"]
num_patches = int(_fixed_grid_thw[0, 0] * _fixed_grid_thw[0, 1] * _fixed_grid_thw[0, 2])
num_image_tokens = num_patches // (SPATIAL_MERGE_SIZE ** 2)

# Load schema
schema_path = SPACE_ROOT / "labvla_schema.json"
with open(schema_path) as f:
    schema = json.load(f)

# Schema-derived dimensions
state_dim = sum(schema["state_dims"])
action_dim = sum(schema["action_dims"])
delta_mask = np.array(schema["delta_mask"], dtype=bool)
max_state_dim = config.max_state_dim
max_action_dim = config.max_action_dim
chunk_size = config.chunk_size

logger.info(
    f"Schema loaded: state_dim={state_dim}, action_dim={action_dim}, "
    f"chunk_size={chunk_size}, image_size={image_size}, "
    f"num_image_tokens={num_image_tokens}"
)


def resize_with_pad(image_tensor, target_h, target_w):
    """Resize keeping aspect ratio + zero-pad to target size."""
    from src.transforms.utils import resize_with_pad as _train_resize_with_pad
    return _train_resize_with_pad(image_tensor, target_h, target_w, "bilinear")


def process_image(image_hwc):
    """Convert HWC uint8 image to Qwen3-VL pixel_values."""
    target_h, target_w = image_size
    img_tensor = torch.from_numpy(image_hwc.copy()).float() / 255.0
    img_tensor = img_tensor.permute(2, 0, 1)
    img_tensor = resize_with_pad(img_tensor, target_h, target_w)
    img_inputs = vlm_processor.image_processor([img_tensor], do_rescale=False, return_tensors="pt")
    pixel_values = img_inputs["pixel_values"]
    image_grid_thw = img_inputs["image_grid_thw"]
    return pixel_values, image_grid_thw


def build_batch(images, prompt, state):
    """Build input batch from camera images, prompt, and state."""
    all_pixel_values = []
    all_image_grid_thw = []
    input_ids = []
    attention_mask = []

    # Pre-process first valid image for placeholder
    first_valid_idx = next((i for i, img in enumerate(images) if img is not None), None)
    if first_valid_idx is None:
        raise ValueError("No valid images provided")

    placeholder_pv, placeholder_grid = process_image(images[first_valid_idx])

    for i, img in enumerate(images):
        if img is not None:
            pv, grid = process_image(img)
            all_pixel_values.append(pv)
            all_image_grid_thw.append(grid)
            vision_ids = (
                [vision_start_token_id]
                + [image_token_id] * num_image_tokens
                + [vision_end_token_id]
            )
            input_ids += vision_ids
            attention_mask += [1] * len(vision_ids)
        else:
            all_pixel_values.append(placeholder_pv)
            all_image_grid_thw.append(placeholder_grid)
            vision_ids = (
                [vision_start_token_id]
                + [image_token_id] * num_image_tokens
                + [vision_end_token_id]
            )
            input_ids += vision_ids
            attention_mask += [0] * len(vision_ids)

    # Tokenize language instruction
    lang_inputs = vlm_processor.tokenizer(
        prompt,
        max_length=config.tokenizer_max_length,
        padding="max_length",
        truncation=True,
    )
    input_ids += lang_inputs.input_ids
    attention_mask += lang_inputs.attention_mask

    pixel_values = torch.cat(all_pixel_values, dim=0)
    image_grid_thw = torch.cat(all_image_grid_thw, dim=0)

    # Pad state to max_state_dim
    state_padded = np.zeros(max_state_dim, dtype=np.float32)
    state_padded[:min(len(state), max_state_dim)] = state[:max_state_dim]
    state_tensor = torch.from_numpy(state_padded).to(dtype=torch.bfloat16)

    batch = {
        "observation.pixel_values": pixel_values.unsqueeze(0).to("cuda"),
        "observation.image_grid_thw": image_grid_thw.to("cuda"),
        "observation.input_ids": torch.tensor(input_ids, dtype=torch.long).unsqueeze(0).to("cuda"),
        "observation.attention_mask": torch.tensor(attention_mask, dtype=torch.long).unsqueeze(0).to("cuda"),
        "observation.state": state_tensor.unsqueeze(0).to("cuda"),
    }
    return batch


# Default robot state used whenever a caller doesn't provide one (e.g. the
# gr.Examples rows below only populate camera_1 + instruction). This mirrors
# the default values of the Robot State sliders in the UI, so results are
# consistent regardless of entry point (example click, manual button click,
# or a direct API call that omits the state args).
DEFAULT_STATE_J = 0.0
DEFAULT_GRIPPER = 0.04


@spaces.GPU(duration=120)
def predict_actions(
    camera_1: Image.Image,
    instruction: str,
    state_j1: float = DEFAULT_STATE_J, state_j2: float = DEFAULT_STATE_J, state_j3: float = DEFAULT_STATE_J,
    state_j4: float = DEFAULT_STATE_J, state_j5: float = DEFAULT_STATE_J, state_j6: float = DEFAULT_STATE_J,
    state_j7: float = DEFAULT_STATE_J, gripper: float = DEFAULT_GRIPPER,
):
    """Predict a robot action chunk from laboratory camera views and a language instruction.

    LabVLA is a Vision-Language-Action model that takes camera images, a natural
    language instruction, and the current robot state (7 joint angles + 1 gripper
    width) as input, and predicts a chunk of 50 future action steps.

    Args:
        camera_1: Camera view of the laboratory workspace.
        instruction: Natural language task instruction (e.g. "Pick up the beaker").
        state_j1..j7: Franka Panda 7-DOF arm joint angles (radians). Defaults to 0.0.
        gripper: Gripper width in meters (0.0 = closed, 0.04 = fully open). Defaults to 0.04.

    Returns:
        A matplotlib figure visualizing the predicted action trajectory,
        and a JSON dict with the raw action values.
    """
    if not instruction or not instruction.strip():
        return None, {"error": "Please provide a task instruction."}

    # Convert PIL image to numpy
    img1 = np.array(camera_1.convert("RGB"))
    img1 = parse_image_to_uint8_hwc(img1)

    # Use the same image for all 3 cameras (the model supports 3 camera views;
    # with only 1, slots 2/3 are masked out)
    images = [img1, None, None]

    # Guard against any caller (or gr.Examples cache) passing an empty/missing
    # value for a state component. NaN must never reach the model.
    def _clean(value, default):
        return default if value is None else value

    state_j1 = _clean(state_j1, DEFAULT_STATE_J)
    state_j2 = _clean(state_j2, DEFAULT_STATE_J)
    state_j3 = _clean(state_j3, DEFAULT_STATE_J)
    state_j4 = _clean(state_j4, DEFAULT_STATE_J)
    state_j5 = _clean(state_j5, DEFAULT_STATE_J)
    state_j6 = _clean(state_j6, DEFAULT_STATE_J)
    state_j7 = _clean(state_j7, DEFAULT_STATE_J)
    gripper = _clean(gripper, DEFAULT_GRIPPER)

    # Build state vector
    state = np.array([state_j1, state_j2, state_j3, state_j4,
                       state_j5, state_j6, state_j7, gripper], dtype=np.float32)

    start_time = time.perf_counter()

    # Build batch
    batch = build_batch(images, instruction, state)

    # Run inference
    with torch.no_grad():
        action_chunk = policy.predict_action_chunk(batch)

    if isinstance(action_chunk, torch.Tensor):
        actions = action_chunk.detach().float().cpu().numpy()
    else:
        actions = np.asarray(action_chunk)

    if actions.ndim == 3:
        actions = actions[0]

    # Truncate to action_dim
    actions = actions[:, :action_dim]

    # Delta -> absolute for arm dims
    arm_mask = delta_mask[:actions.shape[-1]]
    arm_idxs = np.where(arm_mask)[0]
    if arm_idxs.size > 0:
        n_arm = len(arm_mask)
        state_for_add = state[:n_arm]
        if state_for_add.shape[0] < n_arm:
            state_for_add = np.concatenate([
                state_for_add,
                np.zeros(n_arm - state_for_add.shape[0], dtype=state_for_add.dtype),
            ])
        state_delta = state_for_add[arm_mask]
        actions[:, arm_mask] = actions[:, arm_mask] + state_delta[np.newaxis, :]

    infer_time = time.perf_counter() - start_time

    # Create visualization
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    fig, axes = plt.subplots(2, 4, figsize=(16, 8))
    fig.suptitle(
        f"LabVLA Predicted Action Chunk ({chunk_size} steps, {infer_time:.2f}s)\n"
        f"Instruction: \"{instruction}\"",
        fontsize=12, fontweight="bold",
    )

    labels = ["J1", "J2", "J3", "J4", "J5", "J6", "J7", "Gripper"]
    for i in range(8):
        ax = axes[i // 4, i % 4]
        ax.plot(actions[:, i], linewidth=2, color="#4C72B0")
        ax.set_title(labels[i], fontsize=11, fontweight="bold")
        ax.set_xlabel("Step")
        ax.set_ylabel("Value" if i < 7 else "Width (m)")
        ax.grid(True, alpha=0.3)
        ax.axhline(y=state[i], color="r", linestyle="--", alpha=0.5, label="Current state")
        if i == 0:
            ax.legend(fontsize=8)

    plt.tight_layout()

    # Save figure
    fig_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
    fig.savefig(fig_path, dpi=150, bbox_inches="tight")
    plt.close(fig)

    # Prepare JSON output (rounded for readability; gr.JSON renders the dict natively)
    def _round_list(values, ndigits=4):
        return [round(float(v), ndigits) for v in values]

    actions_list = actions.tolist()
    result = {
        "instruction": instruction,
        "state_input": _round_list(state.tolist()),
        "action_chunk_shape": list(actions.shape),
        "num_steps": int(actions.shape[0]),
        "action_dim": int(actions.shape[1]),
        "inference_time_s": round(infer_time, 3),
        "first_action": _round_list(actions_list[0]) if len(actions_list) > 0 else None,
        "last_action": _round_list(actions_list[-1]) if len(actions_list) > 0 else None,
    }

    return fig_path, result


# ---- Gradio UI ----

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

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    gr.Markdown("""
    # 🧪 LabVLA: Vision-Language-Action Model for Scientific Laboratories

    LabVLA is the first VLA foundation model designed specifically for scientific laboratory environments.
    It combines a **Qwen3-VL-4B** vision-language backbone with a **DiT flow-matching action expert** to
    predict robotic action trajectories from laboratory camera views and language instructions.

    📄 [Paper](https://huggingface.co/papers/2606.13578) •
    💻 [GitHub](https://github.com/zjunlp/LabVLA) •
    🤗 [Model](https://huggingface.co/zjunlp/LabVLA-5B-Base)
    """)

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column(scale=1):
                camera_1 = gr.Image(
                    label="Camera View (Laboratory Workspace)",
                    type="pil",
                    height=280,
                )
                instruction = gr.Textbox(
                    label="Task Instruction",
                    placeholder="e.g. Pick up the beaker and pour it into the flask",
                    value="Pick up the beaker",
                    lines=2,
                )
                with gr.Accordion("Robot State (Franka Panda 7-DOF + Gripper)", open=False):
                    with gr.Row():
                        state_j1 = gr.Slider(-6.28, 6.28, value=0.0, step=0.01, label="Joint 1 (rad)")
                        state_j2 = gr.Slider(-6.28, 6.28, value=0.0, step=0.01, label="Joint 2 (rad)")
                    with gr.Row():
                        state_j3 = gr.Slider(-6.28, 6.28, value=0.0, step=0.01, label="Joint 3 (rad)")
                        state_j4 = gr.Slider(-6.28, 6.28, value=0.0, step=0.01, label="Joint 4 (rad)")
                    with gr.Row():
                        state_j5 = gr.Slider(-6.28, 6.28, value=0.0, step=0.01, label="Joint 5 (rad)")
                        state_j6 = gr.Slider(-6.28, 6.28, value=0.0, step=0.01, label="Joint 6 (rad)")
                    with gr.Row():
                        state_j7 = gr.Slider(-6.28, 6.28, value=0.0, step=0.01, label="Joint 7 (rad)")
                        gripper = gr.Slider(0.0, 0.04, value=0.04, step=0.001, label="Gripper Width (m)")

                run_btn = gr.Button("Predict Actions", variant="primary", size="lg")

            with gr.Column(scale=1):
                output_plot = gr.Image(label="Predicted Action Trajectory (50 steps)")
                output_json = gr.JSON(label="Action Details (JSON)")

    gr.Examples(
        examples=[
            ["examples/lab_scene_view.jpg", "Pick up the beaker and pour it into the flask"],
            ["examples/lab_scene_view.jpg", "Transfer the solution from the test tube to the beaker"],
            ["examples/lab_scene_alt.jpg", "Press the button on the instrument"],
        ],
        inputs=[camera_1, instruction],
        outputs=[output_plot, output_json],
        fn=predict_actions,
        cache_examples=True,
        cache_mode="lazy",
    )

    run_btn.click(
        fn=predict_actions,
        inputs=[
            camera_1, instruction,
            state_j1, state_j2, state_j3, state_j4,
            state_j5, state_j6, state_j7, gripper,
        ],
        outputs=[output_plot, output_json],
        api_name="predict",
    )

demo.launch(mcp_server=True)