import urx
import math3d as m3d
from scipy.optimize import fsolve
import math
import numpy as np
import time
import os
import logging
from urx.robotiq_two_finger_gripper import Robotiq_Two_Finger_Gripper
import sys
from util import fprint




rob = None


def init(ip):
    global rob
    #sys.stdout = Logger()
    fprint("Starting UR5 power up...")
    
    # power up robot here

    # wait for power up (this function runs async)


    # trigger auto-initialize

    # wait for auto-initialize

    # init urx
    fprint("Connecting to arm at " + ip)
    trying = True
    while trying:
        try:
            rob = urx.Robot(ip)
            trying = False
        except:
            time.sleep(1)
    robotiqgrip = Robotiq_Two_Finger_Gripper(rob)

    # Sets robot arm endpoint offset (x,y,z,rx,ry,rz)
    rob.set_tcp((0, 0, 0.15, 0, 0, 0))

    # Set weight
    rob.set_payload(2, (0, 0, 0.1))
    #rob.set_payload(2, (0, 0, 0.1))
    time.sleep(0.2)
    fprint("UR5 ready.")

def set_pos_abs(x, y, z, xb, yb, zb, threshold=None):
    global rob
    new_orientation = m3d.Transform()
    new_orientation.orient.rotate_xb(xb)  # Replace rx with the desired rotation around X-axis
    new_orientation.orient.rotate_yb(yb)  # Replace ry with the desired rotation around Y-axis
    new_orientation.orient.rotate_zb(zb)  # Replace rz with the desired rotation around Z-axis

    # Get the current pose
    trans = rob.getl()
    # Apply the new orientation while keeping the current position
    new_trans = m3d.Transform(new_orientation.orient, m3d.Vector(trans[0:3]))
    new_trans.pos.x = x
    new_trans.pos.y = y
    new_trans.pos.z = z
    #rob.speedj(0.2, 0.5, 99999)
    rob.set_pose(new_trans, acc=2, vel=2, command="movej", threshold=threshold)  # apply the new pose

def set_pos_rel_rot_abs(x, y, z, xb, yb, zb):
    global rob
    new_orientation = m3d.Transform()
    new_orientation.orient.rotate_xb(xb)  # Replace rx with the desired rotation around X-axis
    new_orientation.orient.rotate_yb(yb)  # Replace ry with the desired rotation around Y-axis
    new_orientation.orient.rotate_zb(zb)  # Replace rz with the desired rotation around Z-axis

    # Get the current pose
    trans = rob.getl()
    
    # Apply the new orientation while keeping the current position
    new_trans = m3d.Transform(new_orientation.orient, m3d.Vector(trans[0:3]))
    new_trans.pos.x += x
    new_trans.pos.y += y
    new_trans.pos.z += z
    #rob.speedj(0.2, 0.5, 99999)
    rob.set_pose(new_trans, acc=0.1, vel=0.4, command="movej")  # apply the new pose

def set_pos_abs_rot_rel(x, y, z, xb, yb, zb):
    global rob
    new_orientation = m3d.Transform()
    new_orientation.orient.rotate_xb(xb)  # Replace rx with the desired rotation around X-axis
    new_orientation.orient.rotate_yb(yb)  # Replace ry with the desired rotation around Y-axis
    new_orientation.orient.rotate_zb(zb)  # Replace rz with the desired rotation around Z-axis

    # Get the current pose
    trans = rob.getl()
    
    # Apply the new orientation while keeping the current position
    new_trans = m3d.Transform(new_orientation.orient, m3d.Vector(trans[0:3]))
    new_trans.pos.x = x
    new_trans.pos.y = y
    new_trans.pos.z = z
    #rob.speedj(0.2, 0.5, 99999)
    rob.set_pose(new_trans, acc=0.1, vel=0.4, command="movej")  # apply the new pose


def is_safe_move(start_pose, end_pose, r=0.25):
    start_x, start_y = (start_pose[0], start_pose[1])
    end_x, end_y = (end_pose[0], end_pose[1])

    try:
        m = (end_y-start_y)/(end_x-start_x)
        b = start_y - m*start_x
        # print('m = y/x =', m)
        # print('b =', b)
    except:
        m = (end_x-start_x)/(end_y-start_y)
        b = start_x - m*start_y
        # print('m = x/y =', m)
        # print('b =', b)
    
    return r**2 - b**2 + m**2 * r**2 < 0

