import numpy as np
from sympy import symbols, Matrix, pi, sin, cos, eye, pretty_print 

# using sumpy to be able to use symbols and print expressions

# for the use of atan2 in python, check out: https://numpy.org/doc/stable/reference/generated/numpy.arctan2.html


def main():
    joints = 3
    variable_names = ['theta', 'd', 'a', 'alpha']
    DH_variables = [f'{i} '.join(variable_names) + str(i) for i in range(1, joints + 1)]
    symbolic_DH = [symbols(DH) for DH in DH_variables]

    # The RPP manipulator used as example in previous group sessions
    # Defining its DH table as a matrix
    # symbols("variable") = treating the variable as a symbol, so we dont have to give it any value
    DH_table = [[symbols('theta_1'), symbols('d_1'), 0, 0],
                [pi/2, symbols('d_2'), 0, pi/2],
                [0, symbols('d_3'), 0, 0]]

    # FK = forward_kinematics(symbolic_DH)  # symbolic DH
    FK = forward_kinematics(DH_table)  # actual DH
    pretty_print(FK)  # print function from sympy to make it look 'pretty'

    # calling the forward kinematics function with joint variables to get values for x,y,z (the end effector coordinates)
    FK_numpy = forward_kinematics_numpy([0, 100, 200])
    x, y, z = FK_numpy[0:3, 3]  # [0:3, 3] is called a slice and extracts column [0, 3> and row 3 (the x, y, z coordinates)
    print(f'\n{x = }\n{y = }\n{z = }')


# help function for the forward kinematics
def A_matrix(theta_i, d_i, a_i, alpha_i):
    theta, d, a, alpha = symbols('theta d a alpha') # defining the symbols

    # basic homogeneous matrices
    rot_z = Matrix([[cos(theta), -sin(theta), 0, 0],
                    [sin(theta), cos(theta), 0, 0],
                    [0, 0, 1, 0],
                    [0, 0, 0, 1]])

    trans_z = Matrix([[1, 0, 0, 0],
                      [0, 1, 0, 0],
                      [0, 0, 1, d],
                      [0, 0, 0, 1]])

    trans_x = Matrix([[1, 0, 0, a],
                      [0, 1, 0, 0],
                      [0, 0, 1, 0],
                      [0, 0, 0, 1]])

    rot_x = Matrix([[1, 0, 0, 0],
                    [0, cos(alpha), -sin(alpha), 0],
                    [0, sin(alpha), cos(alpha), 0],
                    [0, 0, 0, 1]])

    # multiplying basic matrices to get the basic A-matrix (that we used for the forward kinematics)
    # PS: use correct operation for matrix multiplication(NOT elementwise multiplication)
    A = rot_z @ trans_z @ trans_x @ rot_x 

    # substituting the symbols with the given values as parameters
    A = A.subs([(theta, theta_i), (d, d_i), (a, a_i), (alpha, alpha_i)])

    return A # getting one A matrix


# symbolic forward kinematics with sympy
def forward_kinematics(DH_table):
    A_matrices = [A_matrix(*DH) for DH in DH_table] # calculating all A matrices from the parameters in the DH table

    FK = eye(4)  # identity matrix
    for A in A_matrices: FK @= A #3 multiplying the matrices together
    return FK


# if we know the numerical values, we can use numpy and fill in the values directly in the forward kinematics
def forward_kinematics_numpy(joint_variables):
    # joint_variables = [theta_1, d_2, d_3]
    t_1, d_2, d_3 = joint_variables

    d_1 = 100 # setting d1 = 100

    FK = np.array([[-np.sin(t_1), 0, np.cos(t_1), np.cos(t_1)*d_3],
                   [np.cos(t_1), 0, np.sin(t_1), np.sin(t_1)*d_3],
                   [0, 1, 0, d_1 + d_2],
                   [0, 0, 0, 1]])

    return FK


if __name__ == '__main__':
    main()