#!/usr/bin/env python
"""Kinematic function skeleton code for HW2, Part 3.
Course: EE 125, Fall 2014
Author: Aaron Bestick, 9/10/14

This Python file is a code skeleton for Part 3 of Homework 2. You should fill in 
the body of the eight empty methods below so that they implement the kinematic 
functions described in the homework assignment.

When you think you have the methods implemented correctly, you can test your 
code by running "python kin_func_skeleton.py at the command line.

This code requires the NumPy and SciPy libraries. If you don't already have 
these installed on your personal computer, you can use the lab machines or 
the Ubuntu+ROS VM on the course page to complete this portion of the homework.
"""

import numpy as np
import scipy as sp
from scipy import linalg

np.set_printoptions(precision=4,suppress=True)

def skew_3d(omega):
    """
    Converts a rotation vector in 3D to its corresponding skew-symmetric matrix.
    
    Args:
    omega - (3,) ndarray: the rotation vector
    
    Returns:
    omega_hat - (3,3) ndarray: the corresponding skew symmetric matrix
    """
    if not omega.shape == (3,):
        raise TypeError('omega must be a 3-vector')
    
    omega_hat = np.array([[0,-omega[2],omega[1]],
                          [omega[2],0,-omega[0]],
                          [-omega[1],omega[0],0]])

    return omega_hat

def rotation_2d(theta):
    """
    Computes a 2D rotation matrix given the angle of rotation.
    
    Args:
    theta: the angle of rotation
    
    Returns:
    rot - (2,2) ndarray: the resulting rotation matrix
    """
    
    rot = np.array([[np.cos(theta),-np.sin(theta)],
                   [np.sin(theta),np.cos(theta)]])

    return rot

def rotation_3d(omega, theta):
    """
    Computes a 3D rotation matrix given a rotation axis and angle of rotation.
    
    Args:
    omega - (3,) ndarray: the axis of rotation
    theta: the angle of rotation
    
    Returns:
    rot - (3,3) ndarray: the resulting rotation matrix
    """
    if not omega.shape == (3,):
        raise TypeError('omega must be a 3-vector')
    
    omega_hat = skew_3d(omega)
    rot = linalg.expm(omega_hat*theta)

    return rot

def hat_2d(xi):
    """
    Converts a 2D twist to its corresponding 3x3 matrix representation
    
    Args:
    xi - (3,) ndarray: the 2D twist
    
    Returns:
    xi_hat - (3,3) ndarray: the resulting 3x3 matrix
    """
    if not xi.shape == (3,):
        raise TypeError('omega must be a 3-vector')

    xi_hat = np.array([[0,-xi[2],xi[0]],
                       [xi[2],0,xi[1]],
                       [0,0,0]])

    return xi_hat

def hat_3d(xi):
    """
    Converts a 3D twist to its corresponding 4x4 matrix representation
    
    Args:
    xi - (6,) ndarray: the 3D twist
    
    Returns:
    xi_hat - (4,4) ndarray: the corresponding 4x4 matrix
    """
    if not xi.shape == (6,):
        raise TypeError('xi must be a 6-vector')

    xi_hat = np.array([[0,-xi[5],xi[4],xi[0]],
                       [xi[5],0,-xi[3],xi[1]],
                       [-xi[4],xi[3],0,xi[2]],
                       [0,0,0,0]])

    return xi_hat

def homog_2d(xi, theta):
    """
    Computes a 3x3 homogeneous transformation matrix given a 2D twist and a 
    joint displacement
    
    Args:
    xi - (3,) ndarray: the 2D twist
    theta: the joint displacement
    
    Returns:
    g - (3,3) ndarray: the resulting homogeneous transformation matrix
    """
    if not xi.shape == (3,):
        raise TypeError('xi must be a 3-vector')

    xi_hat = hat_2d(xi)
    g = linalg.expm(theta*xi_hat)

    return g

