import urx
import time
import numpy as np

# Fixed parameters
PI = 3.1415926535
movement_plane_height = 0.4
source_coordinate = [0, 0.1, -0.035]
in_to_m = 0.0254
voxel_xy_offset = [0, -0.1, source_coordinate[2]]

# Load in xyz coordinates of voxels, organized from bottom to top layer
voxel_coordinates = [
    [6.0, 2, 4], [10, 2, 4], [14, 14, 4], [2, 14, 4], [6.0, 2, 8.0], [10, 2, 8.0],
    [14, 14, 8.0], [2, 14, 8.0], [14, 10, 8.0], [2, 10, 8.0], [14, 10, 12], [6.0, 2, 12],
    [10, 2, 12], [2, 10, 12], [14, 14, 12], [2, 14, 12], [10, 6.0, 12], [6.0, 6.0, 12],
    [2, 14, 16], [2, 10, 16], [10, 6.0, 16], [14, 10, 16], [6.0, 6.0, 16], [10, 2, 16],
    [14, 14, 16], [6.0, 2, 16], [14, 2, 16], [2, 2, 16], [14, 6.0, 16], [10, 14, 16],
    [6.0, 14, 16], [6.0, 10, 16], [2, 6.0, 16], [10, 10, 16],
]
for i in range(len(voxel_coordinates)):
    for xyz in range(3):
        if xyz == 1:
            voxel_coordinates[i][xyz] = -voxel_coordinates[i][xyz]
        elif xyz == 2:
            voxel_coordinates[i][xyz] -= 4
        # scale:
        voxel_coordinates[i][xyz] = voxel_coordinates[i][xyz] * in_to_m
        # offset
        voxel_coordinates[i][xyz] += voxel_xy_offset[xyz]
print(voxel_coordinates)

# Setup robot as needed
robot = urx.Robot("192.168.1.53")
robot.set_tcp((0, 0, 0.1, 0, 0, 0))
robot.set_payload(2, (0, 0, 0.1))
acc_cmd = 7 # acc. and vel. for all movements
vel_cmd = 7
time.sleep(0.2)  # leave some time to robot to process the setup commands

# Helper functions for convenience
def deg2rad(deg):
    rad = []
    for i in range(len(deg)):
        rad.append(deg[i] * PI / 180)
    return rad

def translateEndEffectorGlobal(robot, p_des, p_root):
    p_cmd = robot.getl()
    for i in range(len(p_des)):
        p_cmd[i] = p_des[i] + p_root[i]
    p_cmd[3:6] = p_root[3:6]
    print("Moving to global position: ", p_cmd)
    robot.movel(p_cmd, acc=acc_cmd, vel=vel_cmd)

def translateEndEffectorZGlobal(robot, p_z, p_root):
    p_cmd = robot.getl()
    p_cmd[2] = p_z + p_root[2]
    p_cmd[3:6] = p_root[3:6]
    print("Moving to global Z position: ", p_cmd[2])
    robot.movel(p_cmd, acc=acc_cmd, vel=vel_cmd)

def translateEndEffectorXYGlobal(robot, p_xy, p_root):
    p_cmd = robot.getl()
    for i in range(len(p_xy)):
        p_cmd[i] = p_xy[i] + p_root[i]
    p_cmd[3:6] = p_root[3:6]
    print("Moving to global position: ", p_xy)
    robot.movel(p_cmd, acc=acc_cmd, vel=vel_cmd)

def grabSourceBlock(robot, p_root):
    source_coordinate_xy = source_coordinate[0:2]
    source_coordinate_z = source_coordinate[2]
    translateEndEffectorZGlobal(robot, movement_plane_height, p_root)
    translateEndEffectorXYGlobal(robot, source_coordinate_xy, p_root)
    translateEndEffectorZGlobal(robot, source_coordinate_z, p_root)
    robot.send_program('set_tool_digital_out(0, False)')
    time.sleep(0.1)
    translateEndEffectorZGlobal(robot, movement_plane_height, p_root)

def placeSourceBlock(robot, p_des, p_root):
    translateEndEffectorZGlobal(robot, movement_plane_height, p_root)
    translateEndEffectorXYGlobal(robot, p_des[0:2], p_root)
    translateEndEffectorZGlobal(robot, p_des[2], p_root)
    robot.send_program('set_tool_digital_out(0, True)')
    time.sleep(0.1)
    translateEndEffectorZGlobal(robot, movement_plane_height, p_root)

if __name__ == '__main__':
    # Startup procedure
    print("Starting program...")
    j_root = deg2rad([0, -90, 135, 225, -90, 45])
    print("Robot root joint angles: ", j_root)
    print("Moving to home...")
    robot.movej(j_root, acc=acc_cmd, vel=vel_cmd)
    robot.send_program('set_tool_digital_out(0, True)')  # open tool as default
    p_root = robot.getl()  # global root position on startup
    print("Robot global root position: ", p_root)

    # Main code
    for coord in voxel_coordinates:
        grabSourceBlock(robot, p_root)
        placeSourceBlock(robot, coord, p_root)
    # translateEndEffectorXYGlobal(robot, [0, 0.1], p_root)
    # translateEndEffectorZGlobal(robot, 0.3, p_root)

    # Shutdown procedure
    input("Press enter to shutdown...")
    print("Returning to root position and orientation...")
    translateEndEffectorZGlobal(robot, movement_plane_height, p_root)
    translateEndEffectorXYGlobal(robot, [0, 0], p_root)
    robot.movej(j_root, acc=acc_cmd, vel=vel_cmd)
    robot.send_program('set_tool_digital_out(0, True)')  # open tool as default

    print("Program finished.")
    time.sleep(0.5)
    robot.close()