import numpy as np

"The homogeneous transformation matrix from the curriculum"
"""Params:
           Rot_z   = rotation along the z-axis
           Trans_z = translation along the z-axis
           Trans_x = translation along the x-axis
           Rot_x   = rotation along the x-axis
"""
def A_i_matrix(rot_z, trans_z, trans_x, rot_x):
    row_1 = np.array([np.cos(rot_z), -np.sin(rot_z)*np.cos(rot_x), np.sin(rot_z)*np.sin(rot_x), trans_x*np.cos(rot_z)])
    row_2 = np.array([np.sin(rot_z), np.cos(rot_z)*np.cos(rot_x), -np.cos(rot_z)*np.sin(rot_x), trans_x*np.sin(rot_z)])
    row_3 = np.array([0, np.sin(rot_x), np.cos(rot_x), trans_z])
    row_4 = np.array([0, 0, 0, 1])

    A_i = [row_1,
           row_2,
           row_3,
           row_4]

    return np.array(A_i)

"Forward kinematics for the spherical robot (RRP)"
"""Params:
        q = List of joint variables. REMEMBER the joint angles should arrive in radians.
"""
def forward_kinematics(q):
    L1, L2 = 1, 1

    A_1 = A_i_matrix(q[0], L1, 0, np.pi/2)
    A_2 = A_i_matrix(q[1]+np.pi/2, 0, 0, np.pi/2)
    A_3 = A_i_matrix(0, L2+q[2], 0, 0)
    A_02 = np.matmul(A_1, A_2)
    A_03 = np.matmul(A_02, A_3)

    return A_03

"Inverse kinematics for the Spherical robot (RRP)"
"""Params:
        cart_coordinates = point in cartesian space in (x,y,z) coordinates.
"""
def inverse_kinematics(cart_coordinates):
    x, y, z = cart_coordinates[0], cart_coordinates[1], cart_coordinates[2]
    L1, L2 = 1, 1

    s = z-L1
    r = np.sqrt(x**2+y**2)

    # First solution
    theta_1 = np.arctan2(y,x)
    theta_2 = np.arctan2(s,r)
    d_3 = np.sqrt(r**2+s**2)-L2

    # Second solution
    theta_1_rotated = np.arctan2(y,x)+np.pi
    theta_2_rotated = np.arctan2(s,-r)
    d_3 = np.sqrt(r**2+s**2)-L2

    first_solution = [np.rad2deg(theta_1), np.rad2deg(theta_2), d_3]
    second_solution = [np.rad2deg(theta_1_rotated), np.rad2deg(theta_2_rotated), d_3]

    return [first_solution, second_solution]

"Jacobian kinematics for the spherical Robot (RRP)"
"""Params:
    q =  List of joint variables. REMEMBER the joint angles should arrive in radians.
"""
def jacobian_kinematics(q):
    L1, L2 = 1,1

    c1, s1 = np.cos(q[0]), np.sin(q[0])
    c2, s2 = np.cos(q[1]), np.sin(q[1])
    d3 = q[2]

    z_0 = np.array([0, 0, 1])
    z_1 = np.array([s1, -c1, 0])
    z_2 = np.array([c1*c2, c2*s1, s2])

    O3_0 = np.array([c1*c2*(L2+d3), c2*s1*(L2+d3), L1+s2*(d3)])
    O3_1 = np.array([c1*c2*(L2+d3), c2*s1*(L2+d3), s2*(d3)])

    jv_1 = np.cross(z_0, O3_0)
    jv_2 = np.cross(z_1, O3_1)
    jv_3 = z_2

    jw_1 = z_0
    jw_2 = z_1
    jw_3 = [0, 0, 0]

    Jv = np.array([jv_1, jv_2, jv_3])
    Jw = np.array([jw_1, jw_2, jw_3])

    jacobian_matrix = np.array([Jv, Jw])
    return jacobian_matrix

"The determinant of jacobian for the Spherical Robot(RRP)"
"""Params:
    J = The jacobian matrix
"""
def determinant_jacobian(J):
    #TODO: Since python does not have an efficient symbolic approach such as matlab, but we can still use it by using the python package
    #Sympy.
    # I will update this later with sympy. 

    return

def main(_):
    joint_angle1 = np.pi/4
    joint_angle2 = np.pi/4
    d_3 = 1.5
    q = [joint_angle1, joint_angle2, d_3]
    q_dot = [0.05, 1.0, 0.05]
    p_xyz = [1.25, 1.25, 2.7678]

    fk = forward_kinematics(q)
    #ik = inverse_kinematics(p_xyz)
    ik = inverse_kinematics([fk[0][3], fk[1][3], fk[2][3]])
    jacobian = jacobian_kinematics(q)
    v = jacobian[0]*q_dot
    print(fk)
    print(ik)
    print(jacobian)
    print(v)
if __name__ == '__main__':
    main(None)