import numpy as np

############################## Here comes the Python code ##############################
def rotation_matrix_from_vectors(vec1, vec2):

    """

    Find the rotation matrix that aligns vec1 to vec2

    This function is taken from Peter's answer on stackoverflow:
    `Question Link <https://stackoverflow.com/questions/45142959/calculate-rotation-matrix-to-align-two-vectors-in-3d-space>`_

    Parameters:

    vec1:
        Numpy array. A 3d "source" vector

    vec2:
        Numpy array. A 3d "destination" vector

    Return:
        A transform matrix (3x3) which when applied to vec1, aligns it with vec2.

    """

    a, b = (vec1 / np.linalg.norm(vec1)).reshape(3), (vec2 / np.linalg.norm(vec2)).reshape(3)

    print(a)
    print(b)

    tmp = 0

    for item, num_a in enumerate(a):

        num_b = b[item]

        if (num_a != 0) and (num_b != 0):

            tmp += num_a / num_b

    tmp = tmp / 3 # calculate the average

    print(tmp)

    v = np.cross(a, b)

    c = np.dot(a, b)

    s = np.linalg.norm(v)

    if abs(s) < 0.0000000000001: # which means vec1 and vec2 are on the same line, so their cross is 0

        if tmp > 0:

            return [[1, 0, 0],
                    [0, 1, 0],
                    [0, 0, 1]]

        else:

            rotation_matrix = [[-1, 0, 0],
                               [0, -1, 0],
                               [0, 0, -1]]

            return rotation_matrix

    kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])

    rotation_matrix = np.eye(3) + kmat + kmat.dot(kmat) * ((1 - c) / (s ** 2))

    return rotation_matrix

def matrix_multiple(matrix_list):

    """
    This function can help to return the multiplication of a series of matrixes.

    Parameters:

    matrix_list:
        The list of matrix that we want to multiply

    Return:
        Numpy array. The result that we want
    """

    # set the identity matrix as the same dimension as the input matrix
    return_matrix = np.identity(np.shape(matrix_list[0])[0])

    # start multiplication
    for matrix in matrix_list:

        return_matrix = np.matmul(return_matrix, matrix)

    return return_matrix

def rotation_along_two_points(p, p1, p2, theta):

    """
    This function can rotate a point (p) at an angle (theta) along a line generated by two point (p1, p2)

    Parameters:

    p:
        Numpy Array (3x1). The point that we want to rotate

    p1:
        Numpy Array (3x1). First point of the line

    p2:
        Numpy Array (3x1). Second point of the line

    theta:
        Real. Angle that we want to rotate (0 - 360 degree)

    Return:
        Numpy Array (3x1). The point from rotation.

    Reference:

        The algorithm of this function is copied from two sources:

        1. https://robotics.stackexchange.com/questions/12782/how-rotate-a-point-around-an-arbitrary-line-in-3d

        2. http://paulbourke.net/geometry/rotate/

        Many thanks to the authors
    """

    # change the angle to radian
    rad = np.radians(theta)

    # calculate the unit vector of p1 -> p2
    vector_p1_p2 = p2 - p1

    norm_vector_p1_p2 = vector_p1_p2 / np.linalg.norm(vector_p1_p2)

    # the parameter are copied from the answer
    a = norm_vector_p1_p2[0]

    b = norm_vector_p1_p2[1]

    c = norm_vector_p1_p2[2]

    d = np.sqrt(np.power(b, 2) + np.power(c, 2))

    x1 = p1[0]

    y1 = p1[1]

    z1 = p1[2]

    # construct all the matrix

    if d != 0:

        matrix_T = np.array([
                            [1, 0, 0, -x1],
                            [0, 1, 0, -y1],
                            [0, 0, 1, -z1],
                            [0, 0, 0, 1]
                            ])

        matrix_Rx = np.array([
                             [1, 0, 0, 0],
                             [0, c/d, -b/d, 0],
                             [0, b/d, c/d, 0],
                             [0, 0, 0, 1]
                             ])

        matrix_Ry = np.array([
                             [d, 0, -a, 0],
                             [0, 1, 0, 0],
                             [a, 0, d, 0],
                             [0, 0, 0, 1]
                             ])

        matrix_Rz = np.array([
                             [np.cos(rad), -np.sin(rad), 0, 0],
                             [np.sin(rad), np.cos(rad), 0, 0],
                             [0, 0, 1, 0],
                             [0, 0, 0, 1]
                             ])

        # create 4x1 vector for p
        p_4_1 = np.array([p[0], p[1], p[2], 1])

        # do the matrix inverse
        matrix_inv_T = np.linalg.inv(matrix_T)

        matrix_inv_Rx = np.linalg.inv(matrix_Rx)

        matrix_inv_Ry = np.linalg.inv(matrix_Ry)

        matrix_inv_Rz = np.linalg.inv(matrix_Rz)

        # multiple all matrixes
        matrix_mul = matrix_multiple([matrix_inv_T, matrix_inv_Rx, matrix_inv_Ry, matrix_Rz, matrix_Ry, matrix_Rx, matrix_T])

        # calculate the final point
        p_final = np.matmul(matrix_mul, p_4_1)

        # calcualte the return
        p_return = np.array([p_final[0], p_final[1], p_final[2]])

        return p_return

    else:

        matrix_T = np.array([
                    [1, 0, 0, -x1],
                    [0, 1, 0, -y1],
                    [0, 0, 1, -z1],
                    [0, 0, 0, 1]
                    ])

        matrix_Ry = np.array([
                     [d, 0, -a, 0],
                     [0, 1, 0, 0],
                     [a, 0, d, 0],
                     [0, 0, 0, 1]
                     ])

        matrix_Rz = np.array([
                             [np.cos(rad), -np.sin(rad), 0, 0],
                             [np.sin(rad), np.cos(rad), 0, 0],
                             [0, 0, 1, 0],
                             [0, 0, 0, 1]
                             ])

        # create 4x1 vector for p
        p_4_1 = np.array([p[0], p[1], p[2], 1])

        # do the matrix inverse
        matrix_inv_T = np.linalg.inv(matrix_T)

        matrix_inv_Ry = np.linalg.inv(matrix_Ry)

        matrix_inv_Rz = np.linalg.inv(matrix_Rz)

        # multiple all matrixes
        matrix_mul = matrix_multiple([matrix_inv_T, matrix_inv_Ry, matrix_Rz, matrix_Ry, matrix_T])

        # calculate the final point
        p_final = np.matmul(matrix_mul, p_4_1)

        # calcualte the return
        p_return = np.array([p_final[0], p_final[1], p_final[2]])

        return p_return
