"""
Relative action representation — ported from Cosmos-Predict2.5.

Why relative actions:
  In robotics, actions can be represented as:
  1. Absolute positions (joint angles, end-effector pose)
  2. Relative displacements (delta from previous state)

  Cosmos-Predict2.5 uses relative end-effector actions:
    action = [relative_xyz(3), relative_rotation(3), gripper_state(1)]
  where relative transforms are computed in the gripper coordinate frame:
    rel_xyz = prev_rotm.T @ (curr_xyz - prev_xyz)
    rel_rot = rotm2euler(prev_rotm.T @ curr_rotm)

  This representation is better for cross-embodiment transfer because:
  - Same delta action = same physical movement regardless of robot geometry
  - Scale-invariant with proper normalization
  - Naturally handles different workspace origins
  - Matches how Cosmos action-conditioned models are trained

Adapted from: cosmos_predict2/_src/predict2/action/datasets/dataset_utils.py
License: Apache-2.0 (NVIDIA CORPORATION)
"""

from __future__ import annotations

import math
from typing import Optional, Tuple

import numpy as np


def euler2rotm(euler_angles: np.ndarray) -> np.ndarray:
    """Euler angle (ZYX convention) to 3×3 rotation matrix.

    Args:
        euler_angles: (3,) array [roll, pitch, yaw] in radians

    Returns:
        (3, 3) rotation matrix
    """
    alpha, beta, gamma = euler_angles[0], euler_angles[1], euler_angles[2]

    # Roll (X-axis)
    Rx = np.array([
        [1, 0, 0],
        [0, np.cos(alpha), -np.sin(alpha)],
        [0, np.sin(alpha), np.cos(alpha)],
    ])
    # Pitch (Y-axis)
    Ry = np.array([
        [np.cos(beta), 0, np.sin(beta)],
        [0, 1, 0],
        [-np.sin(beta), 0, np.cos(beta)],
    ])
    # Yaw (Z-axis)
    Rz = np.array([
        [np.cos(gamma), -np.sin(gamma), 0],
        [np.sin(gamma), np.cos(gamma), 0],
        [0, 0, 1],
    ])

    return Rz @ Ry @ Rx


def rotm2euler(R: np.ndarray) -> np.ndarray:
    """3×3 rotation matrix to Euler angles (ZYX convention).

    Args:
        R: (3, 3) rotation matrix

    Returns:
        (3,) array [roll, pitch, yaw] in radians, range (-π, π]
    """
    sy = math.sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0])
    singular = sy < 1e-6

    if not singular:
        x = math.atan2(R[2, 1], R[2, 2])
        y = math.atan2(-R[2, 0], sy)
        z = math.atan2(R[1, 0], R[0, 0])
    else:
        x = math.atan2(-R[1, 2], R[1, 1])
        y = math.atan2(-R[2, 0], sy)
        z = 0

    # Normalize to (-π, π]
    for val in [x, y, z]:
        while val > np.pi:
            val -= 2 * np.pi
        while val <= -np.pi:
            val += 2 * np.pi

    return np.array([x, y, z])


def rotm2quat(R: np.ndarray) -> np.ndarray:
    """3×3 rotation matrix to quaternion (w, x, y, z).

    Args:
        R: (3, 3) rotation matrix

    Returns:
        (4,) quaternion [w, x, y, z]
    """
    R = np.array(R, dtype=float)
    trace = np.trace(R)

    if trace > 0:
        s = 0.5 / np.sqrt(trace + 1.0)
        w = 0.25 / s
        x = (R[2, 1] - R[1, 2]) * s
        y = (R[0, 2] - R[2, 0]) * s
        z = (R[1, 0] - R[0, 1]) * s
    elif R[0, 0] > R[1, 1] and R[0, 0] > R[2, 2]:
        s = 2.0 * np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2])
        w = (R[2, 1] - R[1, 2]) / s
        x = 0.25 * s
        y = (R[0, 1] + R[1, 0]) / s
        z = (R[0, 2] + R[2, 0]) / s
    elif R[1, 1] > R[2, 2]:
        s = 2.0 * np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2])
        w = (R[0, 2] - R[2, 0]) / s
        x = (R[0, 1] + R[1, 0]) / s
        y = 0.25 * s
        z = (R[1, 2] + R[2, 1]) / s
    else:
        s = 2.0 * np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1])
        w = (R[1, 0] - R[0, 1]) / s
        x = (R[0, 2] + R[2, 0]) / s
        y = (R[1, 2] + R[2, 1]) / s
        z = 0.25 * s

    return np.array([w, x, y, z])


