#!/usr/bin/env python

# Copyright (C) 2017 Udacity 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 acos 
from sympy import symbols, cos, sin, pi, simplify, atan2, sqrt, Matrix


def handle_calculate_IK(req):
    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:

        ### Your FK code here
        # Create symbols
	q1, q2, q3, q4, q5, q6, q7 = symbols('q1:8')
	d1, d2, d3, d4, d5, d6, d7 = symbols('d1:8')
	a0, a1, a2, a3, a4, a5, a6 = symbols('a0:7')
	alpha0, alpha1, alpha2, alpha3, alpha4, alpha5, alpha6 = symbols('alpha0:7')
	# Create Modified DH parameters
	s = {	alpha0: 0	,a0:0		,d1: 	0.75, 	
		alpha1: -pi/2	,a1:0.35	,d2:	0, 	q2: q2-pi/2,  
     		alpha2: 0	,a2:1.25	,d3:	0, 	
     		alpha3: -pi/2	,a3:-0.054	,d4:	1.5,	
     		alpha4:  pi/2	,a4:0		,d5:	0, 	
     		alpha5: -pi/2	,a5:0		,d6:	0, 	
     		alpha6: 0	,a6:0		,d7:	0.303, 	q7: 0}


	# Define Modified DH Transformation matrix
	#### Homogeneous Transforms

	def TF_MAT (alpha, a, d, q):
            TF = Matrix([[             cos(q),           -sin(q),           0,             a],
                          [ sin(q)*cos(alpha), cos(q)*cos(alpha), -sin(alpha), -sin(alpha)*d],
                          [ sin(q)*sin(alpha), cos(q)*sin(alpha),  cos(alpha),  cos(alpha)*d],
                          [                 0,                 0,           0,             1]])
            return TF

	T0_1 = TF_MAT(alpha0,a0,d1,q1).subs(s)
        T1_2 = TF_MAT(alpha1,a1,d2,q2).subs(s)
        T2_3 = TF_MAT(alpha2,a2,d3,q3).subs(s)
        T3_4 = TF_MAT(alpha3,a3,d4,q4).subs(s)
        T4_5 = TF_MAT(alpha4,a4,d5,q5).subs(s)
        T5_6 = TF_MAT(alpha5,a5,d6,q6).subs(s)
        T6_G = TF_MAT(alpha6,a6,d7,q7).subs(s)
	
	# Create individual transformation matrices
	T0_2 = T0_1 * T1_2 
	T0_3 = T0_2 * T2_3
	T0_4 = T0_3 * T3_4
	T0_5 = T0_4 * T4_5
	T0_6 = T0_5 * T5_6
	T0_G = T0_6 * T6_G
	# Extract rotation matrices from the transformation matrices
	R0_1 = T0_1[0:3,0:3]
	R0_2 = T0_2[0:3,0:3]
	R0_3 = T0_3[0:3,0:3]
	R0_4 = T0_4[0:3,0:3]
	R0_5 = T0_5[0:3,0:3]
	R0_6 = T0_6[0:3,0:3]
	R0_G = T0_G[0:3,0:3]
	#
        ###

        # Initialize service response
        joint_trajectory_list = []
        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
            px = req.poses[x].position.x
            py = req.poses[x].position.y
            pz = req.poses[x].position.z
	    # roll(r), pitch(p), yaw(y) = end-effector orientation
            (r, p, y) = 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])

	    # Compensate for rotation discrepancy between DH parameters and Gazebo

	    # to reduce calculation time R_corr = (R_z(pi)* R_y(-pi/2)) is pre-solved
	    R_corr = Matrix([[0,0,1],
                    	     [0,-1,0],
                             [1,0,0]])
	    # As for the Rrpy = R_z(yaw) * R_y(pitch) * R_x(roll) is also pre-solved
	    Rrpy = Matrix([[cos(y)*cos(p), cos(y)*sin(p)*sin(r)-sin(y)*cos(r), cos(y)*sin(p)*cos(r)+sin(y)*sin(r)],
                            [sin(y)*cos(p), sin(y)*sin(p)*sin(r)+cos(y)*cos(r), sin(y)*sin(p)*cos(r)-cos(y)*sin(r)],
                            [-sin(p),         cos(p)*sin(r),                             cos(p)*cos(r)]])
 	    #to distinguish between R0_6 drived above I used Rrpy0_6 for IK calculations
	    Rrpy = Rrpy * R_corr.T
            Rrpy0_6 = Rrpy

            nx = Rrpy0_6[0,2]
    	    ny = Rrpy0_6[1,2]
    	    nz = Rrpy0_6[2,2]


            EE = Matrix([[px],
                         [py],
                         [pz]])
		

	    # Calculate joint angles using Geometric IK method

	    #calculate the wrist center position using end effector position and correct orientation.
	    d7 = 0.303 #from DH table
	    wc = EE - d7 * Rrpy0_6[:,2]
				
	    a2 = 1.25 #from DH table
	    d4 = 1.504 #from DH table
	    d1 = 0.75 #from DH table

	    B1 = wc[2] - d1
	    B2 = sqrt(wc[0]**2+wc[1]**2)- 0.35

	    B = sqrt(B2**2+B1**2) 
	    A = d4 
	    C = a2 
		
	    a = acos((B**2+C**2 -A**2)/(2*B*C))
	    b = acos((A**2+C**2 -B**2)/(2*A*C))
	    c = acos((A**2+B**2 -C**2)/(2*A*B))

	    theta1 = atan2(wc[1],wc[0])
	    theta2 = pi/2 - a - (atan2(B1,B2))
	    theta3 = pi/2 - b - atan2(0.0054,1.5)

	    Rrpy0_3 = R0_3.evalf(subs={q1: theta1, q2: theta2, q3: theta3})
    	    R3_6 = Rrpy0_3.T * Rrpy0_6

    	    theta4 = atan2(R3_6[2,2], -R3_6[0,2])
    	    theta5 = atan2(sqrt(R3_6[0,2]**2 + R3_6[2,2]**2), R3_6[1,2])
    	    theta6 = atan2(-R3_6[1,1],R3_6[1,0])
	
	    #
            ###

            # 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)

        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()
