import numpy as np
import pickle
import time
import copy
import sys

# ROS Libraries
import rospy
from std_msgs.msg import Bool
from trajectory_msgs.msg import (
    JointTrajectoryPoint,
    JointTrajectory
)

## Our libraries
from utils import load_plan, JTraj2ROS, speed_multiplier
from Controller import Controller
import gcode2contour as gc


def execute_plan(plan,
        r1_con,
        r2_con,
        contours=None,
        contour_speed=0.2,
        move_speed=1,
        retry_time=0.1,
        max_traj_retries=20):
    """
    Executes a plan on MoveIt

    Single arm only
    skips the first n trajectories in the plan
    """
    # plan to the start position of the plan
    r1_con.group.go(plan.trajs[0][0].positions[0])
    r2_con.group.go(plan.trajs[1][0].positions[0])
    controllers = [r1_con, r2_con]

    # execute trajectories in plan
    plan_len = max([a[-1] for a in plan.cumul_time.values()])/move_speed

    s = time.time()
    expected_end_time = s + plan_len
    delays = [0,0] # measure delays of each arm
    relative_delay = 0 # how far arm 0 is behind arm 2
    i_r = [0,0] # index of which traj we're up to for each arm
    arm_complete= None
    while time.time() < expected_end_time + max(delays):

        print("----------")

        # Find which arm and which trajectory is next to be executed
        if arm_complete is not None:
            arm = 1 - arm_complete # if one arm complete, only need to do other
        else:
            arm = np.argmin([plan.cumul_time[0][i_r[0]], plan.cumul_time[1][i_r[1]]])
        next_time = plan.cumul_time[arm][i_r[arm]]/move_speed
        i_r[arm] += 1

        # if the last trajectory for this arm has already been sent
        if i_r[arm] == len(plan.trajs[arm]):
            print("Last Trajectory for this arm has been sent")
            if arm_complete is None:
                arm_complete = arm # mark this arm as complete
                continue
            else:
                print("All arms done (unless Logic error), leaving main loop")
                break

        t = plan.trajs[arm][i_r[arm]]

        # wait till it's time to execute the next trajectory
        rel_delay_arm = max(0, -relative_delay * (-1)**arm)
        sleep_time = next_time - (time.time() - s ) + rel_delay_arm
        if sleep_time > 0:
            print("Sleeping for {}s until time for next traj".format(
                sleep_time))
            time.sleep(sleep_time)

        # Execute next trajectory
        print("Printing {}th trajectory of arm {}, which is contour {}".format(
            i_r[arm], arm, t.contour
            ))
        retries = 0
        if t.contour is not None and contours is not None:
            traj2send = speed_multiplier(t, 1/contour_speed)
            while not controllers[arm].exec_ctraj(
                    traj2send,
                    contour=contours[t.contour],
                    wait=False) and retries < max_traj_retries:
                retries += 1
                time.sleep(retry_time)
        else:
            traj2send = speed_multiplier(t, 1/move_speed)
            while not controllers[arm].exec_ctraj(traj2send, wait=False) \
                    and retries < max_traj_retries:
                retries += 1
                time.sleep(retry_time)

        if retries == max_traj_retries:
            print("---\nFailed after {} retries, quitting".format(max_traj_retries))
            sys.exit(0)

        # Traj started now, measure delay
        delays[arm] = (time.time() - s ) - next_time
        print("Arm {} is behind by {}s when executing {}th Traj".format(
            arm,
            delays[arm],
            i_r[arm]))
        relative_delay = delays[0]-delays[1]
        print("RELATIVE DELAY: arm 0 is {}s behind".format(relative_delay))

    return


if __name__ == '__main__':
    plan2arms = load_plan('two_arm_test2')

    r1_con = Controller(disable_extruder=True, robot=1)
    r2_con = Controller(disable_extruder=True, robot=2)

    execute_plan(plan2arms, r1_con, r2_con, move_speed=0.5)