def states_to_relative_actions(
    arm_states: np.ndarray,
    gripper_states: np.ndarray,
    action_scaler: float = 20.0,
    gripper_scale: float = 1.0,
    use_quat: bool = False,
) -> np.ndarray:
    """Convert end-effector state sequence to relative actions.

    This is the core action representation from Cosmos-Predict2.5.
    Each action is the relative displacement between consecutive states,
    computed in the gripper's local coordinate frame.

    Args:
        arm_states: (T, 6) end-effector poses [x, y, z, roll, pitch, yaw]
        gripper_states: (T,) gripper widths [0=open, 1=closed]
        action_scaler: Scale factor for xyz/rotation deltas (Cosmos default: 20.0)
        gripper_scale: Scale factor for gripper state
        use_quat: If True, use quaternion (4D) for rotation instead of euler (3D)

    Returns:
        (T-1, 7) or (T-1, 8) relative actions:
            [relative_xyz(3), relative_rotation(3 or 4), gripper_state(1)]
    """
    T = len(arm_states)
    act_dim = 8 if use_quat else 7
    actions = np.zeros((T - 1, act_dim), dtype=np.float32)

    for k in range(1, T):
        prev_xyz = arm_states[k - 1, :3]
        prev_rpy = arm_states[k - 1, 3:6]
        prev_rotm = euler2rotm(prev_rpy)

        curr_xyz = arm_states[k, :3]
        curr_rpy = arm_states[k, 3:6]
        curr_rotm = euler2rotm(curr_rpy)

        # Relative position in gripper frame
        rel_xyz = prev_rotm.T @ (curr_xyz - prev_xyz)
        # Relative rotation
        rel_rotm = prev_rotm.T @ curr_rotm

        if use_quat:
            rel_rot = rotm2quat(rel_rotm)
            actions[k - 1, :3] = rel_xyz
            actions[k - 1, 3:7] = rel_rot
            actions[k - 1, 7] = gripper_states[k]
        else:
            rel_rot = rotm2euler(rel_rotm)
            actions[k - 1, :3] = rel_xyz
            actions[k - 1, 3:6] = rel_rot
            actions[k - 1, 6] = gripper_states[k]

    # Scale
    scale = np.ones(act_dim, dtype=np.float32)
    scale[:6] = action_scaler
    scale[-1] = gripper_scale
    actions *= scale

    return actions


def get_action_sequence_from_states(
    data: dict,
    fps_downsample_ratio: int = 1,
    state_key: str = "state",
    gripper_key: str = "continuous_gripper_state",
    action_scaler: float = 20.0,
    gripper_scale: float = 1.0,
    use_quat: bool = False,
) -> np.ndarray:
    """Extract relative action sequence from a Cosmos-style data dict.

    Compatible with Bridge dataset format and Cosmos action-conditioned inference.

    Args:
        data: Dict with state_key and gripper_key entries
        fps_downsample_ratio: Temporal downsampling ratio
        state_key: Key for arm states in data dict
        gripper_key: Key for gripper states in data dict
        action_scaler: Scale factor for relative actions
        gripper_scale: Scale factor for gripper
        use_quat: Use quaternion rotation representation

    Returns:
        (T-1, 7 or 8) relative action sequence
    """
    arm_states = np.array(data[state_key])[::fps_downsample_ratio]
    gripper_states = np.array(data[gripper_key])[::fps_downsample_ratio]

    return states_to_relative_actions(
        arm_states=arm_states,
        gripper_states=gripper_states,
        action_scaler=action_scaler,
        gripper_scale=gripper_scale,
        use_quat=use_quat,
    )
