#!/usr/bin/env python

# Copyright (C) 2017 Electric Movement Inc.
#
# This file is part of Robotic Arm: Pick and Place project for Udacity
# Robotics nano-degree program
#
# All Rights Reserved.

# Author: Harsh Pandya

# import modules
import rospy
import tf
from kuka_arm.srv import *
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint
from geometry_msgs.msg import Pose
#from mpmath import *
#from sympy import *
import numpy as np

# Define functions for Rotation Matrices about x, y, and z given the angle of rotation
# From lesson 2 part 11
def rot_x(q):
    R_x = np.float64([[1,      0,       0], 
                  [0, np.cos(q), -np.sin(q)],
                  [0, np.sin(q),  np.cos(q)]])
    
    return R_x


def rot_y(q):              
    R_y = np.float64([[ np.cos(q), 0, np.sin(q)],
                  [      0, 1,      0],
                  [-np.sin(q), 0, np.cos(q)]])
    
    return R_y


def rot_z(q):    
    R_z = np.float64([[np.cos(q), -np.sin(q), 0], 
                  [np.sin(q),  np.cos(q), 0],
                  [     0,       0, 1]])
    
    return R_z



def handle_calculate_IK(req, test = 'no'):
    rospy.loginfo("Received %s eef-poses from the plan" % len(req.poses))
    if len(req.poses) < 1:
        print "No valid poses received"
        return -1
    else:
        # Initialize service response
        joint_trajectory_list = []


	alpha0, alpha1, alpha2, alpha3, alpha4, alpha5, alpha6 = 0, -np.pi/2, 0, -np.pi/2, np.pi/2, -np.pi/2, 0
	a0, a1, a2, a3, a4, a5, a6 = 0, 0.35, 1.25, -0.054, 0, 0, 0
	d1, d2, d3, d4, d5, d6, d7 = 0.75, 0,0,1.5, 0,0,0.303
	q1, q2, q3, q4, q5, q6, q7 = 0, -np.pi/2, 0,0,0,0,0

        T0_1 = np.float64([[            np.cos(q1),            -np.sin(q1),            0,              a0],
                       [np.sin(q1)*np.cos(alpha0), np.cos(q1)*np.cos(alpha0), -np.sin(alpha0), -np.sin(alpha0)*d1],
                       [np.sin(q1)*np.sin(alpha0), np.cos(q1)*np.sin(alpha0),  np.cos(alpha0),  np.cos(alpha0)*d1],
                       [                  0,                   0,            0,               1]])


        T1_2 = np.float64([[            np.cos(q2),            -np.sin(q2),            0,              a1],
                       [np.sin(q2)*np.cos(alpha1), np.cos(q2)*np.cos(alpha1), -np.sin(alpha1), -np.sin(alpha1)*d2],
                       [np.sin(q2)*np.sin(alpha1), np.cos(q2)*np.sin(alpha1),  np.cos(alpha1),  np.cos(alpha1)*d2],
                       [                  0,                   0,            0,               1]])

        T2_3 = np.float64([[            np.cos(q3),            -np.sin(q3),            0,              a2],
                       [np.sin(q3)*np.cos(alpha2), np.cos(q3)*np.cos(alpha2), -np.sin(alpha2), -np.sin(alpha2)*d3],
                       [np.sin(q3)*np.sin(alpha2), np.cos(q3)*np.sin(alpha2),  np.cos(alpha2),  np.cos(alpha2)*d3],
                       [                  0,                   0,            0,               1]])

        T3_4 = np.float64([[            np.cos(q4),            -np.sin(q4),            0,              a3],
                       [np.sin(q4)*np.cos(alpha3), np.cos(q4)*np.cos(alpha3), -np.sin(alpha3), -np.sin(alpha3)*d4],
                       [np.sin(q4)*np.sin(alpha3), np.cos(q4)*np.sin(alpha3),  np.cos(alpha3),  np.cos(alpha3)*d4],
                       [                  0,                   0,            0,               1]])

        T4_5 = np.float64([[            np.cos(q5),            -np.sin(q5),            0,              a4],
                       [np.sin(q5)*np.cos(alpha4), np.cos(q5)*np.cos(alpha4), -np.sin(alpha4), -np.sin(alpha4)*d5],
                       [np.sin(q5)*np.sin(alpha4), np.cos(q5)*np.sin(alpha4),  np.cos(alpha4),  np.cos(alpha4)*d5],
                       [                  0,                   0,            0,               1]])

        T5_6 = np.float64([[            np.cos(q6),            -np.sin(q6),            0,              a5],
                       [np.sin(q6)*np.cos(alpha5), np.cos(q6)*np.cos(alpha5), -np.sin(alpha5), -np.sin(alpha5)*d6],
                       [np.sin(q6)*np.sin(alpha5), np.cos(q6)*np.sin(alpha5),  np.cos(alpha5),  np.cos(alpha5)*d6],
                       [                  0,                   0,            0,               1]])            


        T6_G = np.float64([[            np.cos(q7),            -np.sin(q7),            0,              a6],
                       [np.sin(q7)*np.cos(alpha6), np.cos(q7)*np.cos(alpha6), -np.sin(alpha6), -np.sin(alpha6)*d7],
                       [np.sin(q7)*np.sin(alpha6), np.cos(q7)*np.sin(alpha6),  np.cos(alpha6),  np.cos(alpha6)*d7],
                       [                  0,                   0,            0,               1]])


        T0_2 = np.matmul(T0_1, T1_2)
        T0_3 = np.matmul(T0_2, T2_3)
        T0_4 = np.matmul(T0_3, T3_4)
        T0_5 = np.matmul(T0_4, T4_5)
        T0_6 = np.matmul(T0_5, T5_6)
        T0_G = np.matmul(T0_6, T6_G)

        # Correct orientation of gripper link in URDF vs. DH convention
        # 180 degree rotation about the z-axis
        R_z = np.float64([[np.cos(np.pi), -np.sin(np.pi), 0, 0],
                      [np.sin(np.pi),  np.cos(np.pi), 0, 0],
                      [      0,        0, 1, 0],
                      [      0,        0, 0, 1]])
        # 90 degree rotation about the y-axis
        R_y = np.float64([[ np.cos(-np.pi/2), 0, np.sin(-np.pi/2), 0],
                      [          0, 1,          0, 0],
                      [-np.sin(-np.pi/2), 0, np.cos(-np.pi/2), 0],
                      [          0, 0,          0, 1]])
