import numpy as np
import matplotlib.pyplot as plt
import time


from arm3d import Robot3D
from conveyor_belt import *

def evalRobot3D(l, theta):
    x=(l[0]*np.cos(theta[0]) + l[1]*np.cos(theta[0]+theta[1]))*np.cos(theta[2])
    y=(l[0]*np.cos(theta[0]) + l[1]*np.cos(theta[0]+theta[1]))*np.sin(theta[2])
    z=(l[0]*np.sin(theta[0]) + l[1]*np.sin(theta[0]+theta[1]))

    pos = np.array([x, y, z])
    xp1=l[0]*np.cos(theta[0]) + l[1]*np.cos(theta[0]+theta[1])

    dq0 = -(l[0]*np.sin(theta[0]) + l[1]*np.sin(theta[0]+theta[1]))
    dq1 = -l[1]*np.sin(theta[0]+theta[1])
    j= np.array([[dq0*np.cos(theta[2]), dq1*np.cos(theta[2]), -xp1*np.sin(theta[2])],
        [dq0*np.sin(theta[2]), dq1*np.sin(theta[2]), xp1*np.cos(theta[2])],
        [xp1, l[1]*np.cos(theta[0]+theta[1]), 0]])

    return pos, j


# %%%%            Instructions            %%%%%%%%%%%%%%
# %%% Keep in mind the following lines of codes
# %%% Use 'plt.cla(); plt.pause(0.01);' at the beginning of your primary loop
# %%% Once you pick an object plot it at the end effector using 
# %%%       [pos_hand, J] = evalRobot3D(ls, t); 
# %%%        plotCylinderWithCaps(0.1,pos_hand,0.2,12,[1 0 0],'z');
# %%%        env.plot_cans()
# %%% Plot the complete scene using:
# %%%        env()
# %%% Make sure you update the obj_picked and obj_placed properly

# %%%%%%% Enter your code here    %%%%%%%%%%%%%%%%%%%%%


if __name__ == "__main__":
    plotRobot3D = Robot3D()

    env = ConveyorBeltPickPlace()

    while True:
        plt.cla()

        curr_pos, _ = evalRobot3D(env.robot.ls, env.theta)