import gcode2contour as gc
from pydecmcts import Tree
import utils as tu
from collections import deque

import matplotlib.pyplot as plt
from math import pi
import time
import numpy as np
import pickle
import copy
import warnings
import random


home = [0, -pi/2, pi/4, -pi/4, -pi/2, -pi/4]
SAFE_TIME = 2
J2TIME_C = 2.7


def avail_contours(data, state, robot_id):
    done_contours = [ct.contour for ct in state]
    return [ct for ct in data['actions'][robot_id] \
            if ct.contour not in done_contours]


def state_storer(data, parent_state, action, robot_id):
    """
    state at a node is just a list of ctrajs planned so far
    """
    # root node catch
    if parent_state is None:
        return deque([])
    new_state = copy.copy(parent_state)
    new_state.append(action)
    return new_state


def joint_diff(js1,js2):
    return sum([(j1-j2)**2 for j1,j2 in zip(js1,js2)])


def sim_select(data, options, sim_state, method='greedy'):
    """
    How to select option during simulation
    """
    if method == 'greedy':
        j  = sim_state[-1].positions[-1]
        costs = [joint_diff(j,o.positions[0]) for o in options]
        return options[np.argmin(costs)]
    elif method == 'random':
        return random.choice(options)
    elif method == 'weighted':
        j  = sim_state[-1].positions[-1]
        weights = [1/joint_diff(j,o.positions[0]) for o in options]
        return random.choices(options, weights=weights)[0]
    else:
        print("ERROR: Invalid sim_select method: '{}'".format(method))


class CombinedSchedule:
    """
    combine decmcts rollouts
    """

    def __init__(self, data, n_robots, joint2time_mul=J2TIME_C):
        self.data = data
        self.col_table = data["collisions"]
        self.n = n_robots
        self.trajs = [[] for i in range(self.n)]
        self.start_times = [[0] for i in range(self.n)]
        self.end_times = [[0] for i in range(self.n)]
        self.t_mul = joint2time_mul
        self.safe_time = SAFE_TIME # safety buffer for Collisions
        self.start_next_contour = [0]*self.n


    def planLenOrder(self):
        """
        orders robots by shortest plan first
        """
        plan_times = [self.end_times[r][-1] for r in range(self.n)]
        return np.argsort(plan_times)


    def lastJointPos(self, robot):
        """
        last joint position of a robot
        """
        if len(self.trajs[robot])==0:
            return self.data["env"].home
        return self.trajs[robot][-1].positions[-1]


    def len(self):
        """
        total time of plan is max time of robots
        """
        return max([self.end_times[r][-1] for r in range(self.n)])


    def append(self, ct, robot, force=False):
        """
        add a contour trajectory to the CombinedSchedule
        """

        # catch for n = 1
        if self.n == 1:
            self.trajs[robot].append(ct)
            self.start_times[robot].append(self.end_times[robot][-1])
            self.end_times[robot].append(self.start_times[robot][-1] + ct.time[-1])
            return True

        # travel time based on joint_diff
        trav_time = self.t_mul * joint_diff(ct.positions[-1], self.lastJointPos(robot))
        t_0 = self.start_next_contour[robot] + trav_time

        # for every other robot, get collision free time offsets
        allowed_offsets = []
        for r in range(self.n):
            if r == robot:
                continue

            t_diff = t_0 - self.start_times[r][-1]

            # if no trajectories in other robot, can start immidiately
            if len(self.trajs[r])==0:
                ao = [] # empty allowed because (0 to 1mils will be added below)
            else:

                # get last trajectory on other robot
                ct_r = self.trajs[r][-1]

                # got collision free offsets for that trajectory
                if (ct, ct_r) in self.col_table.keys():
                    ao = copy.deepcopy(self.col_table[ct, ct_r])
                elif (ct_r, ct) in self.col_table.keys():
                    ao = [(-b,-a) for a,b in
                            reversed(self.col_table[ct_r, ct])]
                else:
                    self.col_table[ct, ct_r] = \
                            tu.col_free_offsets(ct, ct_r, self.data["env"], safe_time=self.safe_time)
                    ao = copy.deepcopy(self.col_table[ct, ct_r])

            # apply t_diff and add allowed offset for when other robot is done
            allowed_offsets.append([(-1e6, t_0 - self.end_times[r][-1]-self.safe_time)] + \
                [(a+t_diff, b+t_diff) for a,b in ao if a+t_diff<=0])

        # easy for n = 2
        if self.n == 2:
            t_wait = max(trav_time, -min(0, allowed_offsets[0][-1][1]))

            #if this contour starts after other arm is finished
            # and it isnt the first contour of the plan
            # and it isn't being added regardless
            if t_0 + t_wait > self.end_times[1-robot][-1] \
                    and not (self.end_times[robot][-1] == 0 and self.end_times[1-robot][-1] == 0) \
                    and not force:
                # reject contour if it cannot be parallelised with current plan
                self.start_next_contour[robot] = self.end_times[1-robot][-1] + self.safe_time
                return False
            self.trajs[robot].append(ct)
            self.start_times[robot].append(t_0 + t_wait)
            self.end_times[robot].append(self.start_times[robot][-1] + ct.time[-1])
            self.start_next_contour[robot] = self.end_times[robot][-1]
            return True

        # for n > 2
        t_wait_last = -1
        while True:
            t_wait = max(trav_time,-min([0] + [o[-1][1] for o in allowed_offsets]))
            if t_wait_last == t_wait:
                Exception("Logic error, this should never happen")
            for ro in allowed_offsets:
                while ro[-1][0] > -t_wait:
                    ro.pop()
            valid = [ro[-1][0] <= -t_wait <= ro[-1][1] for ro in allowed_offsets]
            if all(valid):
                break
            t_wait_last = t_wait

        self.trajs[robot].append(ct)
        self.start_times[robot].append(self.end_times[robot][-1] + t_wait)
        self.end_times[robot].append(self.start_times[robot][-1] + ct.time[-1])
        return t_wait


