from functools import partial
import multiprocessing as mp
import numpy as np
import os
import sys

from dowel import logger, tabular
from garage.sampler.on_policy_vectorized_sampler import OnPolicyVectorizedSampler

from ast_toolbox.rewards import ExampleAVReward
from ast_toolbox.simulators import ExampleAVSimulator

os.sched_setaffinity(0, {i for i in range(mp.cpu_count())})

def myexcepthook(exctype, value, traceback):
    for p in mp.active_children():
       p.terminate()

def slice_dict(in_dict, slice_idx):
    """Helper function to recursively parse through a dictionary of dictionaries and arrays to slice \
    the arrays at a certain index.

    Parameters
    ----------
    in_dict : dict
        Dictionary where the values are arrays or other dictionaries that follow this stipulation.
    slice_idx : int
        Index to slice each array at.

    Returns
    -------
    dict
        Dictionary where arrays at every level are sliced.

    """
    for key, value in in_dict.items():
        # pdb.set_trace()
        if isinstance(value, dict):
            in_dict[key] = slice_dict(value, slice_idx)
        else:
            in_dict[key][slice_idx + 1:, ...] = np.zeros_like(value[slice_idx + 1:, ...])

    return in_dict

def simulate_single_path_worker(sim_func, slice_func, reward_func, reward_info_func, sim_from_s0, path):
    if sim_from_s0:
        s_0 = path["observations"][0]
    else:
        s_0 = None
    actions = path['env_infos']['actions']

    end_idx, info = sim_func(actions=actions, s_0=s_0)

    if end_idx >= 0:
        slice_func(path, end_idx)


    rewards = reward_func(
        action=actions[end_idx],
        info=reward_info_func()
    )

    path["rewards"][end_idx] = rewards
    path['env_infos']['cache'] = np.zeros_like(path["rewards"])
    for i in range(len(info)):
        path['env_infos']['state'][i,:-2] = info[i]['state']
    return path

