#!/usr/bin/env python3

import urx
import math3d as m3d
import math
import time
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)
    rob.set_tcp((0, 0, 0.15, 0, 0, 0))
    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):
    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=5.0, vel=5.0, command="movej")  # 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




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")
    fprint("Current tool pose is: ",  rob.getl())
    #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)
    fprint("Current tool pose is: ",  rob.getl())
    sys.exit(0)
    rob.stop()