# Source code for rl_coach.memories.episodic.episodic_hindsight_experience_replay

    #
    # Copyright (c) 2017 Intel Corporation
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    #      http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.
    #

    import copy
    from enum import Enum
    from typing import Tuple, List

    import numpy as np

    from rl_coach.core_types import Episode, Transition
    from rl_coach.memories.episodic.episodic_experience_replay import EpisodicExperienceReplayParameters, \
        EpisodicExperienceReplay
    from rl_coach.memories.non_episodic.experience_replay import MemoryGranularity
    from rl_coach.spaces import GoalsSpace

    class HindsightGoalSelectionMethod(Enum):
        Future = 0
        Final = 1
        Episode = 2
        Random = 3

    class EpisodicHindsightExperienceReplayParameters(EpisodicExperienceReplayParameters):
        def __init__(self):
            super().__init__()
            self.hindsight_transitions_per_regular_transition = None
            self.hindsight_goal_selection_method = None
            self.goals_space = None

        @property
        def path(self):
            return 'rl_coach.memories.episodic.episodic_hindsight_experience_replay:EpisodicHindsightExperienceReplay'

    class EpisodicHindsightExperienceReplay(EpisodicExperienceReplay):
        """
        Implements Hindsight Experience Replay as described in the following paper: https://arxiv.org/pdf/1707.01495.pdf

        """
        def __init__(self, max_size: Tuple[MemoryGranularity, int],
                     hindsight_transitions_per_regular_transition: int,
                     hindsight_goal_selection_method: HindsightGoalSelectionMethod,
                     goals_space: GoalsSpace):
            """
            :param max_size: The maximum size of the memory. should be defined in a granularity of Transitions
            :param hindsight_transitions_per_regular_transition: The number of hindsight artificial transitions to generate
                                                                 for each actual transition
            :param hindsight_goal_selection_method: The method that will be used for generating the goals for the
                                                    hindsight transitions. Should be one of HindsightGoalSelectionMethod
            :param goals_space: A GoalsSpace which defines the base properties of the goals space
            """
            super().__init__(max_size)

            self.hindsight_transitions_per_regular_transition = hindsight_transitions_per_regular_transition
            self.hindsight_goal_selection_method = hindsight_goal_selection_method
            self.goals_space = goals_space
            self.last_episode_start_idx = 0

        def _sample_goal(self, episode_transitions: List, transition_index: int):
            """
            Sample a single goal state according to the sampling method
            :param episode_transitions: a list of all the transitions in the current episode
            :param transition_index: the transition to start sampling from
            :return: a goal corresponding to the sampled state
            """
            if self.hindsight_goal_selection_method == HindsightGoalSelectionMethod.Future:
                # states that were observed in the same episode after the transition that is being replayed
                selected_transition = np.random.choice(episode_transitions[transition_index+1:])
            elif self.hindsight_goal_selection_method == HindsightGoalSelectionMethod.Final:
                # the final state in the episode
                selected_transition = episode_transitions[-1]
            elif self.hindsight_goal_selection_method == HindsightGoalSelectionMethod.Episode:
                # a random state from the episode
                selected_transition = np.random.choice(episode_transitions)
            elif self.hindsight_goal_selection_method == HindsightGoalSelectionMethod.Random:
                # a random state from the entire replay buffer
                selected_transition = np.random.choice(self.transitions)
            else:
                raise ValueError("Invalid goal selection method was used for the hindsight goal selection")
            return self.goals_space.goal_from_state(selected_transition.state)

        def _sample_goals(self, episode_transitions: List, transition_index: int):
            """
            Sample a batch of goal states according to the sampling method
            :param episode_transitions: a list of all the transitions in the current episode
            :param transition_index: the transition to start sampling from
            :return: a goal corresponding to the sampled state
            """
            return [
                self._sample_goal(episode_transitions, transition_index)
                for _ in range(self.hindsight_transitions_per_regular_transition)
            ]

        def store_episode(self, episode: Episode, lock: bool=True) -> None:
            # generate hindsight transitions only when an episode is finished
            last_episode_transitions = copy.copy(episode.transitions)

            # cannot create a future hindsight goal in the last transition of an episode
            if self.hindsight_goal_selection_method == HindsightGoalSelectionMethod.Future:
                relevant_base_transitions = last_episode_transitions[:-1]
            else:
                relevant_base_transitions = last_episode_transitions

            # for each transition in the last episode, create a set of hindsight transitions
            for transition_index, transition in enumerate(relevant_base_transitions):
                sampled_goals = self._sample_goals(last_episode_transitions, transition_index)
                for goal in sampled_goals:
                    hindsight_transition = copy.copy(transition)

                    if hindsight_transition.state['desired_goal'].shape != goal.shape:
                        raise ValueError((
                            'goal shape {goal_shape} already in transition is '
                            'different than the one sampled as a hindsight goal '
                            '{hindsight_goal_shape}.'
                        ).format(
                            goal_shape=hindsight_transition.state['desired_goal'].shape,
                            hindsight_goal_shape=goal.shape,
                        ))

                    # update the goal in the transition
                    hindsight_transition.state['desired_goal'] = goal
                    hindsight_transition.next_state['desired_goal'] = goal

                    # update the reward and terminal signal according to the goal
                    hindsight_transition.reward, hindsight_transition.game_over = \
                        self.goals_space.get_reward_for_goal_and_state(goal, hindsight_transition.next_state)

                    hindsight_transition.n_step_discounted_rewards = None
                    episode.insert(hindsight_transition)

            super().store_episode(episode)

        def store(self, transition: Transition):
            raise ValueError("An episodic HER cannot store a single transition. Only full episodes are to be stored.")