# Two functions
# convert (rotation matrix,translation vector) to (global 3d position, euler angles)
# convert (global 3d position, euler angles) to (rotation matrix,translation vector)
# Used functions from 'Rotation Matrix To Euler Angles' by Satya Mallick
# https://learnopencv.com/rotation-matrix-to-euler-angles/
import math
import numpy as np


# Calculates Rotation Matrix given euler angles.
def eulerAnglesToRotationMatrix(theta):
    # inputs
    # -- theta: tuple/list (theta_x,theta_y,theta_z) euler angles (in radians)

    R_x = np.array([[1, 0, 0],
                    [0, math.cos(theta[0]), -math.sin(theta[0])],
                    [0, math.sin(theta[0]), math.cos(theta[0])]
                    ])

    R_y = np.array([[math.cos(theta[1]), 0, math.sin(theta[1])],
                    [0, 1, 0],
                    [-math.sin(theta[1]), 0, math.cos(theta[1])]
                    ])

    R_z = np.array([[math.cos(theta[2]), -math.sin(theta[2]), 0],
                    [math.sin(theta[2]), math.cos(theta[2]), 0],
                    [0, 0, 1]
                    ])

    R = np.dot(R_z, np.dot(R_y, R_x))

    return R


# Checks if a matrix is a valid rotation matrix.
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


# Calculates rotation matrix to euler angles
# The result is the same as MATLAB except the order
# of the euler angles ( x and z are swapped ).
def rotationMatrixToEulerAngles(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])


# convert (rotation matrix,translation vector) to (global 3d position, euler angles)
# output is the pose at the center, which is used as the origin of its own frame
# e.g. tag center of apriltag
def transformation_to_pose(transformation):
    # input
    # -- transformation: tuple/list (3x3 rotation matrix R, 3x1 translation vector t)
    # output
    # -- pose: tuple (global position (x,y,z), euler angles (theta_x,theta_y,theta_z))
    assert isinstance(transformation, tuple) or isinstance(transformation, list)
    assert len(transformation) == 2
    R, t = transformation
    assert R.shape == (3, 3) and t.size == 3
    pos = t.reshape(3, )
    angles = rotationMatrixToEulerAngles(R)

    return (pos, angles)


# convert (global 3d position, euler angles) to (rotation matrix,translation vector)
# must be the pose data at the center, which is used as the origin of its own frame
def pose_to_transformation(pose):
    # input
    # -- pose: tuple (global position (x,y,z), euler angles (theta_x,theta_y,theta_z))
    # output
    # -- transformation: tuple/list (3x3 rotation matrix R, 3x1 translation vector t)
    assert isinstance(pose, tuple) or isinstance(pose, list)
    assert len(pose) == 2
    pos, angles = pose
    assert len(pos) == 3 and len(angles) == 3

    t = (np.array(pos)).reshape(3, )

    R = eulerAnglesToRotationMatrix(angles)

    return (R, t)


def rotation_align(z,d):
    """
    https://iquilezles.org/www/articles/noacos/noacos.htm
    compute rotation matrix to align vector z to vector d
    :param z:
    :param d:
    :return: rotation matrix M
    """
    assert np.isclose(np.linalg.norm(z),1)
    assert np.isclose(np.linalg.norm(d), 1)
    v = np.cross(z,d)
    c = np.sum(z*d)

    M = v.reshape(-1,1) @ v.reshape(1,-1) / (1+c) + np.array([
        [c, -v[2], v[1]],
        [v[2], c, -v[0]],
        [-v[1], v[0], c]
    ])

    return M