def cartesian_to_polar(x, y):
    r = np.sqrt(x**2 + y**2)
    theta = np.arctan2(y, x)
    return r, theta

def polar_to_cartesian(r, theta):
    x = r * np.cos(theta)
    y = r * np.sin(theta)
    return x, y


def move_to_polar(start_pos, end_pos):
    global rob

    # Convert to polar coordinates
    start_r, start_theta = cartesian_to_polar(start_pos[0], start_pos[1])
    end_r, end_theta = cartesian_to_polar(end_pos[0], end_pos[1])

    # Interpolate for xy (spiral arc)
    n_points = 30
    r_intermediate = np.linspace(start_r, end_r, n_points)
    theta_intermediate = np.linspace(start_theta, end_theta, n_points)

    # Interpolate for z (height)
    start_z = start_pos[2]
    end_z = end_pos[2]

    z_intermediate = np.linspace(start_z, end_z, n_points)

    # Interpolate for rz (keep tool rotation fixed relative to robot)


    curr_rot = rob.getl()

    theta_delta = theta_intermediate[1]-theta_intermediate[0]
    
    rx_intermediate = [curr_rot[5] + theta_delta*i for i in range(n_points)]

    # curr_rot = rob.getj()
    # start_rz = curr_rot[5]
    # rot = end_theta - start_theta
    # end_base_joint = curr_rot[0]-start_theta + rot

    # end_rz = curr_rot[0] + rot
    # # rob.movel([*polar_to_cartesian(end_r, end_theta), *rob.getl()[2:]], acc=2, vel=2)

    # print('start_theta = ', math.degrees(start_theta))
    # print('end_theta = ', math.degrees(curr_rot[0]-start_theta+rot))
    # print('start_rz =', math.degrees(start_rz))
    # print('rot =', math.degrees(rot))
    # print('end_rz =', math.degrees(end_rz))

    # rz_intermediate = np.linspace(start_rz, end_rz, n_points)

    # Convert back to cartesian coordinates
    curr_pos = rob.getl()

    intermediate_points = [[*polar_to_cartesian(r, theta), z, *curr_pos[3:]]
                            for r, theta, z, rx in zip(r_intermediate, 
                                                       theta_intermediate, 
                                                       z_intermediate, 
                                                       rx_intermediate)]

    # Move robot
    rob.movels(intermediate_points, acc=2, vel=2, radius=0.1)
    

    return rx_intermediate

def move_to_home():
    global rob

    # Home position in degrees
    home_pos = [0.10421807948612624, 
                -2.206111555015423, 
                1.710679229503537, 
                -1.075834511928354, 
                -1.569301366430687, 
                1.675098295930943]

    # Move robot
    rob.movej(home_pos, acc=2, vel=2)

def normalize_degree(theta):
    # Normalizes degree theta from -1.5pi to 1.5pi
    multiplier = 1
    normalized_theta = theta % (math.pi * multiplier)
    
    # Maintain the negative sign if the original angle is negative
    if theta < 0:
        normalized_theta -= math.pi * multiplier

    # Return angle
    return normalized_theta


