import numpy as np
from collections import deque, OrderedDict
from rlkit.core import logger
from rlkit.core.eval_util import create_stats_ordered_dict
from rlkit.samplers.rollout_functions import rollout, multitask_rollout, multitask_rollout_with_relabeler
from rlkit.samplers.data_collector.base import PathCollector
from rlkit.torch.multitask.pointmass_rewards import PointMassBestRandomRelabeler

import os.path as osp


class MdpPathCollector(PathCollector):
    def __init__(
            self,
            env,
            policy,
            max_num_epoch_paths_saved=None,
            render=False,
            render_kwargs=None,
    ):
        if render_kwargs is None:
            render_kwargs = {}
        self._env = env
        self._policy = policy
        self._max_num_epoch_paths_saved = max_num_epoch_paths_saved
        self._epoch_paths = deque(maxlen=self._max_num_epoch_paths_saved)
        self._render = render
        self._render_kwargs = render_kwargs

        self._num_steps_total = 0
        self._num_paths_total = 0

    def collect_new_paths(
            self,
            max_path_length,
            num_steps,
            discard_incomplete_paths,
    ):
        paths = []
        num_steps_collected = 0
        while num_steps_collected < num_steps:
            max_path_length_this_loop = min(  # Do not go over num_steps
                max_path_length,
                num_steps - num_steps_collected,
            )
            path = rollout(
                self._env,
                self._policy,
                max_path_length=max_path_length_this_loop,
                render=self._render,
                render_kwargs=self._render_kwargs
            )
            path_len = len(path['actions'])
            if (
                    path_len != max_path_length
                    and not path['terminals'][-1]
                    and discard_incomplete_paths
            ):
                break
            num_steps_collected += path_len
            paths.append(path)
        self._num_paths_total += len(paths)
        self._num_steps_total += num_steps_collected
        self._epoch_paths.extend(paths)
        return paths

    def get_epoch_paths(self):
        return self._epoch_paths

    def end_epoch(self, epoch):
        self._epoch_paths = deque(maxlen=self._max_num_epoch_paths_saved)

    def get_diagnostics(self):
        path_lens = [len(path['actions']) for path in self._epoch_paths]
        stats = OrderedDict([
            ('num steps total', self._num_steps_total),
            ('num paths total', self._num_paths_total),
        ])
        stats.update(create_stats_ordered_dict(
            "path length",
            path_lens,
            always_show_all_stats=True,
        ))
        return stats

    def get_snapshot(self):
        return dict(
            env=self._env,
            policy=self._policy,
        )


class GoalConditionedPathCollector(PathCollector):
    def __init__(
            self,
            env,
            policy,
            max_num_epoch_paths_saved=None,
            render=False,
            render_kwargs=None,
            observation_key='observation',
            desired_goal_key='desired_goal',
    ):
        if render_kwargs is None:
            render_kwargs = {}
        self._env = env
        self._policy = policy
        self._max_num_epoch_paths_saved = max_num_epoch_paths_saved
        self._render = render
        self._render_kwargs = render_kwargs
        self._epoch_paths = deque(maxlen=self._max_num_epoch_paths_saved)
        self._observation_key = observation_key
        self._desired_goal_key = desired_goal_key

        self._num_steps_total = 0
        self._num_paths_total = 0

    def collect_new_paths(
            self,
            max_path_length,
            num_steps,
            discard_incomplete_paths,
    ):
        paths = []
        num_steps_collected = 0
        while num_steps_collected < num_steps:
            max_path_length_this_loop = min(  # Do not go over num_steps
                max_path_length,
                num_steps - num_steps_collected,
            )
            path = multitask_rollout(
                self._env,
                self._policy,
                max_path_length=max_path_length_this_loop,
                render=self._render,
                render_kwargs=self._render_kwargs,
                observation_key=self._observation_key,
                desired_goal_key=self._desired_goal_key,
                return_dict_obs=True,
            )
            path_len = len(path['actions'])
            if (
                    path_len != max_path_length
                    and not path['terminals'][-1]
                    and discard_incomplete_paths
            ):
                break
            num_steps_collected += path_len
            paths.append(path)
        self._num_paths_total += len(paths)
        self._num_steps_total += num_steps_collected
        self._epoch_paths.extend(paths)
        return paths

    def get_epoch_paths(self):
        return self._epoch_paths

    def end_epoch(self, epoch):
        self._epoch_paths = deque(maxlen=self._max_num_epoch_paths_saved)

    def get_diagnostics(self):
        path_lens = [len(path['actions']) for path in self._epoch_paths]
        stats = OrderedDict([
            ('num steps total', self._num_steps_total),
            ('num paths total', self._num_paths_total),
        ])
        stats.update(create_stats_ordered_dict(
            "path length",
            path_lens,
            always_show_all_stats=True,
        ))
        return stats

    def get_snapshot(self):
        return dict(
            env=self._env,
            policy=self._policy,
            observation_key=self._observation_key,
            desired_goal_key=self._desired_goal_key,
        )


