#!/usr/bin/env python
import rospy
import tf_conversions
import tf2_ros
import copy
from math import fmod, pi, sin, cos, atan2, sqrt, pow, acos, floor
from sensor_msgs.msg import JointState
from geometry_msgs.msg import Pose, Point, Quaternion, TransformStamped

l1 = 3.0
l2 = 3.0
l3 = 2.0

# Returns a 2D list:
# [[solA], [solB]]
def solveTheta2And3(r, z):
    print('In solveTheta23')

    #******************************************************
    # theta2 + theta3
    #******************************************************
    A = -2.0*l3*r
    B = -2.0*l3*(z - l1)
    C = pow(r,2) + pow(z-l1,2) + pow(l3,2) - pow(l2,2)
    print('A: %s B: %s C: %s' % (A,B,C))
    print('-C/(sqrt(pow(A,2) + pow(B,2))): %s' % (-C/(sqrt(pow(A,2) + pow(B,2)))))

    psi = atan2(B,A)
    print('psi: %s' % psi)

    # Two pairs of values of theta2 and theta3 based on +-acos
    # Do fmod because we can end up with values outside of [-pi,pi]
    # because acos returns [-pi,pi] and psi can be [-pi,pi]
    theta23_a = fmod((acos( -C/(sqrt(pow(A,2) + pow(B,2))) ) + psi), pi)
    theta23_b = fmod((-acos( -C/(sqrt(pow(A,2) + pow(B,2))) ) + psi), pi)
    print('theta23_a: %s theta23_b: %s' % (theta23_a, theta23_b))


    #******************************************************
    # theta2
    #******************************************************
    theta2_a = atan2( z-l1 - (l3*sin(theta23_a)), r - (l3*cos(theta23_a)) )
    theta2_b = atan2( z-l1 - (l3*sin(theta23_b)), r - (l3*cos(theta23_b)) )
    print('theta2_b: "%s' % theta2_b)

    #******************************************************
    # theta3
    #******************************************************
    theta3_a = theta23_a - theta2_a
    theta3_b = theta23_b - theta2_b

    print('Exiting solveTheta23')
    return [[theta2_a, theta3_a], [theta2_b, theta3_b]]



def IK(eePose):
    print('In IK')

    x = eePose.position.x
    y = eePose.position.y
    z = eePose.position.z

    #******************************************************
    # theta1
    # theta 1 is based on an overhead view. z is a dot.
    # Two values for theta1 based on +r and -r
    #******************************************************
    r_a = sqrt( pow(x,2) + pow(y,2) )
    r_b = -r_a
    theta1_a = atan2( y/r_a, x/r_a )
    theta1_b = atan2( y/r_b, x/r_b )
    print('r_a: %s r_b: %s theta1_a: %s theta1_b: %s' % (r_a, r_b, theta1_a, theta1_b))


    # Use a function for the other theta values

    print('*************2nd SOL*****************')
    solsRA = solveTheta2And3(r_a, z)
    print('*************END 2nd SOL*****************')
    solsRB = solveTheta2And3(r_b, z)

    sol1 = [theta1_a, solsRA[0][0], solsRA[0][1]]
    sol2 = [theta1_a, solsRA[1][0], solsRA[1][1]]
    sol3 = [theta1_b, solsRB[0][0], solsRB[0][1]]
    sol4 = [theta1_b, solsRB[1][0], solsRB[1][1]]

    print('Exiting IK')
    return [sol1, sol2, sol3, sol4]