# navigate the robot to the target 3D coordinates

import open3d as o3d
import math3d as m3d
import URBasic
import time
from utils.pose_conversion import *
from utils.trajectory_io import *
import pickle
from utils.triangulation import *

from subject_info import SUBJECT_NAME, SCAN_POSE
import argparse

parser = argparse.ArgumentParser(description='Compute target')
parser.add_argument('--pose_model', type=str, default='ViTPose_large', help='pose model')
args = parser.parse_args()
POSE_MODEL = args.pose_model

folder_path = 'src/data/' + SUBJECT_NAME + '/' + SCAN_POSE + '/'

# read intrinsics
with open(folder_path + 'intrinsics/cam_1_intrinsics.pickle', 'rb') as f:
    cam1_intr = pickle.load(f)
    print("cam1_intr: \n", cam1_intr)

with open(folder_path + 'intrinsics/cam_2_intrinsics.pickle', 'rb') as f:
    cam2_intr = pickle.load(f)
    print("cam2_intr: \n", cam2_intr)

# read extrinsics
camera_poses = read_trajectory(folder_path + "odometry.log")

# TSDF volume
volume = o3d.pipelines.integration.ScalableTSDFVolume(
    voxel_length=1 / 512.0,
    sdf_trunc=0.04,
    color_type=o3d.pipelines.integration.TSDFVolumeColorType.RGB8)

# show RGBD images from all the views
for i in range(len(camera_poses)):
    print("Integrate {:d}-th image into the volume.".format(i))
    color = o3d.io.read_image(folder_path + "color_images/cam_{}.jpg".format(i+1))
    depth = o3d.io.read_image(folder_path + "depth_images/cam_{}.png".format(i+1))

    rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(
        color, depth, depth_trunc=4.0, convert_rgb_to_intensity=False)

    cam_intr = cam1_intr if i == 0 else cam2_intr

    intr = o3d.camera.PinholeCameraIntrinsic(
        width=640,
        height=480,
        fx=cam_intr[0,0],
        fy=cam_intr[1,1],
        cx=cam_intr[0,2],
        cy=cam_intr[1,2]
    )

    volume.integrate(rgbd,
                     intr,
                     np.linalg.inv(camera_poses[i].pose))

# point cloud generation
pcd = volume.extract_point_cloud()
downpcd = pcd.voxel_down_sample(voxel_size=0.01)

downpcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))
# o3d.visualization.draw_geometries([downpcd])
# o3d.visualization.draw_geometries([downpcd], point_show_normal=True)
coordinates = np.asarray(downpcd.points)
normals = np.asarray(downpcd.normals)

# UR3e configuration
ROBOT_IP = '169.254.147.11'  # real robot IP
ACCELERATION = 0.5  # robot acceleration
VELOCITY = 0.5  # robot speed value


robot_start_position = (np.radians(-99.73), np.radians(-171.50), np.radians(85.98),
                        np.radians(-1.81), np.radians(90.23), np.radians(221.59))   # front

# add a waypoint to avoid collision
robot_waypoint = (np.radians(98.57), np.radians(-113.38), np.radians(-38.86),
                 np.radians(-90.), np.radians(91.28), np.radians(155.32))


# initialise robot with URBasic
print("initialising robot")
robotModel = URBasic.robotModel.RobotModel()
robot = URBasic.urScriptExt.UrScriptExt(host=ROBOT_IP, robotModel=robotModel)

robot.reset_error()
print("robot initialised")
time.sleep(1)
robot.init_realtime_control()
time.sleep(1)
robot.movej(q=robot_start_position, a=ACCELERATION, v=VELOCITY)


with open(folder_path + POSE_MODEL + '/final_target.pickle', 'rb') as f:
    X_targets =  pickle.load(f)

for i, X_target in enumerate(X_targets):
    print("========================================")
    print(f"{i}th target point: {X_target.squeeze()}")

    # need to convert robot base coordinate to camera coordinate
    init_estimate = X_target.squeeze()  # an initial estimate of the target point
    idx = np.argmin(np.square(coordinates[:, 0:2] - init_estimate[0:2]).sum(axis=1))

    t_target_base = coordinates[idx]
    target_normal_base = normals[idx]

    print('closest point in base frame:', t_target_base)
    print('closest point normal in base frame:', target_normal_base)

    TCP_6d_pose_base = robot.get_actual_tcp_pose()
    print("robot pose: ", TCP_6d_pose_base)

    # Want to align the orientation of TCP to be (0,0,1)
    T_tcp_base = np.asarray(m3d.Transform(TCP_6d_pose_base).get_matrix())  # 4x4 matrix
    R_tcp_base = T_tcp_base[0:3, 0:3]
    R_base_tcp = np.linalg.inv(R_tcp_base)

    target_normal_tcp = R_base_tcp @ target_normal_base.reshape(-1,1)

    # compute the rotation from (0,0,1)/tcp's z-axis to the negative direction of the target normal in tcp frame
    R_target_tcp = rotation_align(np.array([0, 0, 1]), -target_normal_tcp.squeeze())

    R_target_base = R_tcp_base @ R_target_tcp
    T_target_base = np.vstack([np.hstack([R_target_base, t_target_base.reshape(-1, 1)]), np.array([0, 0, 0, 1])])

    target_6d_pose_base = m3d.Transform(T_target_base).get_pose_vector()
    print("Final 6d pose base: ", target_6d_pose_base)

    robot.movej(q=robot_waypoint, a=ACCELERATION, v=VELOCITY)
    robot.movej(pose=target_6d_pose_base, a=ACCELERATION, v=VELOCITY)

    time.sleep(3)
    robot.movej(q=robot_waypoint, a=ACCELERATION, v=VELOCITY)
    robot.movej(q=robot_start_position, a=ACCELERATION, v=VELOCITY)


robot.close()

exit(0)