# Source code for rl_coach.memories.non_episodic.prioritized_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 operator
    import random
    from enum import Enum
    from typing import List, Tuple, Any

    import numpy as np

    from rl_coach.core_types import Transition
    from rl_coach.memories.memory import MemoryGranularity
    from rl_coach.memories.non_episodic.experience_replay import ExperienceReplayParameters, ExperienceReplay
    from rl_coach.schedules import Schedule, ConstantSchedule

    class PrioritizedExperienceReplayParameters(ExperienceReplayParameters):
        def __init__(self):
            super().__init__()
            self.max_size = (MemoryGranularity.Transitions, 1000000)
            self.alpha = 0.6
            self.beta = ConstantSchedule(0.4)
            self.epsilon = 1e-6

        @property
        def path(self):
            return 'rl_coach.memories.non_episodic.prioritized_experience_replay:PrioritizedExperienceReplay'

    class SegmentTree(object):
        """
        A tree which can be used as a min/max heap or a sum tree
        Add or update item value - O(log N)
        Sampling an item - O(log N)
        """
        class Operation(Enum):
            MAX = {"operator": max, "initial_value": -float("inf")}
            MIN = {"operator": min, "initial_value": float("inf")}
            SUM = {"operator": operator.add, "initial_value": 0}

        def __init__(self, size: int, operation: Operation):
            self.next_leaf_idx_to_write = 0
            self.size = size
            if not (size > 0 and size & (size - 1) == 0):
                raise ValueError("A segment tree size must be a positive power of 2. The given size is {}".format(self.size))
            self.operation = operation
            self.tree = np.ones(2 * size - 1) * self.operation.value['initial_value']
            self.data = [None] * size

        def _propagate(self, node_idx: int) -> None:
            """
            Propagate an update of a node's value to its parent node
            :param node_idx: the index of the node that was updated
            :return: None
            """
            parent = (node_idx - 1) // 2

            self.tree[parent] = self.operation.value['operator'](self.tree[parent * 2 + 1], self.tree[parent * 2 + 2])

            if parent != 0:
                self._propagate(parent)

        def _retrieve(self, root_node_idx: int, val: float)-> int:
            """
            Retrieve the first node that has a value larger than val and is a child of the node at index idx
            :param root_node_idx: the index of the root node to search from
            :param val: the value to query for
            :return: the index of the resulting node
            """
            left = 2 * root_node_idx + 1
            right = left + 1

            if left >= len(self.tree):
                return root_node_idx

            if val <= self.tree[left]:
                return self._retrieve(left, val)
            else:
                return self._retrieve(right, val-self.tree[left])

        def total_value(self) -> float:
            """
            Return the total value of the tree according to the tree operation. For SUM for example, this will return
            the total sum of the tree. for MIN, this will return the minimal value
            :return: the total value of the tree
            """
            return self.tree[0]

        def add(self, val: float, data: Any) -> None:
            """
            Add a new value to the tree with data assigned to it
            :param val: the new value to add to the tree
            :param data: the data that should be assigned to this value
            :return: None
            """
            self.data[self.next_leaf_idx_to_write] = data
            self.update(self.next_leaf_idx_to_write, val)

            self.next_leaf_idx_to_write += 1
            if self.next_leaf_idx_to_write >= self.size:
                self.next_leaf_idx_to_write = 0

        def update(self, leaf_idx: int, new_val: float) -> None:
            """
            Update the value of the node at index idx
            :param leaf_idx: the index of the node to update
            :param new_val: the new value of the node
            :return: None
            """
            node_idx = leaf_idx + self.size - 1
            if not 0 <= node_idx < len(self.tree):
                raise ValueError("The given left index ({}) can not be found in the tree. The available leaves are: 0-{}"
                                 .format(leaf_idx, self.size - 1))

            self.tree[node_idx] = new_val
            self._propagate(node_idx)

        def get_element_by_partial_sum(self, val: float) -> Tuple[int, float, Any]:
            """
            Given a value between 0 and the tree sum, return the object which this value is in it's range.
            For example, if we have 3 leaves: 10, 20, 30, and val=35, this will return the 3rd leaf, by accumulating
            leaves by their order until getting to 35. This allows sampling leaves according to their proportional
            probability.
            :param val: a value within the range 0 and the tree sum
            :return: the index of the resulting leaf in the tree, its probability and
                     the object itself
            """
            node_idx = self._retrieve(0, val)
            leaf_idx = node_idx - self.size + 1
            data_value = self.tree[node_idx]
            data = self.data[leaf_idx]

            return leaf_idx, data_value, data

        def __str__(self):
            result = ""
            start = 0
            size = 1
            while size <= self.size:
                result += "{}\n".format(self.tree[start:(start + size)])
                start += size
                size *= 2
            return result

    class PrioritizedExperienceReplay(ExperienceReplay):
        """
        This is the proportional sampling variant of the prioritized experience replay as described
        in https://arxiv.org/pdf/1511.05952.pdf.
        """
        def __init__(self, max_size: Tuple[MemoryGranularity, int], alpha: float=0.6, beta: Schedule=ConstantSchedule(0.4),
                     epsilon: float=1e-6, allow_duplicates_in_batch_sampling: bool=True):
            """
            :param max_size: the maximum number of transitions or episodes to hold in the memory
            :param alpha: the alpha prioritization coefficient
            :param beta: the beta parameter used for importance sampling
            :param epsilon: a small value added to the priority of each transition
            :param allow_duplicates_in_batch_sampling: allow having the same transition multiple times in a batch
            """
            if max_size[0] != MemoryGranularity.Transitions:
                raise ValueError("Prioritized Experience Replay currently only support setting the memory size in "
                                 "transitions granularity.")
            self.power_of_2_size = 1
            while self.power_of_2_size < max_size[1]:
                self.power_of_2_size *= 2
            super().__init__((MemoryGranularity.Transitions, self.power_of_2_size), allow_duplicates_in_batch_sampling)
            self.sum_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.SUM)
            self.min_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.MIN)
            self.max_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.MAX)
            self.alpha = alpha
            self.beta = beta
            self.epsilon = epsilon
            self.maximal_priority = 1.0

        def _update_priority(self, leaf_idx: int, error: float) -> None:
            """
            Update the priority of a given transition, using its index in the tree and its error
            :param leaf_idx: the index of the transition leaf in the tree
            :param error: the new error value
            :return: None
            """
            if error < 0:
                raise ValueError("The priorities must be non-negative values")
            priority = (error + self.epsilon)
            self.sum_tree.update(leaf_idx, priority ** self.alpha)
            self.min_tree.update(leaf_idx, priority ** self.alpha)
            self.max_tree.update(leaf_idx, priority)
            self.maximal_priority = self.max_tree.total_value()

        def update_priorities(self, indices: List[int], error_values: List[float]) -> None:
            """
            Update the priorities of a batch of transitions using their indices and their new TD error terms
            :param indices: the indices of the transitions to update
            :param error_values: the new error values
            :return: None
            """
            self.reader_writer_lock.lock_writing_and_reading()

            if len(indices) != len(error_values):
                raise ValueError("The number of indexes requested for update don't match the number of error values given")
            for transition_idx, error in zip(indices, error_values):
                self._update_priority(transition_idx, error)

            self.reader_writer_lock.release_writing_and_reading()

        def sample(self, size: int) -> List[Transition]:
            """
            Sample a batch of transitions form the replay buffer. If the requested size is larger than the number
            of samples available in the replay buffer then the batch will return empty.
            :param size: the size of the batch to sample
            :return: a batch (list) of selected transitions from the replay buffer
            """

            self.reader_writer_lock.lock_writing()

            if self.num_transitions() >= size:
                # split the tree leaves to equal segments and sample one transition from each segment
                batch = []
                segment_size = self.sum_tree.total_value() / size

                # get the maximum weight in the memory
                min_probability = self.min_tree.total_value() / self.sum_tree.total_value()  # min P(j) = min p^a / sum(p^a)
                max_weight = (min_probability * self.num_transitions()) ** -self.beta.current_value  # max wi

                # sample a batch
                for i in range(size):
                    segment_start = segment_size * i
                    segment_end = segment_size * (i + 1)

                    # sample leaf and calculate its weight
                    val = random.uniform(segment_start, segment_end)
                    leaf_idx, priority, transition = self.sum_tree.get_element_by_partial_sum(val)
                    priority /= self.sum_tree.total_value()   # P(j) = p^a / sum(p^a)
                    weight = (self.num_transitions() * priority) ** -self.beta.current_value  # (N * P(j)) ^ -beta
                    normalized_weight = weight / max_weight  # wj = ((N * P(j)) ^ -beta) / max wi

                    transition.info['idx'] = leaf_idx
                    transition.info['weight'] = normalized_weight

                    batch.append(transition)

                self.beta.step()

            else:
                raise ValueError("The replay buffer cannot be sampled since there are not enough transitions yet. "
                                 "There are currently {} transitions".format(self.num_transitions()))

            self.reader_writer_lock.release_writing()
            return batch

        def store(self, transition: Transition, lock=True) -> None:
            """
            Store a new transition in the memory.
            :param transition: a transition to store
            :return: None
            """
            # Calling super.store() so that in case a memory backend is used, the memory backend can store this transition.
            super().store(transition)

            if lock:
                self.reader_writer_lock.lock_writing_and_reading()

            transition_priority = self.maximal_priority
            self.sum_tree.add(transition_priority ** self.alpha, transition)
            self.min_tree.add(transition_priority ** self.alpha, transition)
            self.max_tree.add(transition_priority, transition)
            super().store(transition, False)

            if lock:
                self.reader_writer_lock.release_writing_and_reading()

        def clean(self, lock=True) -> None:
            """
            Clean the memory by removing all the episodes
            :return: None
            """
            if lock:
                self.reader_writer_lock.lock_writing_and_reading()

            super().clean(lock=False)
            self.sum_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.SUM)
            self.min_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.MIN)
            self.max_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.MAX)

            if lock:
                self.reader_writer_lock.release_writing_and_reading()