def get_joints_from_xyz_rel(x, y, z, rx=0, ry=-math.pi/2, rz=0, initial_guess = (math.pi/2, math.pi/2, 0)):
    # Get limbs and offsets
    offset_x, offset_y, offset_z = (0, 0, 0.14)     # Tool offset
    l_bs, l1, l2, l3, l_wt = (0.1333, .425, .39225, .1267, .0997)    # Limb lengths
    #l3=0.15
    
    # Calculate base angle and r relative to shoulder joint
    def calculate_theta(x, y, a):
        # Calculate if we need the + or - in our equations
        if (x>-a and y>=0) or (x>a and y<0):
            flip = 1
        elif (x<-a and y>=0) or (x<a and y<0):
            flip = -1
        else: 
            # Critical section (x=a, or x=-a). Infinite slope
            # Return 0 or 180 depending on sign
            return math.atan2(y, 0)
        
        # Calculate tangent line y = mx + b
        m = (x*y - math.sqrt(x*x*y*y-(x*x-a*a)*(y*y-a*a)))/(x*x-a*a)
        b = flip * a * math.sqrt(1+m*m)

        # Calculate equivalent tangent point on circle
        cx = (-flip*m*b)/(1+m*m)
        cy = m*cx + flip*b

        # Calculate base angle, make angle negative if flip=1
        theta = math.atan2(cy, cx) + (-math.pi if flip==1 else 0)

        return theta 
    
    base_theta = calculate_theta(x, y, l_bs)
    cx, cy = l_bs*math.cos(base_theta), l_bs*math.sin(l_bs)
    r = math.sqrt((x-cx)**2 + (y-cy)**2) 


    # Formulas to find out joint positions for (r, z)
    def inv_kin_r_z(p):
        a, b, c = p            

        return (l1*math.cos(a) + l2*math.cos(a-b) + l3*math.cos(a-b-c) - r, # r
                l1*math.sin(a) + l2*math.sin(a-b) - l3*math.sin(a-b-c) - (l3*math.sin(a-b-c)) - (z + offset_z),  # z
                a-b-c) # wrist angle


    # Normalize angles
    base, shoulder, elbow, wrist1 = [normalize_degree(deg) for deg in [base_theta, *fsolve(inv_kin_r_z, initial_guess)]]

    # Return result
    return base, shoulder, elbow, wrist1, ry, rz

def get_joints_from_xyz_abs(x, y, z, rx=0, ry=-math.pi/2, rz=math.pi/2):
    joints = get_joints_from_xyz_rel(x, y, z, rx, ry, rz)

    # Joint offsets
    # Base, Shoulder, Elbow, Wrist
    inverse = [1, -1, 1, 1, 1, 1]
    offsets = [-math.pi/2, 0, 0, -math.pi/2, 0, 0]

    # Return adjusted joint positions
    return [o+j*i for j, o, i in zip(joints, offsets, inverse)]

if __name__ == "__main__":
    
    #rob.movej((0, 0, 0, 0, 0, 0), 0.1, 0.2)
    #rob.movel((x, y, z, rx, ry, rz), a, v)
    init("192.168.1.145")
    print("Current tool pose is: ",  rob.getl())
    move_to_home()

    home_pose = [-0.4999999077032916, 
               -0.2000072960336574, 
               0.40002172976662786, 
               0, 
               -3.14152741295329, 
               math.radians(62)]

    # time.sleep(.5)

    p1 = [0,

        0.6,
        .4,
        0.2226,
        3.1126,
        0.0510]
    
    p2 = [0.171,
        -0.115,
        0.2,
        0.2226,
        3.1126,
        0.0510]

    curr_pos = rob.getl()
    # up/down, 
    # tool rotation
    # tool angle (shouldn't need)
    # rob.set_pos(p1[0:3], acc=0.5, vel=0.5)

    # set_pos_abs(*home_pose)


    
    angles = get_joints_from_xyz_abs(-0.7, 0, 0)
    rob.movej(angles, acc=2, vel=2)

    # joints = []
    # for i in np.linspace(-0.3, -0.7, 50):
    #     joints.append(get_joints_from_xyz_abs(i, 0, 0))
    # rob.movejs(joints, acc=2, vel=2)

    # time.sleep(5)


    # angles = get_joints_from_xyz_abs(0, -0.6, 0)
    # rob.movej(angles, acc=2, vel=2)
    


    # set_pos_abs(*p1)
    # move = move_to_polar(p1, p2)


    # for p in move:
    #     print(math.degrees(p))
    # print("Safe? :", is_safe_move(p1, p2))


    # #set_pos_rel_rot_abs(0, 0, -0.2, math.pi, 0, -math.pi)
    # set_pos_abs(0.3, -0.2, 0.5, math.pi, 0, -math.pi)
    # set_pos_abs(0, 0.2, 0.6, math.pi, 0, -math.pi)
    # set_pos_abs(-0.5, -0.2, 0.4, math.pi, 0, -math.pi)
    # #set_pos_rel_rot_abs(0, 0, 0, math.pi, 0, -math.pi)

    # print("Current tool pose is: ",  rob.getl())
    # print("getj(): ",  rob.getj())

    # move_to_home()

    rob.stop()
    os.kill(os.getpid(), 9) # dirty kill of self
    sys.exit(0)