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

# PLAN
# 1. numpy, especially arrays
# 2. sympy
# -- General A_matrix
# -- 'Special' A_matrix
# -- Forward kinematics
# Numerical forward kinematics (numpy)


def main():
    # just some 'fancy' code so i don't have to write out all the names myself
    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
    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'

    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')

    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]])

    A = rot_z @ trans_z @ trans_x @ rot_x

    A = A.subs([(theta, theta_i), (d, d_i), (a, a_i), (alpha, alpha_i)])

    return A


# symbolic forward kinematics with sympy
def forward_kinematics(DH_table):
    A_matrices = [A_matrix(*DH) for DH in DH_table]

    FK = eye(4)  # identity matrix
    for A in A_matrices: FK @= A
    return FK


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

    # antar at d1 = 100
    d_1 = 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()