def homog_3d(xi, theta):
    """
    Computes a 4x4 homogeneous transformation matrix given a 3D twist and a 
    joint displacement.
    
    Args:
    xi - (6,) ndarray: the 3D twist
    theta: the joint displacement

    Returns:
    g - (4,4) ndarary: the resulting homogeneous transformation matrix
    """
    if not xi.shape == (6,):
        raise TypeError('xi must be a 6-vector')

    xi_hat = hat_3d(xi)
    g = linalg.expm(theta*xi_hat)

    return g

def prod_exp(xi, theta):
    """
    Computes the product of exponentials for a kinematic chain, given 
    the twists and displacements for each joint.
    
    Args:
    xi - (6,N) ndarray: the twists for each joint
    theta - (N,) ndarray: the displacement of each joint
    
    Returns:
    g - (4,4) ndarray: the resulting homogeneous transformation matrix
    """
    if not xi.shape[0] == 6:
        raise TypeError('xi must be a 6xN')
    N = xi.shape[1]
    g = np.eye(4,4)
    for i in range(0,N):
        xi_hat = hat_3d(xi[:,i])
        g = g.dot(linalg.expm(theta[i]*xi_hat))

    return g

def left_arm(theta):
    """
    Computes the coordinate tranformation between the base and tool frames
    Args:
    theta - (7,) ndarray: joint angles
    Returns:
    g - (4,4) ndarray: the resulting homogenous transformation matrix
    """
    if not theta.shape == (7,):
        raise TypeError('xi must be a 7-vector')

    q1 = np.array([0.0635,0.2598,0.1188])
    q2 = np.array([0.1106,0.3116,0.3885])
    q3 = np.array([0.1827,0.3838,0.3881])
    q4 = np.array([0.3682,0.5684,0.3181])
    q5 = np.array([0.4417,0.6420,0.3177])
    q6 = np.array([0.6332,0.8337,0.3067])
    q7 = np.array([0.7152,0.9158,0.3063])
    w1 = np.array([-0.0059,0.0113,0.9999])
    w2 = np.array([-0.7088,0.7065,-0.0122])
    w3 = np.array([0.7065,0.7077,-0.0038])
    w4 = np.array([-0.7077,0.7065,-0.0122])
    w5 = np.array([0.7065,0.7077,-0.0038])
    w6 = np.array([-0.7077,0.7065,-0.0122])
    w7 = np.array([0.7065,0.7077,-0.0038])
    gst0 = np.array([[1,0,0,0.7957],
                     [0,1,0,0.9965],
                     [0,0,1,0.3058],
                     [0,0,0,1]])
    v1 = -np.cross(w1,q1)
    v2 = -np.cross(w2,q2)
    v3 = -np.cross(w3,q3)
    v4 = -np.cross(w4,q4)
    v5 = -np.cross(w5,q5)
    v6 = -np.cross(w6,q6)
    v7 = -np.cross(w7,q7)
    xi1 = np.reshape(np.array([v1,w1]),6)
    xi2 = np.reshape(np.array([v2,w2]),6)
    xi3 = np.reshape(np.array([v3,w3]),6)
    xi4 = np.reshape(np.array([v4,w4]),6)
    xi5 = np.reshape(np.array([v5,w5]),6)
    xi6 = np.reshape(np.array([v6,w6]),6)
    xi7 = np.reshape(np.array([v7,w7]),6)
    xi = np.array([xi1,xi2,xi3,xi4,xi5,xi6,xi7]).T
    product = prod_exp(xi, theta)
    g = product.dot(gst0)
    return g
       

#-----------------------------Testing code--------------------------------------
#-------------(you shouldn't need to modify anything below here)----------------

