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

# ============================================================================
# PyTorch versions (for training with gradients)
# ============================================================================


def rotation_6d_to_matrix_torch(rot_6d: torch.Tensor) -> torch.Tensor:
    """
    将 6D 旋转表示转换为 3x3 旋转矩阵 (PyTorch 版本，支持梯度)

    参数:
        rot_6d: torch tensor of shape (..., 6)

    返回:
        torch tensor of shape (..., 3, 3)
    """
    assert rot_6d.shape[-1] == 6, "输入必须是 6D 向量"

    # 拆分为两个 3D 向量
    a1 = rot_6d[..., 0:5:2]  # indices 0, 2, 4
    a2 = rot_6d[..., 1:6:2]  # indices 1, 3, 5

    # 第一个向量归一化
    b1 = F.normalize(a1, dim=-1)

    # 第二个向量去掉在 b1 上的投影，然后归一化
    dot = (b1 * a2).sum(dim=-1, keepdim=True)  # 点积
    proj = dot * b1
    b2 = a2 - proj
    b2 = F.normalize(b2, dim=-1)

    # 第三个向量是叉乘
    b3 = torch.cross(b1, b2, dim=-1)

    # 拼接为旋转矩阵，列向量形式
    matrix = torch.stack([b1, b2, b3], dim=-1)

    return matrix


def geodesic_loss_torch(R1: torch.Tensor, R2: torch.Tensor) -> torch.Tensor:
    """
    计算两个旋转矩阵之间的测地距离 (PyTorch 版本，支持梯度)

    参数:
        R1: torch tensor of shape (..., 3, 3)
        R2: torch tensor of shape (..., 3, 3)

    返回:
        geodesic distance of shape (...)
    """
    R_diff = torch.matmul(R1.transpose(-1, -2), R2)

    trace = R_diff[..., 0, 0] + R_diff[..., 1, 1] + R_diff[..., 2, 2]

    cos = ((trace - 1) / 2).clamp(-1 + 1e-6, 1 - 1e-6)

    loss = 1 - cos

    return loss


# ============================================================================
# NumPy versions (original functions, kept for compatibility)
# ============================================================================


def rotation_6d_to_matrix(rot_6d: np.ndarray) -> np.ndarray:
    """
    将 6D 旋转表示转换为 3x3 旋转矩阵 (使用 Gram-Schmidt 正交化)

    参数:
        rot_6d: numpy array of shape (..., 6)

    返回:
        numpy array of shape (..., 3, 3)
    """
    assert rot_6d.shape[-1] == 6, "输入必须是 6D 向量"

    # 拆分为两个 3D 向量
    a1 = rot_6d[..., 0:5:2]
    a2 = rot_6d[..., 1:6:2]

    # 第一个向量归一化
    b1 = a1 / np.linalg.norm(a1, axis=-1, keepdims=True)

    # 第二个向量去掉在 b1 上的投影，然后归一化
    dot = np.sum(b1 * a2, axis=-1, keepdims=True)  # 点积
    proj = dot * b1
    b2 = a2 - proj
    b2 = b2 / np.linalg.norm(b2, axis=-1, keepdims=True)

    # 第三个向量是叉乘
    b3 = np.cross(b1, b2)

    # 拼接为旋转矩阵，列向量形式
    # 最终形状 (..., 3, 3)
    # print(b1.shape, b2.shape, b3.shape)
    matrix = np.stack([b1, b2, b3], axis=-1)

    return matrix


def abs_6d_2_abs_euler(action):
    # left
    # print(action.shape)
    left_xyz = action[0:3]
    left_6d = action[3:9]
    left_grip = action[9]

    # right
    right_xyz = action[10:13]
    right_6d = action[13:19]
    right_grip = action[19]

    # 6d to euler
    left_matrix = rotation_6d_to_matrix(left_6d)
    right_matrix = rotation_6d_to_matrix(right_6d)
    # left_matrix = np.dot(rotation_6d_to_matrix(left_6d),np.array([[0,0,1],[0,1,0],[-1,0,0]]))
    # right_matrix = np.dot(rotation_6d_to_matrix(right_6d),np.array([[0,0,1],[0,1,0],[-1,0,0]]))

    left_euler = R.from_matrix(left_matrix).as_euler("xyz", degrees=False)
    right_euler = R.from_matrix(right_matrix).as_euler("xyz", degrees=False)

    return np.concatenate(
        [left_xyz, left_euler, [left_grip], right_xyz, right_euler, [right_grip]]
    )


def rotation_matrix_to_6d(R):
    """
    将旋转矩阵 R 转换为 6D 旋转表示
    通过提取旋转矩阵的前六个元素
    """
    return np.concatenate([R[0, :2], R[1, :2], R[2, :2]])


def eef_euler_to_6d(puppet_arm_left_pose, puppet_arm_right_pose):
    position_left = np.array(puppet_arm_left_pose[0:3])
    position_right = np.array(puppet_arm_right_pose[0:3])
    orientation_euler_left = puppet_arm_left_pose[3:6]
    orientation_euler_right = puppet_arm_right_pose[3:6]
    gripper_left = np.array([puppet_arm_left_pose[6]])
    gripper_right = np.array([puppet_arm_right_pose[6]])
    matrix_left = R.from_euler("xyz", orientation_euler_left, degrees=False).as_matrix()
    matrix_right = R.from_euler(
        "xyz", orientation_euler_right, degrees=False
    ).as_matrix()
    left_6d = rotation_matrix_to_6d(matrix_left)
    right_6d = rotation_matrix_to_6d(matrix_right)
    return np.concatenate(
        (position_left, left_6d, gripper_left, position_right, right_6d, gripper_right),
        axis=0,
    )
