# Main script to perform calibration exercise for Galen Mk. 2

from ambf_client import Client
import csv
import pandas as pd
import numpy as np
import argparse
import time

# SIMULATION SET-UP

# Importing .csv
parser = argparse.ArgumentParser()
parser.add_argument("-t")
args = parser.parse_args()
traj_path = args.t

df_in = pd.read_csv(traj_path, header=None) # path should be given as "Trajectories/<desired csv trajectory file>"
joint_inputs = df_in.to_numpy()

N = joint_inputs.shape[0] # N - number of simulation coordinates
print(N)

# Compiling trajectory coordinates 
c0_traj = joint_inputs[:,0]
c1_traj = joint_inputs[:,1]
c2_traj = joint_inputs[:,2]
roll_traj = joint_inputs[:,3]
tilt_traj = joint_inputs[:,4]

# Connecting to AMBF client
c = Client()
c.connect()

time.sleep(1) # Wait to allow client to connect

print("Connection to client successful.")

# Setting object handles
bh = c.get_obj_handle('chassis') # bh - base handle, set to 'chassis' by default
time.sleep(2.0) # Sleep to allow handles to be set
eh = c.get_obj_handle('eelink') # eh - end effector handle, set to 'ee link' by default
time.sleep(2.0)
th = c.get_obj_handle('opticaltracker') # th - optical tracker handle
time.sleep(2.0)

# Check if bh connected - if simulation doesn't run, uncomment the following lines and check to make sure the joint names print
# print("Getting joint names")
# print(bh.get_joint_names())

# Mk.2 home position
carriage0_home = 0.0
carriage1_home = 0.0
carriage2_home = 0.0
roll_home = 0.0
tilt_home = 0.0

# Arrays to compile forward kinematics transformations and optical tracker transformations
E_fwd_kin = np.zeros((N, 6))
E_tracker = np.zeros((N, 6))

# Manually start calibration exercise
raw_input("Press <ENTER> to start calibration.")

for i in range(N):

    # Getting desired joint positions from inverseKinematics
    joint_pos_0_desired = c0_traj[i]
    joint_pos_4_desired = c1_traj[i]
    joint_pos_7_desired = c2_traj[i]
    joint_pos_10_desired = roll_traj[i]
    joint_pos_11_desired = tilt_traj[i]

    joint_pos_0_current = bh.get_joint_pos(0) # get current joint positions
    joint_pos_4_current = bh.get_joint_pos(4)
    joint_pos_7_current = bh.get_joint_pos(7)
    joint_pos_10_current = bh.get_joint_pos(10)
    joint_pos_11_current = bh.get_joint_pos(11)

    e0 = joint_pos_0_desired - joint_pos_0_current # setting joint errors
    e4 = joint_pos_4_desired - joint_pos_4_current
    e7 = joint_pos_7_desired - joint_pos_7_current
    e10 = joint_pos_10_desired - joint_pos_10_current
    e11 = joint_pos_11_desired - joint_pos_11_current

    epsilon_c = 0.0005 # delta joint error threshold
    epsilon_r = 0.001 # roll error threshold
    epsilon_t = 0.001 # tilt error threshold

    while abs(e0) > epsilon_c or abs(e4) > epsilon_c or abs(e7) > epsilon_c or abs(e10) > epsilon_r or abs(e11) > epsilon_t:

        # setting desired joint positions
        bh.set_joint_pos(0, joint_pos_0_desired)
        bh.set_joint_pos(4, joint_pos_4_desired)
        bh.set_joint_pos(7, joint_pos_7_desired)
        bh.set_joint_pos(10, joint_pos_10_desired)
        bh.set_joint_pos(11, joint_pos_11_desired)

        # allow robot to move
        time.sleep(0.001)

        joint_pos_0_current = bh.get_joint_pos(0) # get current joint positions
        joint_pos_4_current = bh.get_joint_pos(4)
        joint_pos_7_current = bh.get_joint_pos(7)
        joint_pos_10_current = bh.get_joint_pos(10)
        joint_pos_11_current = bh.get_joint_pos(11)

        e0 = joint_pos_0_desired - joint_pos_0_current # setting joint errors
        e4 = joint_pos_4_desired - joint_pos_4_current
        e7 = joint_pos_7_desired - joint_pos_7_current
        e10 = joint_pos_10_desired - joint_pos_10_current
        e11 = joint_pos_11_desired - joint_pos_11_current

    # Getting forward kinematics - Option 1 (RPY)
    ee_pose = eh.get_pose()
    E_fwd_kin[i, :] = ee_pose

    # Getting optical tracker transformation - Option 1 (RPY)
    ot_pose = th.get_pose()
    E_tracker[i, :] = ot_pose

    # Small sleep interval for sanity and printing progress
    time.sleep(0.005)
    if i % (N / 100) == 0:
        print(float(i) / float(N) * 100.0, "%")

# Writing data to .csv files - change desired .csv output file name
pd.DataFrame(E_fwd_kin).to_csv('CalibrationOutput/ForwardKinematics/test_fwd_kin.csv', header=None, index=None)
pd.DataFrame(E_tracker).to_csv('CalibrationOutput/OpticalTracker/test_optical_tracker.csv', header=None, index=None)

# SIMULATION SHUTDOWN

# Sending robot to home position
bh.set_joint_pos(0, carriage0_home)
bh.set_joint_pos(4, carriage1_home)
bh.set_joint_pos(7, carriage2_home)
bh.set_joint_pos(10, roll_home)
bh.set_joint_pos(11, tilt_home)

print("Calibration exercise finished. Sending robot to home position.")
c.clean_up()