class TaskConditionedPathCollector(PathCollector):
    def __init__(
            self,
            env,
            policy,
            relabeler,
            max_num_epoch_paths_saved=None,
            render=False,
            render_kwargs=None,
            is_eval=False,
            calculate_r_d=True,
            hide_latent=False,
            normalize_performance=False,
            save_videos=False
    ):
        if render_kwargs is None:
            render_kwargs = {}
        self._env = env
        self._policy = policy
        self._relabeler = relabeler
        self._max_num_epoch_paths_saved = max_num_epoch_paths_saved
        self._render = render
        self._render_kwargs = render_kwargs
        self._epoch_paths = deque(maxlen=self._max_num_epoch_paths_saved)
        self.is_eval = is_eval
        self.plot_pointmass = is_eval and isinstance(relabeler, PointMassBestRandomRelabeler) # False
        if self.plot_pointmass:
            self.latents = np.load(osp.expanduser('~/workspace/rlkit/data/eval_plot_latents.npy'))
            self.latent_grid_rewards = np.load(osp.expanduser('~/workspace/rlkit/data/eval_grid_rew_power1.npy'))
        self.calculate_r_d = calculate_r_d
        self.hide_latent = hide_latent
        self.normalize_performance = normalize_performance
        self._num_steps_total = 0
        self._num_paths_total = 0
        self._epoch = 0

        # for video saving
        self.save_video = save_videos and not self.is_eval and hasattr(relabeler, 'to_save_video')

        # for normalizing later
        self.eval_traj_infos = {}

    def collect_new_paths(
            self,
            max_path_length,
            num_steps,
            discard_incomplete_paths,
    ):
        paths = []
        num_steps_collected = 0

        if self.save_video and self._relabeler.to_save_video(self._epoch):
            render = True
            render_kwargs = dict(mode='rgb_array')
        else:
            render = self._render
            render_kwargs = self._render_kwargs
        while num_steps_collected < num_steps:
            max_path_length_this_loop = min(  # Do not go over num_steps
                max_path_length,
                num_steps - num_steps_collected,
            )
            path = multitask_rollout_with_relabeler(
                self._env,
                self._policy,
                self._relabeler,
                max_path_length=max_path_length_this_loop,
                render=render,
                render_kwargs=render_kwargs,
                return_dict_obs=False,
                calculate_r_d=self.calculate_r_d,
                hide_latent=self.hide_latent
            )
            path_len = len(path['actions'])
            if (
                    path_len != max_path_length
                    and not path['terminals'][-1]
                    and discard_incomplete_paths
            ):
                break
            num_steps_collected += path_len
            paths.append(path)
        self._num_paths_total += len(paths)
        self._num_steps_total += num_steps_collected
        if self.normalize_performance:
            for path in paths:
                path['normalized_rewards'] = self._relabeler.get_normalized_path_rewards(path)
        self._epoch_paths.extend(paths)
        return paths

    def get_epoch_paths(self):
        return self._epoch_paths

    def end_epoch(self, epoch):
        if self.is_eval:
            epoch_path_info = []
            for path in self._epoch_paths:
                latent = path['latents'][0]
                rewards = self._relabeler.calculate_path_reward(path, latent)
                epoch_path_info.append((latent, rewards.sum()))
            self.eval_traj_infos['epoch{}'.format(self._epoch)] = epoch_path_info
            logger.save_extra_data(self.eval_traj_infos, 'eval_traj_infos.pkl')

        self._epoch_paths = deque(maxlen=self._max_num_epoch_paths_saved)
        if self.plot_pointmass:
            paths = []
            for z in self.latents:
                path = multitask_rollout_with_relabeler(
                    self._env,
                    self._policy,
                    self._relabeler,
                    max_path_length=15,
                    render=self._render,
                    render_kwargs=self._render_kwargs,
                    return_dict_obs=False,
                    latent=z
                )
                paths.append(path)
            self._relabeler.plot_multiple_heatmaps(self.latents, paths, "Epoch" + str(self._epoch),
                                                   grid_rewards=self.latent_grid_rewards)
            # save the trajectories (in case I want to do the methods overlaid)
            np.save(osp.join(logger.get_snapshot_dir(), str(self._epoch)),
                    np.array([path['observations'] for path in paths]))
        self._epoch += 1

    def get_diagnostics(self):
        path_lens = [len(path['actions']) for path in self._epoch_paths]
        stats = OrderedDict([
            ('num steps total', self._num_steps_total),
            ('num paths total', self._num_paths_total),
        ])
        stats.update(create_stats_ordered_dict(
            "path length",
            path_lens,
            always_show_all_stats=True,
        ))
        return stats

    def get_snapshot(self):
        return dict(
            env=self._env,
            policy=self._policy,
        )
