# -*- coding: utf-8 -*-
"""
Spyder Editor

This is a temporary script file.
"""

import numpy as np
import matplotlib.pyplot as plt
from math import *
 
a1 = 0.6
a2 = 0.4
def drawRobot (th1, th2, color) :
     point0 = np.array([0, 0])
     point1 = np.array([a1*cos(th1), a1*sin(th1)])
     point2 = point1 + np.array([a2*cos(th1+th2), a2*sin(th1+th2)])
     x_values = [point0[0], point1[0], point2[0]]
     y_values = [point0[1], point1[1], point2[1]]
     plt.plot(x_values, y_values, color)
     plt.plot(point2[0], point2[1], color='red', marker='o')
     plt.gca().axis('equal')
      
 # For Q 5-7
def Jacobian(theta) :
    # Unpack vars
    t1, t2 = theta
    # Return J(theta)
    return np.array([
                [-a1*sin(t1)-a2*sin(t1+t2), -a2*sin(t1+t2)],
                [a1*cos(t1)+a2*cos(t1+t2), a2*cos(t1+t2)]])
# For Q11
def FK(theta):
    th1, th2 = theta
    return np.array([a1*cos(th1), a1*sin(th1)]) + np.array([a2*cos(th1+th2), a2*sin(th1+th2)])

#------------------------------------------------------------------
def Q4():
    plt.figure()
    q_0 = np.array([0, 0]).T
    dq = np.array([3, -2]).T
    # Plot initial pose
    drawRobot(q_0[0], q_0[1], 'black')
    # Iterate for 100 periods
    q = q_0
    dT = 0.01
    for i in range (0,100): 
        q = q + dq * dT
        drawRobot(q[0], q[1], 'blue')
        

def Q6(q_0, dX_desired):
    plt.figure()
    # Plot initial pose
    drawRobot(q_0[0], q_0[1], 'black')
    # Iterate for 100 periods
    q = q_0
    dT = 0.01
    for i in range (0,100): 
        J = Jacobian(q)
        Jinv = np.linalg.inv(J)
        dq = Jinv @ dX_desired
        q = q + dq * dT
        drawRobot(q[0], q[1], 'blue')
        
def Q10(q_0, dX_desired, val_lambda):
    plt.figure()
    # Plot initial pose
    drawRobot(q_0[0], q_0[1], 'black')
    # Iterate for 100 periods
    q = q_0
    dT = 0.01
    for i in range (0,100): 
        J = Jacobian(q)
        # Damped least-square inverse
        # Wampler (1986) / Nakamura and Hanafusa (1986)
        Jinv = J.T @ np.linalg.inv(J@J.T + val_lambda**2 * np.eye(2))
        dq = Jinv @ dX_desired
        q = q + dq * dT
        drawRobot(q[0], q[1], 'blue')
        
def Q11(q_0, X_desired, val_lambda, P):
    plt.figure()
    # Plot initial pose
    drawRobot(q_0[0], q_0[1], 'black')
    # Iterate for 100 periods
    q = q_0
    dT = 0.01
    for i in range (0,100): 
        J = Jacobian(q)
        error = X_desired - FK(q)
        # Proportional gain on position error
        #print("||error|| = ", np.linalg.norm(error))
        dX_desired = P * error
        # Damped least-square inverse for IDKM
        # Wampler (1986) / Nakamura and Hanafusa (1986)
        Jinv = J.T @ np.linalg.inv(J@J.T + val_lambda**2 * np.eye(2))
        dq = Jinv @ dX_desired
        q = q + dq * dT
        drawRobot(q[0], q[1], 'blue')
        
def Q12(q_0, val_lambda, P):
    plt.figure()
    # Plot initial pose
    drawRobot(q_0[0], q_0[1], 'black')
    # Iterate for 100 periods
    q = q_0
    dT = 0.01
    # Define circle
    t = 0
    center = np.array([0.3, 0.3]).T
    radius = 0.1
    def func_Xd(t) :
        return center + radius * np.array([cos(2*t), sin(2*t)])
    for i in range (0,900): 
        t += dT
        X_desired = func_Xd(t)
        J = Jacobian(q)
        error = X_desired - FK(q)
        #print("||error|| = ", np.linalg.norm(error))
        dX_desired = P * error
        # Damped least-square inverse
        # Wampler (1986) / Nakamura and Hanafusa (1986)
        Jinv = J.T @ np.linalg.inv(J@J.T + val_lambda**2 * np.eye(2))
        dq = Jinv @ dX_desired
        q = q + dq * dT
        drawRobot(q[0], q[1], 'blue')
        plt.plot(X_desired[0], X_desired[1], 'green', marker='X')
        
#----------------------------------------------------------------------------
        
if __name__ =='__main__':
    Q4()
    
    Q6( q_0 = np.array([0.2, 0.6]).T,
        dX_desired = np.array([0, -0.5]).T )
    # Q8 -> raise LinAlgError("Singular matrix")
    '''
    Q6( q_0 = np.array([0., 0.]).T,
        dX_desired = np.array([-0.1, 0]).T )
    '''
    J = Jacobian(np.array([0., 0.]).T)
    print("J = \n", J)
    print("det(J) = ", np.linalg.det(J))
    
    # Q9
    q_init = np.array([0.001, -0.001]).T
    dX = np.array([-0.1, 0]).T
    Q6( q_0 = q_init,
        dX_desired = dX )
    J = Jacobian(q_init)
    print("J = \n", J)
    print("det(J) = ", np.linalg.det(J))
    
    # Q10
    Q10( q_0 = q_init,
         dX_desired = dX,
         val_lambda = 0.01)
    
    #Q11
    target = np.array([-0.7, 0.3]).T
    Q11(q_0 = q_init,
        X_desired = target,
        val_lambda = 0.01,
        P= 10)
    plt.plot(target[0], target[1], 'green', marker='X')
    
    #Q12
    Q12(q_0 = q_init,
        val_lambda = 0.1,
        P= 20)