import numpy as np
from sklearn.linear_model import RANSACRegressor
from scipy.spatial.transform import Rotation as R, Slerp


def interpolate_rotation_matrices(rot_matrix1, rot_matrix2, num_interpolations):
    # Convert the rotation matrices to rotation objects
    rot1 = R.from_matrix(rot_matrix1)
    rot2 = R.from_matrix(rot_matrix2)

    # Convert the rotation objects to quaternions
    quat1 = rot1.as_quat()
    quat2 = rot2.as_quat()

    # Define the times of the known rotations
    times = [0, 1]

    # Create the Slerp object
    slerp = Slerp(times, R.from_quat([quat1, quat2]))

    # Define the times of the interpolations
    interp_times = np.linspace(0, 1, num_interpolations)

    # Perform the interpolation
    interp_rots = slerp(interp_times)

    # Convert the interpolated rotations to matrices
    interp_matrices = interp_rots.as_matrix()

    return interp_matrices


def is_y_axis_up(pose_matrix):
    """
    判断物体在世界坐标系中的 y 轴是否朝上。

    参数:
    pose_matrix (numpy.ndarray): 4x4 齐次变换矩阵。

    返回:
    bool: True 如果 y 轴朝上，False 如果 y 轴朝下。
    """
    # 提取旋转矩阵的第二列
    y_axis_vector = pose_matrix[:3, 1]

    # 世界坐标系的 y 轴
    world_y_axis = np.array([0, 1, 0])

    # 计算点积
    dot_product = np.dot(y_axis_vector, world_y_axis)

    # 返回 True 如果朝上，否则返回 False
    return dot_product > 0


def is_local_axis_facing_world_axis(pose_matrix, local_axis='y', world_axis='z'):
    # 定义局部坐标系的轴索引
    local_axis_index = {'x': 0, 'y': 1, 'z': 2}

    # 定义世界坐标系的轴向量
    world_axes = {
        'x': np.array([1, 0, 0]),
        'y': np.array([0, 1, 0]),
        'z': np.array([0, 0, 1])
    }

    # 提取局部坐标系的指定轴
    local_axis_vector = pose_matrix[:3, local_axis_index[local_axis]]

    # 获取世界坐标系的指定轴向量
    world_axis_vector = world_axes[world_axis]

    # 计算点积
    dot_product = np.dot(local_axis_vector, world_axis_vector)

    # 返回 True 如果朝向指定世界轴，否则返回 False
    return dot_product > 0


def fix_gripper_rotation(source_affine, target_affine, rot_axis='z'):
    '''
      gripper是对称结构，绕Z轴旋转180度前后是等效的。选择离当前pose更近的target pose以避免不必要的旋转
    '''
    if rot_axis == 'z':
        R_180 = np.array([[-1, 0, 0],
                          [0, -1, 0],
                          [0, 0, 1]])
    elif rot_axis == 'y':
        R_180 = np.array([[-1, 0, 0],
                          [0, 1, 0],
                          [0, 0, -1]])
    elif rot_axis == 'x':
        R_180 = np.array([[1, 0, 0],
                          [0, -1, 0],
                          [0, 0, -1]])
    else:
        assert False, "Invalid rotation axis. Please choose from 'x', 'y', 'z'."

    # 提取source和target的旋转部分（3x3矩阵）
    source_rotation = source_affine[:3, :3]
    target_rotation = target_affine[:3, :3]
    # 将target_rotation绕其自身的Z轴转180度，得到target_rotation_2
    target_rotation_2 = np.dot(target_rotation, R_180)

    # 定义一个函数来计算两个旋转矩阵之间的距离
    def rotation_matrix_distance(R1, R2):
        # 使用奇异值分解来计算两个旋转矩阵之间的距离
        U, _, Vh = np.linalg.svd(np.dot(R1.T, R2))
        # 确保旋转矩阵的行列式为1，即旋转而不是反射
        det_check = np.linalg.det(U) * np.linalg.det(Vh)
        if det_check < 0:
            Vh = -Vh
        return np.arccos(np.trace(Vh) / 2)

    # 计算source_rotation与target_rotation之间的距离
    distance_target_rotation = rotation_matrix_distance(source_rotation, target_rotation)
    # 计算source_rotation与target_rotation_2之间的距离
    distance_target_rotation_2 = rotation_matrix_distance(source_rotation, target_rotation_2)
    # 比较两个距离，确定哪个旋转更接近source_rotation
    if distance_target_rotation < distance_target_rotation_2:
        return target_affine
    else:
        # 重新组合旋转矩阵target_rotation_2和原始的平移部分
        target_affine_2 = np.eye(4)
        target_affine_2[:3, :3] = target_rotation_2
        target_affine_2[:3, 3] = target_affine[:3, 3]  # 保留原始的平移部分
        return target_affine_2


