# -- coding: utf-8 --
import numpy as np
from scipy.spatial.transform import Rotation as rot

import cv2
import numpy as np
import math
 
# 旋转矩阵转旋转向量（Rxyz->Rvector，按特定轴旋转）
def rotvector2rot(rotvector):
    Rm = cv2.Rodrigues(rotvector)[0]
    return Rm

# 旋转矩阵转欧拉角
def isRotationMatrix(R):
    Rt = np.transpose(R)
    shouldBeIdentity = np.dot(Rt, R)
    I = np.identity(3, dtype=R.dtype)
    n = np.linalg.norm(I - shouldBeIdentity)
    return n < 1e-6
def rot2euler(R):
    # assert (isRotationMatrix(R))
    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]) * 180 / np.pi
        y = math.atan2(-R[2, 0], sy) * 180 / np.pi
        z = math.atan2(R[1, 0], R[0, 0]) * 180 / np.pi
    else:
        x = math.atan2(-R[1, 2], R[1, 1]) * 180 / np.pi
        y = math.atan2(-R[2, 0], sy) * 180 / np.pi
        z = 0
    return np.array([x, y, z])
def rot2eulerpi(R):
    # assert (isRotationMatrix(R))
    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
    return np.array([x, y, z])
def rot_to_ypr_zyx(R):
    """
    将旋转矩阵R转换为Yaw-Pitch-Roll角度（Z-Y-X顺序）
    """
    sin_p = 2 * (R[2, 0] ** 2 - R[0, 0] ** 2 + R[1, 0] ** 2) ** 0.5
    if sin_p < 1e-6:
        sin_p = 1e-6
    if R[2, 0] > 0:
        sin_p = -sin_p
    
    roll = np.arctan2(R[2, 0], R[2, 2])
    pitch = np.arcsin(-R[2, 1])
    yaw = np.arctan2(R[1, 0], R[0, 0])
    
    return yaw, pitch, roll
def rot_to_ypr_xyz(R):
    # 计算R的trace值
    tr = np.trace(R)
    # 计算sqrt(R.T.R) - 1的模长，即平方根trace值
    sqrt_tr_m1 = np.sqrt(0.25 * ((tr - 3) ** 2) + (R[2, 2] - 1) ** 2)
 
    # 根据平方根trace值判断Yaw-Pitch-Roll角的计算方法
    yaw = np.arctan2(R[2, 1], R[2, 2])  # 使用arctan2计算Yaw角
    pitch = np.arcsin(-R[2, 0])  # 使用arcsin计算Pitch角
    roll = np.arctan2(R[1, 0], R[0, 0])  # 使用arctan2计算Roll角
 
    return yaw, pitch, roll

# 四元数转欧拉角 
def quaternion2euler(quaternion):
    r = rot.from_quat(quaternion)
    euler = r.as_euler('xyz', degrees=True)
    return euler

# 四元数转旋转矩阵
def quaternion2rot(quaternion):
    r = rot.from_quat(quaternion)
    rot = r.as_matrix()
    return rot

# 欧拉角转四元数
def euler2quaternion(euler):
    r = rot.from_euler('xyz', euler, degrees=True)
    quaternion = r.as_quat()
    return quaternion

# 欧拉角转旋转矩阵
def euler2rot(euler):
    r = rot.from_euler('xyz', euler, degrees=True)
    rotation_matrix = r.as_matrix()
    return rotation_matrix