# SPDX-FileCopyrightText: Copyright (c) 2022-2025, The Isaac Lab Project Developers.
# SPDX-License-Identifier: BSD-3-Clause

"""Coordinate frame utilities for XR teleoperation."""

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


def quat_wxyz_to_rotation(quat_wxyz: np.ndarray) -> R:
    """Convert Isaac-style wxyz quaternion to scipy Rotation."""
    return R.from_quat([quat_wxyz[1], quat_wxyz[2], quat_wxyz[3], quat_wxyz[0]])


def rotation_to_quat_wxyz(rot: R) -> np.ndarray:
    """Convert scipy Rotation to Isaac-style wxyz quaternion."""
    q = rot.as_quat()
    return np.array([q[3], q[0], q[1], q[2]])


def world_pose_to_root_frame(
    pos_w: np.ndarray,
    quat_wxyz: np.ndarray,
    root_pos_w: np.ndarray,
    root_quat_wxyz: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """Express a world-frame pose in the robot root frame."""
    root_rot = quat_wxyz_to_rotation(root_quat_wxyz)
    pose_rot = quat_wxyz_to_rotation(quat_wxyz)
    pos_root = root_rot.inv().apply(pos_w - root_pos_w)
    quat_root = rotation_to_quat_wxyz(root_rot.inv() * pose_rot)
    return pos_root, quat_root


def convert_action_world_to_root(action_tensor: torch.Tensor, robot) -> torch.Tensor:
    """Convert absolute IK action from world frame to robot root frame.

    Args:
        action_tensor: 8D tensor [x, y, z, qw, qx, qy, qz, gripper] in world frame.
        robot: Isaac Lab Articulation object (env.unwrapped.scene["robot"]).

    Returns:
        8D tensor with pos/quat converted to robot root frame.
    """
    root_pos_w = robot.data.root_pos_w[0].detach().cpu().numpy()
    root_quat_w = robot.data.root_quat_w[0].detach().cpu().numpy()
    pos_root, quat_root = world_pose_to_root_frame(
        action_tensor[:3].numpy(), action_tensor[3:7].numpy(),
        root_pos_w, root_quat_w,
    )
    out = action_tensor.clone()
    out[:3] = torch.tensor(pos_root, dtype=torch.float32)
    out[3:7] = torch.tensor(quat_root, dtype=torch.float32)
    return out