def visualiseCombinedSchedule(combined_schedule):
    """
    Plot combined scedule so we can see density of results
    """
    print("green vertical line mark start, red vertical linesmark end of contour")
    fig, ax = plt.subplots()
    for r in range(combined_schedule.n):
        for t in range(1,len(combined_schedule.start_times[r])):
            ax.plot([combined_schedule.start_times[r][t],
                     combined_schedule.end_times[r][t]],
                    [r,r],
                    'bo-',
                    linewidth=2)

            ax.plot([combined_schedule.start_times[r][t],
                     combined_schedule.start_times[r][t]],
                    [0,combined_schedule.n-1],
                    'g--',
                    linewidth=0.5)

            ax.plot([combined_schedule.end_times[r][t],
                     combined_schedule.end_times[r][t]],
                    [0,combined_schedule.n-1],
                    'r--',
                    linewidth=0.5)
    plt.xlabel("Time")
    plt.ylabel("Robots")
    plt.show()
    return


def createCombinedSchedule(data, states):
    """
    given states, create a combined schedule
    """
    plan = CombinedSchedule(data, len(states))
    cts_left = [copy.copy(s) for s in states.values()]

    if len(cts_left)>2:
        raise Exception("New CombinedScheduler does not support >2 arms")

    done_contours = []
    t_used = 0

    robot = 0
    force=False
    assert(len(cts_left[0])>0)
    while sum([len(r) for r in cts_left])>0:
        # get next cts whose contour hasn't been done
        i = 0
        while i < len(cts_left[robot]):
            ct = cts_left[robot][i]

            if ct.contour in done_contours:
                del cts_left[robot][i]
                if len(cts_left[robot]) == 0:
                    robot = 1-robot
                    break
                continue
            else:
                # try adding
                if plan.append(ct, robot, force=force):

                    # log done contours and time
                    done_contours.append(ct.contour)
                    t_used += ct.time[-1]
                    force = False

                    # get next arm to add to
                    rs = plan.planLenOrder()
                    for r in rs:
                        if len(cts_left[r]) != 0:
                            robot = r
                            break
                    break

                # if tried everything, switch robots
                elif i == len(cts_left[robot])-1:
                    # TODO work with more than two robots
                    robot = 1-robot
                    force=True
                    break
                else:
                    # try next contour in list
                    i+=1

    return plan, t_used


def reward(data, states):
    """
    Joint Scheduler + reward function on outputplan
    """
    plan, t_used = createCombinedSchedule(data, states)
    # return t_used/(plan.len()*len(states)), plan # FOR TESTING
    return t_used/plan.len()