def rotate_180_along_axis(pose, rot_axis='z'):
    '''
      gripper是对称结构，绕Z轴旋转180度前后是等效的。选择离当前pose更近的target pose以避免不必要的旋转
    '''
    if rot_axis == 'z':
        R_180 = np.array([[-1, 0, 0],
                          [0, -1, 0],
                          [0, 0, 1]])
    elif rot_axis == 'y':
        R_180 = np.array([[-1, 0, 0],
                          [0, 1, 0],
                          [0, 0, -1]])
    elif rot_axis == 'x':
        R_180 = np.array([[1, 0, 0],
                          [0, -1, 0],
                          [0, 0, -1]])
    else:
        assert False, "Invalid rotation axis. Please choose from 'x', 'y', 'z'."

    single_mode = pose.ndim == 2
    if single_mode:
        pose = pose[np.newaxis, :, :]
    R_180 = np.tile(R_180[np.newaxis, :, :], (pose.shape[0], 1, 1))
    pose[:, :3, :3] = pose[:, :3, :3] @ R_180

    if single_mode:
        pose = pose[0]

    return pose


def translate_along_axis(pose, shift, axis='z', use_local=True):
    '''
    根据指定的角度和旋转轴来旋转target_affine。
    参数:
    - target_affine: 4x4 仿射变换矩阵
    - angle_degrees: 旋转角度（以度为单位）
    - rot_axis: 旋转轴，'x'、'y' 或 'z'
    '''
    pose = pose.copy()
    translation = np.zeros(3)
    if axis == 'z':
        translation[2] = shift
    elif axis == 'y':
        translation[1] = shift
    elif axis == 'x':
        translation[0] = shift
    if len(pose.shape) == 3:
        for i in range(pose.shape[0]):
            if use_local:
                pose[i][:3, 3] += pose[i][:3, :3] @ translation
            else:
                pose[i][:3, 3] += translation
    else:
        if use_local:
            pose[:3, 3] += pose[:3, :3] @ translation
        else:
            pose[:3, 3] += translation

    return pose


def rotate_along_axis(target_affine, angle_degrees, rot_axis='z', use_local=True):
    '''
    根据指定的角度和旋转轴来旋转target_affine。
    参数:
    - target_affine: 4x4 仿射变换矩阵
    - angle_degrees: 旋转角度（以度为单位）
    - rot_axis: 旋转轴，'x'、'y' 或 'z'
    '''
    # 将角度转换为弧度
    angle_radians = np.deg2rad(angle_degrees)

    # 创建旋转对象
    if rot_axis == 'z':
        rotation_vector = np.array([0, 0, angle_radians])
    elif rot_axis == 'y':
        rotation_vector = np.array([0, angle_radians, 0])
    elif rot_axis == 'x':
        rotation_vector = np.array([angle_radians, 0, 0])
    else:
        raise ValueError("Invalid rotation axis. Please choose from 'x', 'y', 'z'.")

    # 生成旋转矩阵
    R_angle = R.from_rotvec(rotation_vector).as_matrix()

    # 提取旋转部分（3x3矩阵）
    target_rotation = target_affine[:3, :3]

    # 将 target_rotation 绕指定轴旋转指定角度，得到 target_rotation_2
    if use_local:
        target_rotation_2 = np.dot(target_rotation, R_angle)
    else:
        target_rotation_2 = np.dot(R_angle, target_rotation)

    # 重新组合旋转矩阵 target_rotation_2 和原始的平移部分
    target_affine_2 = np.eye(4)
    target_affine_2[:3, :3] = target_rotation_2
    target_affine_2[:3, 3] = target_affine[:3, 3]  # 保留原始的平移部分

    return target_affine_2


