#322254-MOD-HW2
!apt update -y
!apt install ffmpeg -y
import numpy as np
from numpy.linalg import inv
import matplotlib.pyplot as plt
from matplotlib.animation import FFMpegWriter
# Defines our kinematic tree
class KinematicTreeNode:
    def __init__(self): # Constructor
        self.translation = None
        self.rotation = None
        self.child = None
# Define our rotation matrix
def RotZ(theta):
    c = np.cos(theta)
    s = np.sin(theta)
    return np.array([[c, -s],[s, c]])
# FORWARD KINEMATICS
# Function that returns the End-effector position
# Takes the base link and returns the end-effector position
def EndEffectorPos(X):
    if X.child == None: # This is the last node in this branch of the tree
        return X.translation
    else: # Return the transformed end effector position
        return X.translation + np.matmul(RotZ(X.rotation),EndEffectorPos(X.child))
# FORWARD KINEMATICS (Returns the positions of ALL the joints, not just the end-effector)
# Takes the base link and returns the position of all joints in a list
# Base link is first, then in order out to the end effector
# Returns a matrix where the first column is the coordinates of the first joint, second column is second joint, etc.
def GetAllJoints(X):
    if X.child == None:
        return X.translation
    else:
        returned_joints = GetAllJoints(X.child)
        temp, translate = np.meshgrid(returned_joints[0,:], X.translation)
        transformed_joints = np.matmul(RotZ(X.rotation),returned_joints)+translate
        
        list_of_joints = np.hstack((X.translation, transformed_joints))
        return list_of_joints
# Build the tree
base = KinematicTreeNode()
L1 = KinematicTreeNode()
L2 = KinematicTreeNode()
L3 = KinematicTreeNode()
dt = 0.01
t_vec = len(np.arange(0,10,dt))
joint_position = []
for i in np.arange(0,10,0.01):
    base.translation = np.array([[0],[0]], dtype=float) # translation to the joint
    base.rotation = 0.1*np.sin(i) # rotation of the joint
    L1.translation = np.array([[1],[0]], dtype=float)
    L1.rotation = 0.2*np.sin(i)
    L2.translation = np.array([[1],[0]], dtype=float)
    L2.rotation = 0.3*np.sin(i)
    L3.translation = np.array([[1],[0]], dtype=float)
    L3.rotation = 0
# Loop for Inverse Kinematics with Newton-Raphson
    base.child = L1
    L1.child = L2
    L2.child = L3
    L3.child = None
    ("This shows that the GetAllJoints command works")
    GetAllJoints(base)
    GetAllJoints(base)
    joint_position.append(GetAllJoints(base))
    
joint_pos = np.array(joint_position)
fig, ax = plt.subplots()
p1, = ax.plot([], [], color='black', linewidth=2)
lb, ub = -5, 5
ax.set_xlim([lb, ub])
ax.set_ylim([lb, ub])
ax.set_xlabel('X')  # X Label
ax.set_ylabel('Y')  # Y Label
ax.set_title('Forward Kinematics:')
video_title = "simulation"
x_pendulum_arm = joint_pos[-0 , 0 ]
y_pendulum_arm = joint_pos[-0 , 1 ]
p1.set_data(x_pendulum_arm, y_pendulum_arm)
print(x_pendulum_arm)
print(y_pendulum_arm)
fig.canvas.draw()
fig, ax = plt.subplots()
# Create a plot on those axes, which is currently empty
p, = ax.plot([],[], color='black',linewidth=2) # initializes an empty plot
lb, ub = -5, 5
ax.set_xlim([lb, ub])
ax.set_ylim([lb, ub])
ax.set_xlabel('X')  # X Label
ax.set_ylabel('Y')  # Y Label
ax.set_title('Forward Kinematics:')
video_title = "Question 1 simulation"
FPS = 20
sample_rate = int(1 / (FPS*dt))
dpi = 300 #Quality of video
writerObj = FFMpegWriter(fps=FPS)
simulation_size = t_vec
# Plot and Create Animation:
with writerObj.saving(fig, video_title+".mp4", dpi):
    # We want to create video that represents the simulation
    # So we need to sample only a few of the frames:
    for i in range(0, simulation_size, sample_rate):
        # Update Pendulum Arm:
        x_data_points = [joint_pos[i ,0 ]]
        y_data_points = [joint_pos[i , 1 ]]
        # We want to avoid creating new plots to make an animation (Very Slow)
        # Instead lets take the plot we made earlier and just update it with new data.
        p.set_data(x_data_points, y_data_points)  # Update plot with set_data
        # Update Pendulum Patch:
        #patch_center = x_bead[i], y_bead[i]
        #c.center = patch_center  # Same idea here, instead of drawing a new patch update its location
        # Update Drawing:
        fig.canvas.draw()  # Update the figure with the new changes
        # Grab and Save Frame:
        writerObj.grab_frame()
# We can also run it inline with the following block of code:
from IPython.display import Video
Video("/work/HW2/Question 1 simulation.mp4", embed=True, width=640, height=480)
# You can find the video's path by right clicking on it in Notebooks & Files 
# and selecting "Copy path to clipboard"