text: from utils import *
  import sys
  import random
  import numpy as np
  from math import log
  
  MAIN_TREE = 0
  
  class SumTree(object):
      """Implementation of a SumTree with multiple trees covering the same data
      array, adapted from: https://github.com/jaara/AI-blog/blob/master/SumTree.py
      Each transition is stored along with its trace_id
      """
  
      def __init__(self, capacity, e, a):
          self.capacity = capacity
          self.write = 0
          self.trees = {}
          self.main_tree = MAIN_TREE
          self.data = np.repeat(((None, None, None), ), capacity, axis=0)
          self.updates = {}
          self.e = e
          self.a = a
  
      def copy_tree(self, trg_i, src_i=MAIN_TREE):
          """Copies src_i's priorities into a new tree trg_i
  
          Arguments:
              trg_i {object} -- Target tree identifier
  
          Keyword Arguments:
              src_i {object} -- Source tree identifier (default: {MAIN_TREE})
          """
  
          if trg_i not in self.trees:
              self.trees[trg_i] = np.copy(self.trees[src_i])
              self.updates[trg_i] = 0
  
      def create(self, i):
          """Create tree i, either by copying the main tree if it exists, or by
              creating a new tree from scratch
  
          Arguments:
              i {object} -- The new tree's identifier
          """
  
          if i is None:
              i = self.main_tree
          if i not in self.trees and self.main_tree in self.trees:
              self.copy_tree(i, self.main_tree)
          elif i not in self.trees:
              self.trees[i] = np.zeros(2 * self.capacity - 1)
              self.updates[i] = 0
  
      def _propagate(self, idx, change, tree_id=None):
          """Propagate priority changes to root
  
          Arguments:
              idx {int} -- Node to propagate from
              change {float} -- Priority difference to propagate
  
          Keyword Arguments:
              tree_id {object} -- Which tree the change applies to (default: {None})
          """
  
          tree_id = tree_id if tree_id is not None else self.main_tree
  
          parent = (idx - 1) // 2
  
          self.trees[tree_id][parent] += change
  
          if parent != 0:
              self._propagate(parent, change, tree_id)
  
      def _retrieve(self, idx, s, tree_id=None):
          """Retrieve the node covering offset s starting from node idx
  
          Arguments:
              idx {int} -- Note to start from
              s {float} -- offset
  
          Keyword Arguments:
              tree_id {object} -- Which tree the priorities relate to (default: {None})
  
          Returns:
              int -- node covering the offset
          """
  
          tree_id = tree_id if tree_id is not None else self.main_tree
  
          left = 2 * idx + 1
          right = left + 1
          if left >= len(self.trees[tree_id]):
              return idx
  
          if s <= self.trees[tree_id][left]:
              return self._retrieve(left, s, tree_id)
          else:
              return self._retrieve(right, s - self.trees[tree_id][left],
                                    tree_id)
  
      def total(self, tree_id=None):
          """Returns the tree's total priority
  
          Keyword Arguments:
              tree_id {object} -- Tree identifier (default: {None})
  
          Returns:
              float -- Total priority
          """
  
          tree_id = tree_id if tree_id is not None else self.main_tree
  
          return self.trees[tree_id][0]
  
      def average(self, tree_id=None):
          """Return the tree's average priority, assumes the tree is full
  
          Keyword Arguments:
              tree_id {object} -- Tree identifier (default: {None})
  
          Returns:
              float -- Average priority
          """
  
          return self.total(tree_id) / self.capacity
  
      def add(self, priorities, data, write=None):
          """Adds a data sample to the SumTree at position write
  
          Arguments:
              priorities {dict} -- Dictionary of priorities, one key per tree
              data {tuple} -- Transition to be added
  
          Keyword Arguments:
              write {int} -- Position to write to (default: {None})
  
          Returns:
              tuple -- Tuple containing the replaced data as well as the node's
                          index in the tree
          """
  
          write = write if write is not None else self.write
          idx = write + self.capacity - 1
  
          # Save replaced data to eventually save in secondary memory
          replaced_data = np.copy(self.data[write])
          replaced_priorities = {
              tree: np.copy(self.trees[tree][idx])
              for tree in self.trees
          }
          replaced = (replaced_data, replaced_priorities)
  
          # Set new priorities
          for i, p in priorities.items():
              self.update(idx, p, i)
  
          # Set new data
          self.data[write] = data
          return replaced, idx
  
      def update(self, idx, p, tree_id=None):
          """For a given index, update priorities for the given trees
  
          Arguments:
              idx {int} -- Node's position in the tree
              p {dict|float} -- Dictionary of priorities or priority for the given tree_id
  
          Keyword Arguments:
              tree_id {object} -- Tree to be updated (default: {None})
          """
  
          if type(p) == dict:
              for k in p:
                  self.update(idx, p[k], k)
              return
          tree_id = tree_id if tree_id is not None else self.main_tree
  
          change = p - self.trees[tree_id][idx]
  
          self.trees[tree_id][idx] = p
          self._propagate(idx, change, tree_id)
  
      def get(self, s, tree_id=None):
          """Get the node covering the given offset
  
          Arguments:
              s {float} -- Offset to retrieve
  
          Keyword Arguments:
              tree_id {object} -- Tree to retrieve from (default: {None})
  
          Returns:
              tuple -- Containing the index, the priority and the transition
          """
  
          tree_id = tree_id if tree_id is not None else self.main_tree
  
          idx = self._retrieve(0, s, tree_id)
  
          return self.get_by_id(idx, tree_id)
  
      def get_by_id(self, idx, tree_id=None):
          tree_id = tree_id if tree_id is not None else self.main_tree
          dataIdx = idx - self.capacity + 1
  
          return (idx, self.trees[tree_id][idx], self.data[dataIdx])
  
  
  class DiverseMemory(object):
      """Prioritized Replay Buffer with integrated secondary Diverse Replay Buffer
  
      Arguments:
          main_capacity {int} -- Normal prioritized replay capacity
  
      Keyword Arguments:
          sec_capacity {int} -- Size of the secondary diverse replay buffer, if 0, the buffer functions as a normal prioritized Replay Buffer (default: {0})
          trace_diversity {bool} -- Wether diversity should be enforced at trace-level (True) or at transition-level (False)
          crowding_diversity {bool} -- Wether a crowding distance is applied to compute diversity
          value_function {function} -- When applied to a trace, this function should return the trace's value to be used in the crowding distance computation
          e {float} -- epsilon to be added to errors (default: {0.01})
          a {int} -- Power to which the error will be raised, if a==0, functionality is reduced to a replay buffer without prioritization (default: {2})
  
      """
  
      def __init__(
              self,
              main_capacity,
              sec_capacity=0,
              trace_diversity=True,
              crowding_diversity=True,
              value_function=lambda trace, trace_id, memory_indices: np.random.random(
                  1),
              e=0.01,
              a=2):
  
          self.len = 0
          self.trace_diversity = trace_diversity
          self.value_function = value_function
          self.crowding_diversity = crowding_diversity
          self.e = e
          self.a = a
          self.capacity = main_capacity + sec_capacity
          self.tree = SumTree(self.capacity, e, a)
          self.main_capacity = main_capacity
          self.sec_capacity = sec_capacity
          self.secondary_traces = []
  
      def _getPriority(self, error):
          """Compute priority from error
  
          Arguments:
              error {float} -- error
  
          Returns:
              float -- priority
          """
  
          return (error + self.e)**self.a
  
      def _getError(self, priority):
          """Given a priority, computes the corresponding error
  
          Arguments:
              priority {float} -- priority
  
          Returns:
              float -- error
          """
  
          return priority**(1 / self.a) - self.e
  
      def dupe(self, trg_i, src_i):
          """Copies the tree src_i into a new tree trg_i
  
          Arguments:
              trg_i {object} -- target tree identifier
              src_i {object} -- source tree identifier
          """
  
          self.tree.copy_tree(trg_i, src_i)
  
      def main_mem_is_full(self):
          """
              Because of the circular way in which we fill the memory, checking
              whether the current write position is free is sufficient to know
              if the memory is full.
          """
          return self.tree.data[self.tree.write][1] is not None
  
      def extract_trace(self, start):
          """Determines the end of the trace starting at position start
  
          Arguments:
              start {int} -- Trace's starting position
  
          Returns:
              int -- The trace's end posiition
          """
          trace_id = self.tree.data[start][0]
  
          end = (start + 1) % self.main_capacity
  
          if not self.trace_diversity:
              return end
          if trace_id is not None:
              while self.tree.data[end][0] == trace_id:
                  end = (end + 1) % self.main_capacity
                  if end == start:
                      break
  
          return end
  
      def add(self, error, sample, trace_id=None, pred_idx=None, tree_id=None):
          """Add the sample to the replay buffer, with a priority proportional to its error
          If trace_id is provided, the sample and the other samples with the same id will be treated as a trace when determining diversity
  
          Arguments:
              error {float} -- Error
              sample {tuple} -- The transition to be stored
  
          Keyword Arguments:
              trace_id {object} -- The trace's identifier (default: {None})
              tree_id {object} -- The tree for which the error is relevant (default: {None})
  
          Returns:
              int -- The index of the node in which the sample was stored
          """
  
          self.len = min(self.len + 1, self.capacity)
          self.tree.create(tree_id)
  
          sample = (trace_id, sample, None
                    if pred_idx is None else (pred_idx - self.capacity + 1))
  
          # Free up space in main memory if necessary
          if self.main_mem_is_full():
              end = self.extract_trace(self.tree.write)
              self.move_to_sec(self.tree.write, end)
  
          # Save sample into main memory
          if type(error) is not dict:
              error = {tree_id: error}
          idx = self.add_sample(sample, error, self.tree.write)
          self.tree.write = (self.tree.write + 1) % self.main_capacity
  
          return idx
  
      def remove_trace(self, trace):
          """Removes the trace from the main memory
  
          Arguments:
              trace {list} -- List of indices for the trace
          """
  
          _, trace_idx = trace
          for i in trace_idx:
              self.tree.data[i] = (None, None, None)
  
              idx = i + self.tree.capacity - 1
              for tree in self.tree.trees:
                  self.tree.update(idx, 0, tree)
  
      def get_trace_value(self, trace_tuple):
          """Applies the value_function to the trace's data to compute its value
  
          Arguments:
              trace_tuple {tuple} -- Tuple containing the trace and the trace's indices
  
          Returns:
              np.array -- The trace's value
          """
  
          trace, write_indices = trace_tuple
          if not self.trace_diversity:
              assert len(trace) == 1
          trace_id = trace[0][0]
          trace_data = [t[1] for t in trace]
  
          return self.value_function(trace_data, trace_id, write_indices)
  
      def sec_distances(self, traces):
          """Give a set of traces, this method computes each trace's crowding distance
  
          Arguments:
              traces {list} -- List of trace tuples
  
          Returns:
              list -- List of distances
          """
  
          values = [self.get_trace_value(tr) for tr in traces]
          if self.crowding_diversity:
              distances = crowd_dist(values)
          else:
              distances = values
          return [(i, d) for i, d in enumerate(distances)], values
  
      def get_sec_write(self, secondary_traces, trace, reserved_idx=None):
          """
              Given a trace, find free spots in the secondary memory to store it
              by recursively removing past traces with a low crowding distance
          """
  
          if reserved_idx is None:
              reserved_idx = []
              
          if len(trace) > self.sec_capacity:
              return None
              
          if len(reserved_idx) >= len(trace):
              return reserved_idx[:len(trace)]
  
          # Find free spots in the secondary memory
          # TODO: keep track of free spots so recomputation isn't necessary
          free_spots = [
              i + self.main_capacity for i in range(self.sec_capacity)
              if (self.tree.data[self.main_capacity + i][1]) is None
          ]
          
          if len(free_spots) > len(reserved_idx):
              return self.get_sec_write(secondary_traces, trace,
                                        free_spots[:len(trace)])
  
          # Get crowding distance of traces stored in the secondary buffer
          idx_dist, _ = self.sec_distances(secondary_traces)
  
          # Highest density = lowest distance
          i, _ = min(idx_dist, key=lambda d: d[1])
  
          _, trace_idx = secondary_traces[i]
          reserved_idx += trace_idx
  
          self.remove_trace(secondary_traces[i])
  
          del secondary_traces[i]
          return self.get_sec_write(secondary_traces, trace, reserved_idx)
  
      def move_to_sec(self, start, end):
          """Move the trace spanning from start to end to the secondary replay
          buffer
  
          Arguments:
              start {int} -- Start position of the trace
              end {int} -- End position of the trace
          """
  
          # Recover trace that needs to be moved
          if end <= start:
              indices = np.r_[start:self.main_capacity, 0:end]
          else:
              indices = np.r_[start:end]
          if not self.trace_diversity:
              assert len(indices) == 1
          trace = np.copy(self.tree.data[indices])
          priorities = {
              tree_id: self.tree.trees[tree_id][indices + self.tree.capacity - 1]
              for tree_id in self.tree.trees
          }
  
          # Get destination indices in secondary memory
          write_indices = self.get_sec_write(self.secondary_traces, trace)
  
          # Move trace to secondary memory if enough space was freed
          if write_indices is not None and len(write_indices) >= len(trace):
              for i, (w, t) in enumerate(zip(write_indices, trace)):
  
                  self.tree.data[w] = t
  
                  idx = w + self.tree.capacity - 1
                  for tree_id in priorities:
                      p = priorities[tree_id][i]
                      self.tree.update(idx, p, tree_id)
  
                  if i > 0:
                      self.tree.data[w][2] = write_indices[i - 1]
              if not self.trace_diversity:
                  assert len(trace) == 1
              self.secondary_traces.append((trace, write_indices))
          # elif self.sec_capacity>0:
          #     print("No space found for trace", trace[0][0],", discarding...",file=sys.stderr)
  
          # Remove trace from main memory
          self.remove_trace((None, indices))
  
  
      def add_sample(self, transition, error, write=None):
          """Stores the transition into the priority tree
          """
  
          p = {k: self._getPriority(error[k]) for k in error}
          _, idx = self.tree.add(p, transition, write=write)
          return idx
  
      def add_tree(self, tree_id):
          """Adds a secondary priority tree
  
          Arguments:
              tree_id {object} -- The secondary tree's id
          """
  
          self.tree.create(tree_id)
  
      def get_data(self, include_indices=False):
          """Get all the data stored in the replay buffer 
  
          Keyword Arguments:
              include_indices {bool} -- Whether to include each sample's position in the replay buffer (default: {False})
  
          Returns:
              The data
          """
  
          all_data = list(np.arange(self.capacity) + self.capacity - 1), list(
              self.tree.data)
          indices = []
          data = []
          for i, d in zip(all_data[0], all_data[1]):
              if (d[1]) is not None:
                  indices.append(i)
                  data.append(d[1])
          if include_indices:
              return indices, data
          else:
              return data
  
      def sample(self, n, tree_id=None):
          """Sample n transitions from the replay buffer, following the priorities
          of the tree identified by tree_id
  
          Arguments:
              n {int} -- Number of transitios to sample
  
          Keyword Arguments:
              tree_id {object} -- identifier of the tree whose priorities should be followed (default: {None})
  
          Returns:
              tuple -- pair of (indices, transitions, priorities)
          """
  
          if n < 1:
              return None, None, None
          batch = np.zeros((n, ), dtype=np.ndarray)
          ids = np.zeros(n, dtype=int)
          priorities = np.zeros(n, dtype=float)
          segment = self.tree.total(tree_id) / n
          for i in range(n):
              a = segment * i
              b = segment * (i + 1)
  
              s = np.random.uniform(a, b)
              (idx, p, data) = self.tree.get(s, tree_id)
              while (data[1]) is None or (
                      idx - self.capacity + 1 >= self.capacity):
                  s = np.random.uniform(0, self.tree.total(tree_id))
                  (idx, p, data) = self.tree.get(s, tree_id)
              ids[i] = idx
              batch[i] = data[1]
              priorities[i] = p
          return ids, batch, priorities
  
      def update(self, idx, error, tree_id=None):
          """Given a node's idx, this method updates the corresponding priority in
          the tree identified by tree_id
  
          Arguments:
              idx {int} -- Node's index
              error {float} -- New error
  
          Keyword Arguments:
              tree_id {object} -- Identifies the tree to update (default: {None})
          """
  
          if tree_id is None:
              tree_id = self.tree.main_tree
          if type(error) is not dict:
              error = {tree_id: error}
          p = {k: self._getPriority(error[k]) for k in error}
          self.tree.update(idx, p, tree_id)
  
      def get(self, indices):
          """Given a list of node indices, this method returns the data stored at
          those indices
  
          Arguments:
              indices {list} -- List of indices
  
          Returns:
              np.array -- array of transitions
          """
  
          indices = np.array(indices, dtype=int) - self.capacity + 1
          return self.tree.data[indices][:, 1]
  
      def get_error(self, idx, tree_id=None):
          tree_id = self.tree.main_tree if tree_id is None else tree_id
          priority = self.tree.trees[tree_id][idx]
          return self._getError(priority)
  