#         R_corr = simplify(R_z * R_y)
        R_corr = np.matmul(R_z, R_y)
        # Total transform with gripper orientation corrected
#         T_total = simplify(T0_G * R_corr)
        T_total = np.matmul(T0_G, R_corr) 


        for x in xrange(0, len(req.poses)):
            # IK code starts here
            joint_trajectory_point = JointTrajectoryPoint()

            # Extract end-effector position and orientation from request
	    # px,py,pz = end-effector position
	    # roll, pitch, yaw = end-effector orientation
            px = req.poses[x].position.x
            py = req.poses[x].position.y
            pz = req.poses[x].position.z

            (roll, pitch, yaw) = tf.transformations.euler_from_quaternion(
                [req.poses[x].orientation.x, req.poses[x].orientation.y,
                    req.poses[x].orientation.z, req.poses[x].orientation.w])
     
            # Calculate joint angles using Geometric IK method

            
            ee_length = d7
            #ee_length = ee_length.subs(s)
            # a = l2
            l2_3 = a2
            #l2_3 = l2_3.subs(s)
            # b = l3 + l4 w/ adjustment
            # labeled as d in writeup
            l3_4 = 0.96 # from URDF file
            # labeled as e in writeup
            l4_5 = 0.54 # from URDF file
            # labeled as a3 in writeup
            l3_4_offset = np.absolute(a3)
            #l3_4_offset = l3_4_offset.subs(s)

            # labeled as B in writeup
            l3_4_angle = np.pi - np.arcsin(l3_4_offset / l3_4)
            # Cosine rule - labeled as b in writeup
            # b = sqrt(d^2 + e^2 - 2de * cos(B)
            dist3_5 = np.sqrt(l3_4**2 + l4_5**2 - 2*l3_4*l4_5*np.cos(l3_4_angle))
            
            # 1. Find total rotation matrix from roll-pitch-yaw data
            #R_total = simplify(rot_z(yaw) * rot_y(pitch) * rot_x(roll))
            R_total = np.matmul(rot_z(yaw), rot_y(pitch), rot_x(roll))
            # 2. Find wrist center position using the end effector position and orientation
                # d6 = 0
                # wx = px - (d6 + l) * nx
                # wy = py - (d6 + l) * ny
                # wz = pz - (d6 + l) * nz

            ee_position = np.float64([[px],
                                  [py],
                                  [pz]])

            ee_adj = np.float64([[ee_length],
                             [0],
                             [0]])

            w_c = ee_position - R_total * ee_adj 

            # 3. theta1 calc
            theta1 = np.arctan2(w_c[1,0], w_c[0,0])

