import numpy as np

from scipy.spatial.transform import Rotation
from anydata.utils.decorators import iterate1


@iterate1
def invert_extrinsics(extrinsics):
    """Inverts a transformation matrix (extrinsics)"""
    inv_extrinsics = np.eye(4)
    inv_extrinsics[:3, :3] = np.transpose(extrinsics[:3, :3])
    inv_extrinsics[:3, -1] = - inv_extrinsics[:3, :3] @ extrinsics[:3, -1]
    return inv_extrinsics


def rotx(t):
    c, s = np.cos(t), np.sin(t)
    return np.array([[1,  0,  0],
                     [0,  c, -s],
                     [0,  s,  c]])


def roty(t):
    c, s = np.cos(t), np.sin(t)
    return np.array([[c,  0,  s],
                     [0,  1,  0],
                     [-s, 0,  c]])


def rotz(t):
    c, s = np.cos(t), np.sin(t)
    return np.array([[c, -s,  0],
                     [s,  c,  0],
                     [0,  0,  1]])


def transform_from_rot_trans(R, t):
    """Transformation matrix from rotation matrix and translation vector"""
    R, t = R.reshape(3, 3), t.reshape(3, 1)
    return np.vstack((np.hstack([R, t]), [0, 0, 0, 1]))


def pose_to_matrix(pose):
    """Convert 7D pose [x, y, z, qx, qy, qz, qw] to 4x4 transformation matrix"""
    T = np.eye(4)
    T[:3, :3] = Rotation.from_quat(pose[3:]).as_matrix()
    T[:3, 3] = pose[:3]
    return T