def rotation_matrix_to_quaternion(R):
    assert R.shape == (3, 3)

    # 计算四元数分量
    trace = np.trace(R)
    if trace > 0:
        S = np.sqrt(trace + 1.0) * 2  # S=4*qw
        qw = 0.25 * S
        qx = (R[2, 1] - R[1, 2]) / S
        qy = (R[0, 2] - R[2, 0]) / S
        qz = (R[1, 0] - R[0, 1]) / S
    elif (R[0, 0] > R[1, 1]) and (R[0, 0] > R[2, 2]):
        S = np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2]) * 2  # S=4*qx
        qw = (R[2, 1] - R[1, 2]) / S
        qx = 0.25 * S
        qy = (R[0, 1] + R[1, 0]) / S
        qz = (R[0, 2] + R[2, 0]) / S
    elif R[1, 1] > R[2, 2]:
        S = np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2]) * 2  # S=4*qy
        qw = (R[0, 2] - R[2, 0]) / S
        qx = (R[0, 1] + R[1, 0]) / S
        qy = 0.25 * S
        qz = (R[1, 2] + R[2, 1]) / S
    else:
        S = np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1]) * 2  # S=4*qz
        qw = (R[1, 0] - R[0, 1]) / S
        qx = (R[0, 2] + R[2, 0]) / S
        qy = (R[1, 2] + R[2, 1]) / S
        qz = 0.25 * S

    return np.array([qw, qx, qy, qz])


def quat_wxyz_to_rotation_matrix(quat):
    qw, qx, qy, qz = quat
    return np.array([
        [1 - 2 * qy ** 2 - 2 * qz ** 2, 2 * qx * qy - 2 * qz * qw, 2 * qx * qz + 2 * qy * qw],
        [2 * qx * qy + 2 * qz * qw, 1 - 2 * qx ** 2 - 2 * qz ** 2, 2 * qy * qz - 2 * qx * qw],
        [2 * qx * qz - 2 * qy * qw, 2 * qy * qz + 2 * qx * qw, 1 - 2 * qx ** 2 - 2 * qy ** 2]
    ])


def estimate_affine_3d_fixed_scale(src_points, dst_points):
    ransac = RANSACRegressor()
    ransac.fit(src_points, dst_points)
    inlier_mask = ransac.inlier_mask_

    src = src_points[inlier_mask]
    dst = dst_points[inlier_mask]

    # Normalize the input points to ensure uniform scaling
    src_mean = np.mean(src, axis=0)
    dst_mean = np.mean(dst, axis=0)
    src_centered = src - src_mean
    dst_centered = dst - dst_mean

    # Estimate rotation using singular value decomposition (SVD)
    U, _, Vt = np.linalg.svd(np.dot(dst_centered.T, src_centered))
    R_est = np.dot(U, Vt)

    # Ensure a right-handed coordinate system
    if np.linalg.det(R_est) < 0:
        Vt[2, :] *= -1
        R_est = np.dot(U, Vt)

    # Calculate the uniform scale
    scale = np.sum(dst_centered * (R_est @ src_centered.T).T) / np.sum(src_centered ** 2)

    # Construct the affine transformation matrix
    transform = np.eye(4)
    transform[:3, :3] = scale * R_est
    transform[:3, 3] = dst_mean - scale * R_est @ src_mean

    return transform