def decmcts_layer(
            contours,
            contour_tracker,
            env,
            plan=None,
            finish_home=True,
            col_print=False,
            mcts_nodes=100, # no. of nodes to grow on each tree
            comm_n=10, # no. of action dists to communicate
            joint2time_mul=J2TIME_C,
            default_wait = 5 # time spent waiting when traj add fails
            ):
    """
    decmcts planner
    """
    # update collision box for printed material
    if col_print:
        bound_mins, bound_maxs = tu.get_bounding_box(contour_tracker)
        bound_maxs[2] = tracker[contours[0]].pos[0][2]/1000 - 0.001
        print_block = env.add_block(bound_mins, bound_maxs)

    if plan is None:
        plan = tu.Plan(env)

    # save calculated trajectories for contours between loops
    c_traj = {}
    n_arms = env.numRobots
    arm_done = [False]*n_arms

    # DecMCTS variables
    data = {'env': env, # data used by reward function
            'collisions': {},
            'actions': [None]*n_arms}

    # DecMCTS main loop

    s_layer = time.time()

    # Update Data for trees to include latest contours
    for cid in contours:
        for arm in range(n_arms):

            if cid not in c_traj:
                c_traj[cid] = [None]*n_arms

            if c_traj[cid][arm] is None:
                multiple_trajs = tu.contour_trajs(
                        contour_tracker[cid],
                        env,
                        arm
                        )

                # only one ik per contour
                if len(multiple_trajs)>0:
                    c_traj[cid][arm] = [tu.get_closest_traj(multiple_trajs,home)]
                else:
                    c_traj[cid][arm] = []

        # is this contour impossible?
        if c_traj[cid] == [[]]*n_arms:
            warnings.warn("Contour {} impossible for both arms, SKIPPING".format(cid))
            #contour_tracker.do_contour(cid)

    # Clear old actions and unroll new ones into flat list
    for r in range(n_arms):
        data['actions'][r] = [ct
                              for c in contours
                              for ct in c_traj[c][r]
                              if ct is not None]

    # Create DecMCTS trees
    trees = []
    for i in range(n_arms):
        trees.append(Tree(data,
                          reward,
                          avail_contours,
                          state_storer,
                          sim_select,
                          avail_contours,
                          comm_n,
                          i))

    # grow mcts_nodes nodes
    s = time.time()
    for i in range(mcts_nodes):
        # grow the trees
        for tree in trees:
            tree.grow(nsims=1, depth=1e6)
        if i%1==0 and i!=0:
            e = time.time()
            print("Completed {} nodes in {:.2g}s".format(i, e-s))

        # send update to other trees
        for tree in trees:
            for other_tree in trees:
                if other_tree is not tree:
                    other_tree.receive_comms(tree.send_comms(), tree.id)

    # Best rollouts
    top_rolls = {}
    for tree in trees:
        top_rolls[tree.id] = tree.my_act_dist.best_action()

    # Get the CombinedSchedule of these rollouts
    combined_schedule, _ = createCombinedSchedule(data, top_rolls)

    # Add contours from combined schedule to real plan
    contours_left = sum([len(l) for l in combined_schedule.trajs])
    inds = [0] * n_arms
    while contours_left>0:

        # select trajectory to add next (as per plan not combined schedule)
        arms = np.argsort([plan.cumul_time[arm][-1] for arm in range(plan.n)])
        arm = None
        for i in arms:
            if inds[i] < len(combined_schedule.trajs[i]):
                arm = i
                break

        assert arm != None, "something's gone wrong here - all arms complete but total is not"

        # Trajectory to add and time to wait for it
        traj = combined_schedule.trajs[arm][inds[arm]]

        # TODO THIS VAIRABLE SHOULD BE RELATIVE TO OTHER ARM
        decmcts_wait_time = combined_schedule.start_times[arm][inds[arm]] - \
                combined_schedule.end_times[arm][inds[arm]-1]

        # Estimated travel times
        time_last2next = joint2time_mul * \
                joint_diff(plan.trajs[arm][-1].positions[-1], traj.positions[0])
                # this should be minimum value of decmcts_wait_time
        time_last2home = joint2time_mul * \
            joint_diff(plan.trajs[arm][-1].positions[-1], home)
        time_home2next = joint2time_mul * \
            joint_diff(home, traj.positions[0])

        # if there's enough time to get home in the wait period, do that
        #athome = list(plan.getLastJoint(arm)) == home
        #additional_time_spent = 0
        #if time_last2home + time_home2next <= decmcts_wait_time:

        #    # try going to home pos
        #    while not athome:
        #        gohome = plan.plan_to(env.home, arm)

        #        # if no travel found, stay at last pos
        #        if gohome is not None and plan.appendTrajectory([gohome], arm)==0:
        #            print(arm, "returning to home to wait")
        #            athome = True
        #        else:
        #            print("There's enough time to get home but couldn't find a path - Unfortunate")
        #            print("Staying put for {}s".format(default_wait))
        #            last_js = plan.trajs[arm][-1].positions[-1]
        #            stayput = tu.JTrajectory([last_js, last_js],[0,default_wait], arm)
        #            if plan.appendTrajectory([stayput], arm)!=0:
        #                print("cannot go home or stay put Logically possible but unlikely - TODO")
        #                raise Exception("cannot go home or stay put Logically possible but unlikely - TODO")
        #            additional_time_spent += default_wait

        #            if time_last2home + additional_time_spent + default_wait + time_home2next > decmcts_wait_time:
        #                # can't go home anymore
        #                break

        # add to plan
        traj_added = False
        retry_counter = 0
        while not traj_added:

            #print("---")
            # create plan to next traj
            travel_traj = plan.plan_to(traj.positions[0], arm)
            # try adding traj + contour to planner
            if travel_traj is not None and plan.appendTrajectory([travel_traj, traj], arm) == 0:
                # if success, break, else loop again
                print("contour", traj.contour, "added to real plan")
                #contour_tracker.do_contour(traj.contour)
                traj_added = True
                break
            else:
                #print("deviating from decmcts plan")
                # If the cumul time of current arm is greater than that of other arms, send other arms home
                if plan.cumul_time[arm][-1] == plan.len():
                    print("current arm has the longest plan but failing, take other arms home to clear")
                    # send other arms home
                    for other_arm in plan.trajs:
                        if other_arm == arm:
                            continue
                        # if already at home, stay at home
                        if all(plan.getLastJoint(other_arm) == home):
                            print("Other arm is already home, this should not occur !!!")
                            raise Exception(" this should not occur")

                        # Otherwise, go home
                        gohome = plan.plan_to(env.home, other_arm)
                        if gohome is not None and plan.appendTrajectory([gohome], other_arm)==0:
                            print("current arm is", arm, ", taking other arm home")
                            print("Arm", other_arm, "returning to home")
                            athome = True
                        else:
                            # other arm cannot go home, send current arm home to make room
                            if all(plan.getLastJoint(arm) == home):
                                print("other arm cannot go home despite current arm already being home - this should not happen")
                                raise Exception("other arm cannot go home despite current arm already being home - this should not happen")
                            gohome = plan.plan_to(env.home, arm)
                            if gohome is not None and plan.appendTrajectory([gohome], arm)==0:
                                print("other arm {} cannot go home, taking arm {} home since it cannot stay put".format(other_arm, arm))
                            else:
                                print("Neither arm can go home!")
                                raise Exception("Neither arm can go home!")

                last_js = plan.trajs[arm][-1].positions[-1]
                stayput = tu.JTrajectory([last_js, last_js],[0,default_wait], arm)
                if plan.appendTrajectory([stayput], arm)!=0:
                    print("cannot stay at last position, need to move out of way, going home")
                    gohome = plan.plan_to(env.home, arm)
                    if gohome is not None and plan.appendTrajectory([gohome], arm)==0:
                        print("taking arm {} home since it cannot stay put".format(arm))
                    else:
                        print("Cannot stay put or go home!")
                        raise Exception("Cannot stay put or go home!")


        contours_left -= 1 # to ensure loop exits
        inds[arm] += 1 # move to next contour next time

    # take arms home
    if finish_home:
        arms_home=[list(plan.getLastJoint(a)) == home for a in range(n_arms)]
        arm = 0
        retries = 0
        while not all(arms_home) and retries < n_arms:
            gohome = plan.plan_to(env.home, arm)
            if list(plan.getLastJoint(arm)) != home and \
                    gohome is not None and \
                    plan.appendTrajectory([gohome], arm)==0:
                print("Arm {} safely returned to home pose".format(arm))
                retries = 0
                arms_home[arm] = True
            else:
                print("Failed to take arm {} home, trying another arm".format(arm))
                retries +=1
            arm += 1
            if arm == n_arms:
                arm = 0

    e_layer = time.time()
    print("last layer completed in {}s".format(e_layer-s_layer))

    return combined_schedule, plan


