'''
获取pose矩阵

required:
    - numpy
'''

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

def qvec2rotmat(qvec):
    '''将四元数转换为旋转矩阵

    Args:
        qvec: 四元数

    Returns:
        旋转矩阵
    '''
    return Rotation.from_quat([qvec[1], qvec[2], qvec[3], qvec[0]]).as_matrix()
    """
    return np.array(
        [
            [
                1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2,
                2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],
                2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2],
            ],
            [
                2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],
                1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2,
                2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1],
            ],
            [
                2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],
                2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],
                1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2,
            ],
        ]
    )
    #"""
def get_w2c(qvec: list, tvec: list) -> np.ndarray:
    '''获取世界坐标系到相机坐标系的变换矩阵

    Args:
        qvec: 四元数
        tvec: 平移向量

    Returns:
        世界坐标系到相机坐标系的变换矩阵
    '''
    #"""
    w2c = np.eye(4).astype(np.float64)
    w2c[:3, :3] = qvec2rotmat(qvec)
    w2c[:3, 3] = tvec
    #"""    
    """
    R = qvec2rotmat(qvec)  # 相机到世界的旋转
    R_w2c = R.T                 # 世界到相机的旋转
    t_w2c = -R_w2c @ tvec       # 平移分量

    w2c = np.eye(4)
    w2c[:3, :3] = R_w2c
    w2c[:3, 3] = t_w2c
    #"""
    """
    R2 = Rotation.from_quat([qvec[0], qvec[1], qvec[2], qvec[3]]).as_matrix()
    # w2c = [R | -R @ tvec]
    w2c2 = np.eye(4)
    w2c2[:3, :3] = R2
    w2c2[:3, 3] = -R2 @ tvec
    print("w2c", w2c)
    print("w2c2", w2c2)
    #"""
    """
    w2c = np.eye(4).astype(np.float64)
    R = qvec2rotmat(qvec)
    w2c[:3, :3] = R.T
    w2c[:3, 3] = -R @ tvec
    #"""
    """
    w2c = np.eye(4).astype(np.float64)
    R = Rotation.from_quat([qvec[1], qvec[2], qvec[3], qvec[0]]).as_matrix()
    # COLMAP的tvec是相机在世界坐标系中的位置，需取反
    t = -R.T @ tvec
    w2c = np.eye(4)
    w2c[:3, :3] = R.T  # 旋转部分
    w2c[:3, 3] = t     # 平移部分
    #"""
    return w2c

def get_c2w(qvec: list, tvec: list) -> np.ndarray:
    '''获取相机坐标系到世界坐标系的变换矩阵

    Args:
        qvec: 四元数
        tvec: 平移向量

    Returns:
        相机坐标系到世界坐标系的变换矩阵
    '''
    c2w = np.eye(4).astype(np.float64)
    c2w[:3, :3] = qvec2rotmat(qvec).T
    c2w[:3, 3] = -qvec2rotmat(qvec).T @ tvec
    return c2w