import numpy as np
from numpy import sin, cos
import matplotlib.pyplot as plt

def dhmat(delta, d, a, alpha):
    return np.array([[cos(delta), -sin(delta)*cos(alpha), sin(delta)*sin(alpha), a*cos(delta)],
                     [sin(delta), cos(delta)*cos(alpha), -cos(delta)*sin(alpha), a*sin(delta)],
                     [0, sin(alpha), cos(alpha), d],
                     [0, 0, 0, 1]])

def main():
    theta_1 = np.radians(np.arange(20, 50, 1))
    theta_2 = np.radians(np.arange(0, 30, 1))
    d_1 = 0
    d_2 = 0
    l_1 = 0.1
    l_2 = 0.2
    alpha_1 = 0
    alpha_2 = 0

    P_0 = np.array([0, 0, 0, 1]).T
    P_1 = np.array([0, 0, 0, 1]).T
    P_2 = np.array([0, 0, 0, 1]).T

    fig, ax = plt.subplots()

    for i in range(len(theta_1)):
        P1_0 = dhmat(theta_1[i], d_1, l_1, alpha_1).dot(P_1)
        P2_0 = dhmat(theta_1[i], d_1, l_1, alpha_1).dot(dhmat(theta_2[i], d_2, l_2, alpha_2)).dot(P_2)
        ax.plot([P_0[0], P1_0[0], P2_0[0]], [P_0[1], P1_0[1], P2_0[1]], '-bo')
    ax.set_xlabel('x (m)')
    ax.set_ylabel('y (m)')
    fig.savefig('./twp_link_planar_robot.png')
    plt.show()

if __name__ == '__main__':
    main()