# coding=utf-8
# Copyright 2020 The TF-Agents Authors.
#
# 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
#
#     https://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.

"""Uniform replay buffer in Python.

The base class provides all the functionalities of a uniform replay buffer:
  - add samples in a First In First Out way.
  - read samples uniformly.

PyHashedReplayBuffer is a flavor of the base class which
compresses the observations when the observations have some partial overlap
(e.g. when using frame stacking).
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import threading

import numpy as np
import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import
from tf_agents.replay_buffers import replay_buffer
from tf_agents.specs import array_spec
from tf_agents.utils import nest_utils
from tf_agents.utils import numpy_storage


class PyUniformReplayBuffer(replay_buffer.ReplayBuffer):
  """A Python-based replay buffer that supports uniform sampling.

  Writing and reading to this replay buffer is thread safe.

  This replay buffer can be subclassed to change the encoding used for the
  underlying storage by overriding _encoded_data_spec, _encode, _decode, and
  _on_delete.
  """

  def __init__(self, data_spec, capacity):
    """Creates a PyUniformReplayBuffer.

    Args:
      data_spec: An ArraySpec or a list/tuple/nest of ArraySpecs describing a
        single item that can be stored in this buffer.
      capacity: The maximum number of items that can be stored in the buffer.
    """
    super(PyUniformReplayBuffer, self).__init__(data_spec, capacity)

    self._storage = numpy_storage.NumpyStorage(
        self._encoded_data_spec(), capacity
    )
    self._lock = threading.Lock()
    self._np_state = numpy_storage.NumpyState()

    # Adding elements to the replay buffer is done in a circular way.
    # Keeps track of the actual size of the replay buffer and the location
    # where to add new elements.
    self._np_state.size = np.int64(0)
    self._np_state.cur_id = np.int64(0)

    # Total number of items that went through the replay buffer.
    self._np_state.item_count = np.int64(0)

  def _encoded_data_spec(self):
    """Spec of data items after encoding using _encode."""
    return self._data_spec

  def _encode(self, item):
    """Encodes an item (before adding it to the buffer)."""
    return item

  def _decode(self, item):
    """Decodes an item."""
    return item

  def _on_delete(self, encoded_item):
    """Do any necessary cleanup."""
    pass

  @property
  def size(self):
    return self._np_state.size

  def _num_frames(self):
    raise NotImplementedError(
        'num_frames is not yet implemented in PyUniformReplayBuffer'
    )

  def _add_batch(self, items):
    outer_shape = nest_utils.get_outer_array_shape(items, self._data_spec)
    if outer_shape[0] != 1:
      raise NotImplementedError(
          'PyUniformReplayBuffer only supports a batch '
          'size of 1, but received `items` with batch '
          'size {}.'.format(outer_shape[0])
      )

    item = nest_utils.unbatch_nested_array(items)
    with self._lock:
      if self._np_state.size == self._capacity:
        # If we are at capacity, we are deleting element cur_id.
        self._on_delete(self._storage.get(self._np_state.cur_id))
      self._storage.set(self._np_state.cur_id, self._encode(item))
      self._np_state.size = np.minimum(self._np_state.size + 1, self._capacity)
      self._np_state.cur_id = (self._np_state.cur_id + 1) % self._capacity
      self._np_state.item_count += 1

  def _get_next(
      self, sample_batch_size=None, num_steps=None, time_stacked=True
  ):
    num_steps_value = num_steps if num_steps is not None else 1

    def get_single():
      """Gets a single item from the replay buffer."""
      with self._lock:
        if self._np_state.size <= 0:

          def empty_item(spec):
            return np.empty(spec.shape, dtype=spec.dtype)

          if num_steps is not None:
            item = [
                tf.nest.map_structure(empty_item, self.data_spec)
                for n in range(num_steps)
            ]
            if time_stacked:
              item = nest_utils.stack_nested_arrays(item)
          else:
            item = tf.nest.map_structure(empty_item, self.data_spec)
          return item
        idx = np.random.randint(self._np_state.size - num_steps_value + 1)
        if self._np_state.size == self._capacity:
          # If the buffer is full, add cur_id (head of circular buffer) so that
          # we sample from the range [cur_id, cur_id + size - num_steps_value].
          # We will modulo the size below.
          idx += self._np_state.cur_id

        if num_steps is not None:
          # TODO(b/120242830): Try getting data from numpy in one shot rather
          # than num_steps_value.
          item = [
              self._decode(self._storage.get((idx + n) % self._capacity))
              for n in range(num_steps)
          ]
        else:
          item = self._decode(self._storage.get(idx % self._capacity))

      if num_steps is not None and time_stacked:
        item = nest_utils.stack_nested_arrays(item)
      return item

    if sample_batch_size is None:
      return get_single()
    else:
      samples = [get_single() for _ in range(sample_batch_size)]
      return nest_utils.stack_nested_arrays(samples)

  def _as_dataset(
      self,
      sample_batch_size=None,
      num_steps=None,
      sequence_preprocess_fn=None,
      num_parallel_calls=None,
  ):
    if sequence_preprocess_fn is not None:
      raise NotImplementedError('sequence_preprocess_fn is not supported.')
    if num_parallel_calls is not None:
      raise NotImplementedError(
          'PyUniformReplayBuffer does not support '
          'num_parallel_calls (must be None).'
      )

    data_spec = self._data_spec
    if sample_batch_size is not None:
      data_spec = array_spec.add_outer_dims_nest(
          data_spec, (sample_batch_size,)
      )
    if num_steps is not None:
      data_spec = (data_spec,) * num_steps
    shapes = tuple(s.shape for s in tf.nest.flatten(data_spec))
    dtypes = tuple(s.dtype for s in tf.nest.flatten(data_spec))

    def generator_fn():
      while True:
        if sample_batch_size is not None:
          batch = [
              self._get_next(num_steps=num_steps, time_stacked=False)
              for _ in range(sample_batch_size)
          ]
          item = nest_utils.stack_nested_arrays(batch)
        else:
          item = self._get_next(num_steps=num_steps, time_stacked=False)
        yield tuple(tf.nest.flatten(item))

    def time_stack(*structures):
      time_axis = 0 if sample_batch_size is None else 1
      return tf.nest.map_structure(
          lambda *elements: tf.stack(elements, axis=time_axis), *structures
      )

    ds = tf.data.Dataset.from_generator(generator_fn, dtypes, shapes).map(
        lambda *items: tf.nest.pack_sequence_as(data_spec, items)
    )
    if num_steps is not None:
      return ds.map(time_stack)
    else:
      return ds

  def _gather_all(self):
    data = [
        self._decode(self._storage.get(idx)) for idx in range(self._capacity)
    ]
    stacked = nest_utils.stack_nested_arrays(data)
    batched = tf.nest.map_structure(lambda t: np.expand_dims(t, 0), stacked)
    return batched

  def _clear(self):
    self._np_state.size = np.int64(0)
    self._np_state.cur_id = np.int64(0)