class ASTVectorizedSampler(OnPolicyVectorizedSampler):
    """A vectorized sampler for AST to handle open-loop simulators.

    Garage usually genearates samples in a closed-loop process. This version of the vectorized sampler instead
    grabs dummy data until the full rollout specification is generated, then goes back and runs the `simulate`
    function to actually obtain results. Rewards are then calculated and the path data is corrected.

    Parameters
    ----------
    algo : :py:class:`garage.np.algos.base.RLAlgorithm`
        The algorithm.
    env : :py:class:`ast_toolbox.envs.ASTEnv`
        The environment.
    n_envs : int
        Number of parallel environments to run.
    open_loop : bool
        True if the simulation is open-loop, meaning that AST must generate all actions ahead of time, instead
        of being able to output an action in sync with the simulator, getting an observation back before
        the next action is generated. False to get interactive control, which requires that `blackbox_sim_state`
        is also False.
    sim : :py:class:`ast_toolbox.simulators.ASTSimulator`
        The simulator wrapper, inheriting from `ast_toolbox.simulators.ASTSimulator`.
    reward_function : :py:class:`ast_toolbox.rewards.ASTReward`
        The reward function, inheriting from `ast_toolbox.rewards.ASTReward`.

    """

    def __init__(self, algo, env, n_envs=1, open_loop=True,  batch_simulate=False,
                 sim=ExampleAVSimulator(), reward_function=ExampleAVReward(),
                 sim_from_s0=False):

        # pdb.set_trace()
        self.open_loop = open_loop
        self.batch_simulate = batch_simulate
        self.sim = sim
        self.reward_function = reward_function
        self.sim_from_s0 = sim_from_s0
        super().__init__(algo, env, n_envs)

    def obtain_samples(self, itr, batch_size=None, whole_paths=False):
        """Sample the policy for new trajectories.

        Parameters
        ----------
        itr : int
            Iteration number.
        batch_size : int
            Number of samples to be collected. If None,
            it will be default [algo.max_path_length * n_envs].
        whole_paths : bool
            Whether return all the paths or not. True
            by default. It's possible for the paths to have total actual
            sample size larger than batch_size, and will be truncated if
            this flag is true.

        Returns
        -------
        : list[dict]
            A list of sampled rollout paths.
            Each rollout path is a dictionary with the following keys:
                - observations (numpy.ndarray)
                - actions (numpy.ndarray)
                - rewards (numpy.ndarray)
                - agent_infos (dict)
                - env_infos (dict)

        """
        # pdb.set_trace()
        paths = super().obtain_samples(itr, batch_size)
        # pdb.set_trace()
        if self.open_loop:
            # mp.log_to_stderr(logging.DEBUG)
            if self.batch_simulate:
                # import pdb; pdb.set_trace()
                def store_process_return_value(func, queue, **kwargs):
                    queue.put(func(**kwargs))


                queue = mp.Queue()
                process_kwargs = {'func':self.sim.batch_simulate_paths,
                                  'queue': queue,
                                  'paths': paths,
                                  'reward_function': self.reward_function}

                p = mp.Process(target=store_process_return_value, kwargs=process_kwargs)
                p.start()
                p.join()

                if p.exitcode == 0:
                    paths = queue.get()
                    p.terminate()

                paths = self.sim.batch_simulate_paths(paths=paths, reward_function=self.reward_function)
            else:
                if self.n_envs > 1:
                    print("Opening Pool with ", self.n_envs, ' processes.')
                    sys.excepthook = myexcepthook

                    simulate_single_path_worker_partial = partial(simulate_single_path_worker,
                                                                  self.sim.simulate,
                                                                  slice_dict,
                                                                  self.reward_function.give_reward,
                                                                  self.sim.get_reward_info,
                                                                  self.sim_from_s0)

                    # pool = mp.Pool(processes=self.n_envs)
                    # paths = pool.map(simulate_single_path_worker_partial, paths)
                    # pool.close()
                    # print('Pool Closed')

                    from pebble import ProcessPool
                    from concurrent.futures import TimeoutError
                    with ProcessPool() as pool:
                        future = pool.map(simulate_single_path_worker_partial, paths, timeout=300)

                        iterator = future.result()
                        paths = []

                        while True:
                            try:
                                paths.append(next(iterator))
                            except StopIteration:
                                break
                            except TimeoutError as error:
                                print("THREAD TIMEOUT: function took longer than %d seconds" % error.args[1])
                            except Exception as error:
                                print("Simulation process raised %s" % error)
                                print(error.traceback)  # Python's traceback of remote process

                else:
                    new_paths = []
                    init_state = self.sim.clone_state().copy()
                    for path in paths:
                        # import pdb
                        # pdb.set_trace()
                        self.sim.restore_state(init_state)
                        new_path = simulate_single_path_worker(sim_func=self.sim.simulate,
                                                    slice_func=slice_dict,
                                                    reward_func=self.reward_function.give_reward,
                                                    reward_info_func=self.sim.get_reward_info,
                                                    sim_from_s0=self.sim_from_s0,
                                                    path=path)
                        new_paths.append(new_path)
                        # print(new_path['env_infos']['state'][:,9:13])

                    paths = new_paths


            # for path in paths:
            #     s_0 = path["observations"][0]
            #
            #     # actions = path['env_infos']['info']['actions']
            #     actions = path['actions']
            #     # pdb.set_trace()
            #     end_idx, info = self.sim.simulate(actions=actions, s_0=s_0)
            #     # print('----- Back from simulate: ', end_idx)
            #     if end_idx >= 0:
            #         # pdb.set_trace()
            #         self.slice_dict(path, end_idx)
            #     rewards = self.reward_function.give_reward(
            #         action=actions[end_idx],
            #         info=self.sim.get_reward_info()
            #     )
            #     # print('----- Back from rewards: ', rewards)
            #     # pdb.set_trace()
            #     path["rewards"][end_idx] = rewards
            #     # info[:, -1] = path["rewards"][:info.shape[0]]
            #     # path['env_infos']['sim_info'] = info
            #     path['env_infos']['sim_info'] = np.zeros_like(path["rewards"])
            #     # import pdb; pdb.set_trace()

        total_steps = 0
        for path in paths:
            total_steps += path['actions'].shape[0]

        tabular.record('TotalSteps', total_steps)

        return paths

    def slice_dict(self, in_dict, slice_idx):
        """Helper function to recursively parse through a dictionary of dictionaries and arrays to slice \
        the arrays at a certain index.

        Parameters
        ----------
        in_dict : dict
            Dictionary where the values are arrays or other dictionaries that follow this stipulation.
        slice_idx : int
            Index to slice each array at.

        Returns
        -------
        dict
            Dictionary where arrays at every level are sliced.

        """
        for key, value in in_dict.items():
            # pdb.set_trace()
            if isinstance(value, dict):
                in_dict[key] = self.slice_dict(value, slice_idx)
            else:
                in_dict[key][slice_idx + 1:, ...] = np.zeros_like(value[slice_idx + 1:, ...])

        return in_dict
