# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Action-token helpers for Cosmos3 action generation.

These helpers cover action modes that use action tokens as an auxiliary output
stream: ``policy``, ``forward_dynamics``, and ``inverse_dynamics``. The pipeline
returns predicted actions through the diffusion output payload/metadata
envelope rather than as a video frame stream.
"""

from __future__ import annotations

from typing import Any

import numpy as np
import torch

from .utils import VIDEO_RES_SIZE_INFO

ACTION_MODE_POLICY = "policy"
ACTION_MODE_FORWARD_DYNAMICS = "forward_dynamics"
ACTION_MODE_INVERSE_DYNAMICS = "inverse_dynamics"
ACTION_MODES = {
    ACTION_MODE_POLICY,
    ACTION_MODE_FORWARD_DYNAMICS,
    ACTION_MODE_INVERSE_DYNAMICS,
}


EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = {
    "no_action": 0,
    "av": 1,
    "camera_pose": 2,
    "hand_pose": 3,
    "pusht": 4,
    "libero": 5,
    "umi": 6,
    "bridge_orig_lerobot": 7,
    "droid_lerobot": 8,
    "robomind-franka": 8,
    "galbot": 9,
    "robomind-franka-dual": 12,
    "robomind-ur": 13,
    "agibotworld": 15,
    "agibot_gear_gripper": 15,
    "agibot_gear_gripper_ext": 15,
    "fractal": 20,
}


def normalize_action_mode(mode: Any) -> str | None:
    if mode is None:
        return None
    normalized = str(mode).strip().lower()
    if not normalized:
        return None
    if normalized not in ACTION_MODES:
        raise ValueError(f"Unsupported Cosmos3 action_mode={mode!r}; expected one of {sorted(ACTION_MODES)}.")
    return normalized


def resolve_domain_id(
    *,
    domain_id: Any = None,
    domain_name: Any = None,
    require_explicit: bool = False,
) -> int:
    if domain_id is not None:
        resolved = int(domain_id)
        if resolved < 0:
            raise ValueError(f"Cosmos3 domain_id must be non-negative, got {resolved}.")
        return resolved

    if domain_name is None or str(domain_name).strip() == "":
        if require_explicit:
            raise ValueError(
                "Cosmos3 action generation requires extra_args['domain_id'] or non-empty extra_args['domain_name']."
            )
        return 0

    key = str(domain_name).strip().lower()
    if key not in EMBODIMENT_TO_DOMAIN_ID:
        raise ValueError(
            f"Unknown Cosmos3 action domain_name={domain_name!r}; "
            f"expected one of {sorted(EMBODIMENT_TO_DOMAIN_ID)} or pass domain_id directly."
        )
    return EMBODIMENT_TO_DOMAIN_ID[key]


def action_condition_indexes(mode: str, action_length: int) -> list[int]:
    mode = normalize_action_mode(mode)
    if mode == ACTION_MODE_FORWARD_DYNAMICS:
        return list(range(action_length))
    if mode in {ACTION_MODE_POLICY, ACTION_MODE_INVERSE_DYNAMICS}:
        return []
    raise AssertionError(f"Unexpected action mode: {mode!r}")


def vision_condition_indexes(mode: str, video_length: int, temporal_compression_factor: int) -> list[int]:
    mode = normalize_action_mode(mode)
    latent_frames = (video_length - 1) // temporal_compression_factor + 1
    if mode in {ACTION_MODE_POLICY, ACTION_MODE_FORWARD_DYNAMICS}:
        return [0]
    if mode == ACTION_MODE_INVERSE_DYNAMICS:
        return list(range(latent_frames))
    raise AssertionError(f"Unexpected action mode: {mode!r}")


def action_start_frame_offset(mode: str, action_length: int, video_length: int) -> int:
    del mode
    if action_length == video_length - 1:
        return 1
    if action_length == video_length:
        return 0
    raise ValueError(
        "Cosmos3 action_chunk_size must equal num_frames - 1 or num_frames; "
        f"got action_chunk_size={action_length}, num_frames={video_length}."
    )


def build_action_condition_mask(
    mode: str,
    action_length: int,
    *,
    device: torch.device,
    dtype: torch.dtype,
) -> torch.Tensor:
    mask = torch.zeros(1, action_length, 1, device=device, dtype=dtype)
    for idx in action_condition_indexes(mode, action_length):
        mask[:, idx, :] = 1.0
    return mask


def build_vision_condition_mask(
    mode: str,
    video_length: int,
    temporal_compression_factor: int,
    *,
    device: torch.device,
    dtype: torch.dtype,
) -> torch.Tensor:
    latent_frames = (video_length - 1) // temporal_compression_factor + 1
    mask = torch.zeros(1, 1, latent_frames, 1, 1, device=device, dtype=dtype)
    for idx in vision_condition_indexes(mode, video_length, temporal_compression_factor):
        mask[:, :, idx, :, :] = 1.0
    return mask


def pad_action_to_dim(action: torch.Tensor, action_dim: int) -> torch.Tensor:
    if action.shape[-1] > action_dim:
        raise ValueError(f"Cosmos3 action dimension {action.shape[-1]} exceeds model action_dim={action_dim}.")
    if action.shape[-1] == action_dim:
        return action
    padding = torch.zeros(*action.shape[:-1], action_dim - action.shape[-1], dtype=action.dtype, device=action.device)
    return torch.cat([action, padding], dim=-1)


def load_action_tensor(action: Any = None) -> torch.Tensor:
    if action is None:
        raise ValueError(
            "Cosmos3 forward_dynamics action mode requires extra_args['action'] "
            "(a JSON array / nested list / tensor of shape [T, D])."
        )
    if isinstance(action, torch.Tensor):
        tensor = action.detach().to(dtype=torch.float32)
    else:
        tensor = torch.as_tensor(np.asarray(action), dtype=torch.float32)
    if tensor.ndim == 3 and tensor.shape[0] == 1:
        tensor = tensor.squeeze(0)
    if tensor.ndim != 2:
        raise ValueError(f"Cosmos3 action must have shape [T, D], got {tuple(tensor.shape)}.")
    return tensor


def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int, int]:
    key = str(resolution)
    if key not in VIDEO_RES_SIZE_INFO:
        raise ValueError(
            f"Unknown Cosmos3 action resolution={resolution!r}; expected one of {sorted(VIDEO_RES_SIZE_INFO)}."
        )
    input_ratio = h / w
    best_size = None
    best_diff = float("inf")
    for cand_w, cand_h in VIDEO_RES_SIZE_INFO[key].values():
        diff = abs(input_ratio - cand_h / cand_w)
        if diff < best_diff:
            best_diff = diff
            best_size = (cand_w, cand_h)
    assert best_size is not None
    return best_size