def array_func_test(func_name, args, ret_desired):
    ret_value = func_name(*args)
    if not isinstance(ret_value, np.ndarray):
        print('[FAIL] ' + func_name.__name__ + '() returned something other than a NumPy ndarray')
    elif ret_value.shape != ret_desired.shape:
        print('[FAIL] ' + func_name.__name__ + '() returned an ndarray with incorrect dimensions')
    elif not np.allclose(ret_value, ret_desired, rtol=1e-3):
        print('[FAIL] ' + func_name.__name__ + '() returned an incorrect value')
    else:
        print('[PASS] ' + func_name.__name__ + '() returned the correct value!')

if __name__ == "__main__":
    print('Testing...')

    #Test skew_3d()
    arg1 = np.array([1.0, 2, 3])
    func_args = (arg1,)
    ret_desired = np.array([[ 0., -3.,  2.],
                            [ 3., -0., -1.],
                            [-2.,  1.,  0.]])
    array_func_test(skew_3d, func_args, ret_desired)

    #Test rotation_2d()
    arg1 = 2.658
    func_args = (arg1,)
    ret_desired = np.array([[-0.8853, -0.465 ],
                            [ 0.465 , -0.8853]])
    array_func_test(rotation_2d, func_args, ret_desired)

    #Test rotation_3d()
    arg1 = np.array([2.0, 1, 3])
    arg2 = 0.587
    func_args = (arg1,arg2)
    ret_desired = np.array([[-0.1325, -0.4234,  0.8962],
                            [ 0.8765, -0.4723, -0.0935],
                            [ 0.4629,  0.7731,  0.4337]])
    array_func_test(rotation_3d, func_args, ret_desired)

    #Test hat_2d()
    arg1 = np.array([2.0, 1, 3])
    func_args = (arg1,)
    ret_desired = np.array([[ 0., -3.,  2.],
                            [ 3.,  0.,  1.],
                            [ 0.,  0.,  0.]])
    array_func_test(hat_2d, func_args, ret_desired)

    #Test hat_3d()
    arg1 = np.array([2.0, 1, 3, 5, 4, 2])
    func_args = (arg1,)
    ret_desired = np.array([[ 0., -2.,  4.,  2.],
                            [ 2., -0., -5.,  1.],
                            [-4.,  5.,  0.,  3.],
                            [ 0.,  0.,  0.,  0.]])
    array_func_test(hat_3d, func_args, ret_desired)

    #Test homog_2d()
    arg1 = np.array([2.0, 1, 3])
    arg2 = 0.658
    func_args = (arg1,arg2)
    ret_desired = np.array([[-0.3924, -0.9198,  0.1491],
                            [ 0.9198, -0.3924,  1.2348],
                            [ 0.    ,  0.    ,  1.    ]])
    array_func_test(homog_2d, func_args, ret_desired)

    #Test homog_3d()
    arg1 = np.array([2.0, 1, 3, 5, 4, 2])
    arg2 = 0.658
    func_args = (arg1,arg2)
    ret_desired = np.array([[ 0.4249,  0.8601, -0.2824,  1.7814],
                            [ 0.2901,  0.1661,  0.9425,  0.9643],
                            [ 0.8575, -0.4824, -0.179 ,  0.1978],
                            [ 0.    ,  0.    ,  0.    ,  1.    ]])
    array_func_test(homog_3d, func_args, ret_desired)

    #Test prod_exp()
    arg1 = np.array([[2.0, 1, 3, 5, 4, 6], [5, 3, 1, 1, 3, 2], [1, 3, 4, 5, 2, 4]]).T
    arg2 = np.array([0.658, 0.234, 1.345])
    func_args = (arg1,arg2)
    ret_desired = np.array([[ 0.4392,  0.4998,  0.7466,  7.6936],
                            [ 0.6599, -0.7434,  0.1095,  2.8849],
                            [ 0.6097,  0.4446, -0.6562,  3.3598],
                            [ 0.    ,  0.    ,  0.    ,  1.    ]])
    array_func_test(prod_exp, func_args, ret_desired)

    #Left_arm Calculation
    theta = np.array([1,2,1,2,1,2,1])
    g = left_arm(theta)
    print g
    
    print('Done!')