def decmcts(
        sim_env,
        tracker,
        nth_layer=1,
        layers=None,
        nodes_per_layer=1000,
        interactive=False):
    """
    plan whole print with decmcts using decmcts_layer
    """
    if layers is None:
        n_layers = len(tracker.layers)
        layers = list(range(0, n_layers, nth_layer))

    plan = None
    fin_home = False
    print("Planning the following layers: {}".format(layers))

    for i in layers:
        print("PLANNING LAYER {}".format(i))

        # take arms home if final layer
        if i == layers[-1]:
            fin_home = True

        cs, plan = decmcts_layer(
                tracker.layers[i],
                tracker,
                sim_env,
                plan=plan,
                finish_home=fin_home,
                mcts_nodes=nodes_per_layer
                )

        if interactive:
            visualiseCombinedSchedule(cs)
            plan.visualise()
    return plan


if __name__ == "__main__":

    # Read gcode file
    model = "CuteOcto_x2"
    tracker = tu.get_tracker(model)

    # 3arms simulation
    env_desc3 = tu.PlanningEnv([[0,0,1,0],[0,0,1,0],[0,0,0,1]],
                       [[0.9, -0.35, 0],[-0.8, -0.35, 0.001],[0, 0.5, 0]],
                       [[0,0,0,1]]*3,
                       [[0.15,0,0]]*3) # WARNING - also edit urdf

    # 2arms real calibration
    env_desc2 = tu.read_env('calibrations/r1_tforms.yaml',
                            'calibrations/r2_tforms.yaml')
    sim_env = tu.SimEnv(env_desc=env_desc2, gui=True, home=home)

    # decmcts plan first n contours
    plan = decmcts(sim_env, tracker, layers=[1,5], interactive=True)