# Copyright (c) 2022-2026, The Mindbot Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
#
# Adapted from XRoboToolkit (https://github.com/NVIDIA-AI-IOT/XRoboToolkit-Teleop-Sample-Python)
# Original geometry utilities by Zhigen Zhao.

"""Geometry utilities for XR-to-World coordinate transformations.

This module provides the core coordinate transformation matrices and functions
needed to convert PICO 4 Ultra (OpenXR) data into Isaac Sim world coordinates.

Coordinate Systems:
    - OpenXR (PICO 4): +X Right, +Y Up, -Z Forward
    - Isaac Sim (World): +X Forward, +Y Left, +Z Up
"""

import numpy as np
from scipy.spatial.transform import Rotation as R


# =====================================================================
# Core Coordinate Transform: OpenXR -> Isaac Sim World
# =====================================================================

R_HEADSET_TO_WORLD = np.array(
    [
        [0, 0, -1],   # XR -Z (forward)  -> World +X (forward)
        [-1, 0, 0],   # XR -X (left)     -> World +Y (left)
        [0, 1, 0],    # XR +Y (up)       -> World +Z (up)
    ],
    dtype=np.float64,
)
"""3x3 rotation matrix that maps OpenXR coordinates to Isaac Sim world coordinates."""


# =====================================================================
# Quaternion Utilities
# =====================================================================

def is_valid_quaternion(quat: np.ndarray, tol: float = 1e-6) -> bool:
    """Check if a quaternion is valid (unit norm).
    
    Args:
        quat: Quaternion array of length 4 (any convention).
        tol: Tolerance for norm deviation from 1.0.
    
    Returns:
        True if the quaternion is valid.
    """
    if not isinstance(quat, (list, tuple, np.ndarray)) or len(quat) != 4:
        return False
    if not np.all(np.isfinite(quat)):
        return False
    magnitude = np.sqrt(np.sum(np.square(quat)))
    return abs(magnitude - 1.0) <= tol


def quaternion_to_angle_axis(quat_wxyz: np.ndarray, eps: float = 1e-6) -> np.ndarray:
    """Convert a quaternion (w, x, y, z) to angle-axis representation.
    
    Args:
        quat_wxyz: Quaternion in (w, x, y, z) format.
        eps: Tolerance for small angles.
    
    Returns:
        Angle-axis vector [ax*angle, ay*angle, az*angle], shape (3,).
    """
    q = np.array(quat_wxyz, dtype=np.float64, copy=True)
    if q[0] < 0.0:
        q = -q

    w = q[0]
    vec_part = q[1:]
    angle = 2.0 * np.arccos(np.clip(w, -1.0, 1.0))

    if angle < eps:
        return np.zeros(3, dtype=np.float64)

    sin_half_angle = np.sin(angle / 2.0)
    if sin_half_angle < eps:
        return np.zeros(3, dtype=np.float64)

    axis = vec_part / sin_half_angle
    return axis * angle


def quat_diff_as_angle_axis(
    q1_wxyz: np.ndarray, q2_wxyz: np.ndarray, eps: float = 1e-6
) -> np.ndarray:
    """Calculate the rotation from q1 to q2 as an angle-axis vector.

    Computes DeltaQ such that q2 = DeltaQ * q1.

    Args:
        q1_wxyz: Source quaternion (w, x, y, z).
        q2_wxyz: Target quaternion (w, x, y, z).
        eps: Tolerance for small rotations.

    Returns:
        Angle-axis vector representing DeltaQ, shape (3,).
    """
    # Use scipy for robust quaternion math (scipy uses xyzw convention)
    r1 = R.from_quat([q1_wxyz[1], q1_wxyz[2], q1_wxyz[3], q1_wxyz[0]])
    r2 = R.from_quat([q2_wxyz[1], q2_wxyz[2], q2_wxyz[3], q2_wxyz[0]])
    delta = r2 * r1.inv()
    rotvec = delta.as_rotvec()
    if np.linalg.norm(rotvec) < eps:
        return np.zeros(3, dtype=np.float64)
    return rotvec


def quat_diff_as_rotvec_xyzw(
    q1_xyzw: np.ndarray, q2_xyzw: np.ndarray, eps: float = 1e-6
) -> np.ndarray:
    """Calculate the rotation from q1 to q2 as a rotation vector.

    Same as quat_diff_as_angle_axis but takes (x, y, z, w) quaternions
    directly (scipy convention), avoiding extra conversions.

    Args:
        q1_xyzw: Source quaternion (x, y, z, w).
        q2_xyzw: Target quaternion (x, y, z, w).
        eps: Tolerance for small rotations.

    Returns:
        Rotation vector (angle * axis), shape (3,).
    """
    r1 = R.from_quat(q1_xyzw)
    r2 = R.from_quat(q2_xyzw)
    delta = r2 * r1.inv()
    rotvec = delta.as_rotvec()
    if np.linalg.norm(rotvec) < eps:
        return np.zeros(3, dtype=np.float64)
    return rotvec


def apply_delta_pose(
    source_pos: np.ndarray,
    source_quat_xyzw: np.ndarray,
    delta_pos: np.ndarray,
    delta_rot: np.ndarray,
    eps: float = 1.0e-6,
) -> tuple:
    """Apply a delta pose to a source pose.

    Args:
        source_pos: Position of source frame, shape (3,).
        source_quat_xyzw: Quaternion of source frame (x, y, z, w), shape (4,).
        delta_pos: Position displacement, shape (3,).
        delta_rot: Rotation displacement in angle-axis format, shape (3,).
        eps: Tolerance to consider rotation as zero.

    Returns:
        Tuple of (target_pos, target_quat_xyzw), both np.ndarray.
    """
    target_pos = source_pos + delta_pos

    angle = np.linalg.norm(delta_rot)
    if angle > eps:
        delta_r = R.from_rotvec(delta_rot)
    else:
        delta_r = R.identity()

    source_r = R.from_quat(source_quat_xyzw)
    target_r = delta_r * source_r
    target_quat_xyzw = target_r.as_quat()

    return target_pos, target_quat_xyzw


# =====================================================================
# XR Pose Transform
# =====================================================================

def transform_xr_pose(
    pos_xr: np.ndarray, quat_xyzw_xr: np.ndarray
) -> tuple:
    """Transform an XR device pose from OpenXR frame to Isaac Sim world frame.

    Uses the R_HEADSET_TO_WORLD rotation matrix to transform both position
    and orientation, following the same approach as XRoboToolkit.

    Args:
        pos_xr: Position in OpenXR frame, shape (3,).
        quat_xyzw_xr: Quaternion in (x, y, z, w) format from XR SDK, shape (4,).

    Returns:
        Tuple of (pos_world, quat_xyzw_world):
            - pos_world: Position in world frame, shape (3,).
            - quat_xyzw_world: Quaternion in (x, y, z, w) for scipy, shape (4,).
    """
    # Transform position
    pos_world = R_HEADSET_TO_WORLD @ pos_xr

    # Transform orientation: R_world = R_transform * R_xr * R_transform^T
    r_xr = R.from_quat(quat_xyzw_xr)  # scipy uses (x, y, z, w)
    r_transform = R.from_matrix(R_HEADSET_TO_WORLD)
    r_world = r_transform * r_xr * r_transform.inv()
    quat_xyzw_world = r_world.as_quat()

    return pos_world, quat_xyzw_world