#             # 4. theta2 calc
#             J2_x = a1 * cos(theta1)
#             J2_x = J2_x.subs(s)
#             J2_y = a1 * sin(theta1)
#             J2_y = J2_y.subs(s)
#             J2_z = d1
#             J2_z = J2_z.subs(s)
            
            J2_x = a1 * np.cos(theta1)
            #J2_x = J2_x.subs(s)
            J2_y = a1 * np.sin(theta1)
            #J2_y = J2_y.subs(s)
            J2_z = d1
            #J2_z = J2_z.subs(s)
            
            
            J5_x = w_c[0,0]
            J5_y = w_c[1,0]
            J5_z = w_c[2,0]
            # labeled g in writeup
            J5_2_z = J5_z - J2_z

            # labeled c in writeup
            dist_J2_J5 = np.sqrt((J5_x - J2_x)**2 + (J5_y - J2_y)**2 + (J5_z - J2_z)**2)
            # labeled f in writeup
            dist_J2_J5_xy = np.sqrt((J5_x - J2_x)**2 + (J5_y - J2_y)**2)

            # theta2 = pi/2 - beta - delta
            theta2 = np.pi/2 - (np.arccos((dist3_5**2 - l2_3**2 - dist_J2_J5**2)/(-2*l2_3*dist_J2_J5))) - np.arctan2(J5_2_z,dist_J2_J5_xy)
           
            # 5. theta3 calc
            # theta3 = pi/2 - delta - gamma
            # delta = asin(a3/b)
            # gamma = acos((c^2 - a^2 - b^2)/(-2ab))
            theta3 = np.pi/2 - np.arcsin(l3_4_offset/dist3_5) - np.arccos((dist_J2_J5**2 - l2_3**2 - dist3_5**2)/(-2*l2_3*dist3_5))
		
            # 6. Find R3_6 from orientation data

            # R_rpy = R_total

            q1,q2,q3 = theta1,theta2,theta3
            T0_1 = np.float64([[            np.cos(q1),            -np.sin(q1),            0,              a0],
                           [np.sin(q1)*np.cos(alpha0), np.cos(q1)*np.cos(alpha0), -np.sin(alpha0), -np.sin(alpha0)*d1],
                           [np.sin(q1)*np.sin(alpha0), np.cos(q1)*np.sin(alpha0),  np.cos(alpha0),  np.cos(alpha0)*d1],
                           [                  0,                   0,            0,               1]])

            T1_2 = np.float64([[            np.cos(q2),            -np.sin(q2),            0,              a1],
                           [np.sin(q2)*np.cos(alpha1), np.cos(q2)*np.cos(alpha1), -np.sin(alpha1), -np.sin(alpha1)*d2],
                           [np.sin(q2)*np.sin(alpha1), np.cos(q2)*np.sin(alpha1),  np.cos(alpha1),  np.cos(alpha1)*d2],
                           [                  0,                   0,            0,               1]])
            T2_3 = np.float64([[            np.cos(q3),            -np.sin(q3),            0,              a2],
                       [np.sin(q3)*np.cos(alpha2), np.cos(q3)*np.cos(alpha2), -np.sin(alpha2), -np.sin(alpha2)*d3],
                       [np.sin(q3)*np.sin(alpha2), np.cos(q3)*np.sin(alpha2),  np.cos(alpha2),  np.cos(alpha2)*d3],
                       [                  0,                   0,            0,               1]])
            T0_2 = np.matmul(T0_1, T1_2)
            T0_3 = np.matmul(T0_2, T2_3)
            
            R0_3 = np.float64([[T0_3[0,0], T0_3[0,1], T0_3[0,2]],
                           [T0_3[1,0], T0_3[1,1], T0_3[1,2]],
                           [T0_3[2,0], T0_3[2,1], T0_3[2,2]]])
            
            # Correct orientation between DH convention and URDF
            #R_total_adjust = simplify(R_total * rot_z(-pi/2) * rot_y(-pi/2))
            R_total_adjust = np.matmul(R_total, rot_z(-np.pi/2), rot_y(-np.pi/2))
            # for a valid rotation matrix the transpose is == to the inverse 
            #R3_6 = simplify(R0_3.T * R_total_adjust)
            R3_6 = np.matmul(np.transpose(R0_3), R_total_adjust)

            # 7. Find alpha, beta, gamma euler angles as done in lesson 2 part 8.

            # Method using euler_from_matrix assuming a yzy rotation rather than an xyz rotation
            alpha, beta, gamma = tf.transformations.euler_from_matrix(R3_6.tolist(), 'ryzy')
            theta4 = alpha
            theta5 = beta
            theta6 = gamma
            
            # Populate response for the IK request
            # In the next line replace theta1,theta2...,theta6 by your joint angle variables
	    joint_trajectory_point.positions = [theta1, theta2, theta3, theta4, theta5, theta6]
	    joint_trajectory_list.append(joint_trajectory_point)

        # Assemble variables if testing
        if test == 'yes':
            test_variables = {'px' : px, 'py' : py, 'pz' : pz, 'roll' : roll, 'pitch' : pitch, 'yaw' : yaw,
                              'wcx' : wx, 'wcy' : wy, 'wcz' : wz, 'theta1' : theta1, 'theta2' : theta2, 
                              'theta3' : theta3, 'theta4' : theta4, 'theta5' : theta5, 'theta6' : theta6 } 

            return test_variables
        else:
            rospy.loginfo("length of Joint Trajectory List: %s" % len(joint_trajectory_list))
            return CalculateIKResponse(joint_trajectory_list)


def IK_server():
    # initialize node and declare calculate_ik service
    rospy.init_node('IK_server')
    s = rospy.Service('calculate_ik', CalculateIK, handle_calculate_IK)
    print "Ready to receive an IK request"
    rospy.spin()

if __name__ == "__main__":
    IK_server()
