# Copyright 2018 The dm_control 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
#
#    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.
# ============================================================================

"""An object that creates and updates buffers for enabled observables."""

import collections
import functools
from absl import logging

from dm_control.composer import variation
from dm_control.composer.observation import obs_buffer
from dm_env import specs
import numpy as np

DEFAULT_BUFFER_SIZE = 1
DEFAULT_UPDATE_INTERVAL = 1
DEFAULT_DELAY = 0


class _EnabledObservable:
  """Encapsulates an enabled observable, its buffer, and its update schedule."""

  __slots__ = ('observable', 'observation_callable',
               'update_interval', 'delay', 'buffer_size',
               'buffer', 'update_schedule')

  def __init__(self, observable, physics, random_state,
               strip_singleton_buffer_dim, pad_with_initial_value):
    self.observable = observable
    self.observation_callable = (
        observable.observation_callable(physics, random_state))

    self._bind_attribute_from_observable('update_interval',
                                         DEFAULT_UPDATE_INTERVAL,
                                         random_state)
    self._bind_attribute_from_observable('delay',
                                         DEFAULT_DELAY,
                                         random_state)
    self._bind_attribute_from_observable('buffer_size',
                                         DEFAULT_BUFFER_SIZE,
                                         random_state)

    obs_spec = self.observable.array_spec
    if obs_spec is None:
      # We take an observation to determine the shape and dtype of the array.
      # This occurs outside of an episode and doesn't affect environment
      # behavior. At this point the physics state is not guaranteed to be valid,
      # so we might get a `PhysicsError` if the observation callable calls
      # `physics.forward`. We suppress such errors since they do not matter as
      # far as the shape and dtype of the observation are concerned.
      with physics.suppress_physics_errors():
        obs_array = self.observation_callable()
      obs_array = np.asarray(obs_array)
      obs_spec = specs.Array(shape=obs_array.shape, dtype=obs_array.dtype)
    self.buffer = obs_buffer.Buffer(
        buffer_size=self.buffer_size,
        shape=obs_spec.shape, dtype=obs_spec.dtype,
        pad_with_initial_value=pad_with_initial_value,
        strip_singleton_buffer_dim=strip_singleton_buffer_dim)
    self.update_schedule = collections.deque()

  def _bind_attribute_from_observable(self, attr, default_value, random_state):
    obs_attr = getattr(self.observable, attr)
    if obs_attr:
      if isinstance(obs_attr, variation.Variation):
        setattr(self, attr,
                functools.partial(obs_attr, random_state=random_state))
      else:
        setattr(self, attr, obs_attr)
    else:
      setattr(self, attr, default_value)


def _call_if_callable(arg):
  if callable(arg):
    return arg()
  else:
    return arg


def _validate_structure(structure):
  """Validates the structure of the given observables collection.

  The collection must either be a dict, or a (list or tuple) of dicts.

  Args:
    structure: A candidate collection of observables.

  Returns:
    A boolean that is `True` if `structure` is either a list or a tuple, or
    `False` otherwise.

  Raises:
    ValueError: If `structure` is neither a dict nor a (list or tuple) of dicts.
  """
  is_nested = isinstance(structure, (list, tuple))
  if is_nested:
    is_valid = all(isinstance(obj, dict) for obj in structure)
  else:
    is_valid = isinstance(structure, dict)
  if not is_valid:
    raise ValueError(
        '`observables` should be a dict, or a (list or tuple) of dicts'
        ': got {}'.format(structure))
  return is_nested


class Updater:
  """Creates and updates buffers for enabled observables."""

  def __init__(self, observables, physics_steps_per_control_step=1,
               strip_singleton_buffer_dim=False,
               pad_with_initial_value=False):
    self._physics_steps_per_control_step = physics_steps_per_control_step
    self._strip_singleton_buffer_dim = strip_singleton_buffer_dim
    self._pad_with_initial_value = pad_with_initial_value
    self._step_counter = 0
    self._observables = observables
    self._is_nested = _validate_structure(observables)
    self._enabled_structure = None
    self._enabled_list = None

  def reset(self, physics, random_state):
    """Resets this updater's state."""

    def make_buffers_dict(observables):
      """Makes observable states in a dict."""
      # Use `type(observables)` so that our output structure respects the
      # original dict subclass (e.g. OrderedDict).
      out_dict = type(observables)()
      for key, value in observables.items():
        if value.enabled:
          out_dict[key] = _EnabledObservable(value, physics, random_state,
                                             self._strip_singleton_buffer_dim,
                                             self._pad_with_initial_value)
      return out_dict

    if self._is_nested:
      self._enabled_structure = type(self._observables)(
          make_buffers_dict(obs_dict) for obs_dict in self._observables)
      self._enabled_list = []
      for enabled_dict in self._enabled_structure:
        self._enabled_list.extend(enabled_dict.values())
    else:
      self._enabled_structure = make_buffers_dict(self._observables)
      self._enabled_list = self._enabled_structure.values()

    self._step_counter = 0
    for enabled in self._enabled_list:
      first_delay = _call_if_callable(enabled.delay)
      enabled.buffer.insert(
          0, first_delay,
          enabled.observation_callable())

  def observation_spec(self):
    """The observation specification for this environment.

    Returns a dict mapping the names of enabled observations to their
    corresponding `Array` or `BoundedArray` specs.

    If an obs has a BoundedArray spec, but uses an aggregator that
    does not preserve those bounds (such as `sum`), it will be mapped to an
    (unbounded) `Array` spec. If using a bounds-preserving custom aggregator
    `my_agg`, give it an attribute `my_agg.preserves_bounds = True` to indicate
    to this method that it is bounds-preserving.

    The returned specification is only valid as of the previous call
    to `reset`. In particular, it is an error to call this function before
    the first call to `reset`.

    Returns:
      A dict mapping observation name to `Array` or `BoundedArray` spec
      containing the observation shape and dtype, and possibly bounds.

    Raises:
      RuntimeError: If this method is called before `reset` has been called.
    """
    if self._enabled_structure is None:
      raise RuntimeError('`reset` must be called before `observation_spec`.')

    def make_observation_spec_dict(enabled_dict):
      """Makes a dict of enabled observation specs from of observables."""
      out_dict = type(enabled_dict)()
      for name, enabled in enabled_dict.items():

        if (enabled.observable.aggregator is None
            and enabled.observable.array_spec is not None):
          # If possible, keep the original array spec, just updating the name
          # and modifying the dimension for buffering. Doing this allows for
          # custom spec types to be exposed by the environment where possible.
          out_dict[name] = enabled.observable.array_spec.replace(
              name=name, shape=enabled.buffer.shape
          )
          continue

        if isinstance(enabled.observable.array_spec, specs.BoundedArray):
          bounds = (enabled.observable.array_spec.minimum,
                    enabled.observable.array_spec.maximum)
        else:
          bounds = None

        if enabled.observable.aggregator:
          aggregator = enabled.observable.aggregator
          aggregated = aggregator(np.zeros(enabled.buffer.shape,
                                           dtype=enabled.buffer.dtype))
          shape = aggregated.shape
          dtype = aggregated.dtype

          # Ditch bounds if the aggregator isn't known to be bounds-preserving.
          if bounds:
            if not hasattr(aggregator, 'preserves_bounds'):
              logging.warning('Ignoring the bounds of this observable\'s spec, '
                              'as its aggregator method has no boolean '
                              '`preserves_bounds` attrubute.')
              bounds = None
            elif not aggregator.preserves_bounds:
              bounds = None
        else:
          shape = enabled.buffer.shape
          dtype = enabled.buffer.dtype

        if bounds:
          spec = specs.BoundedArray(minimum=bounds[0],
                                    maximum=bounds[1],
                                    shape=shape,
                                    dtype=dtype,
                                    name=name)
        else:
          spec = specs.Array(shape=shape, dtype=dtype, name=name)

        out_dict[name] = spec
      return out_dict

    if self._is_nested:
      enabled_specs = type(self._enabled_structure)(
          make_observation_spec_dict(enabled_dict)
          for enabled_dict in self._enabled_structure)
    else:
      enabled_specs = make_observation_spec_dict(self._enabled_structure)

    return enabled_specs

  def prepare_for_next_control_step(self):
    """Simulates the next control step and optimizes the update schedule."""
    if self._enabled_structure is None:
      raise RuntimeError('`reset` must be called before `before_step`.')
    for enabled in self._enabled_list:

      if (enabled.update_interval == DEFAULT_UPDATE_INTERVAL
          and enabled.delay == DEFAULT_DELAY
          and enabled.buffer_size < self._physics_steps_per_control_step):
        for i in reversed(range(enabled.buffer_size)):
          next_step = (
              self._step_counter + self._physics_steps_per_control_step - i)
          next_delay = DEFAULT_DELAY
          enabled.update_schedule.append((next_step, next_delay))
      else:
        if enabled.update_schedule:
          last_scheduled_step = enabled.update_schedule[-1][0]
        else:
          last_scheduled_step = self._step_counter
        max_step = self._step_counter + 2 * self._physics_steps_per_control_step
        while last_scheduled_step < max_step:
          next_update_interval = _call_if_callable(enabled.update_interval)
          next_step = last_scheduled_step + next_update_interval
          next_delay = _call_if_callable(enabled.delay)
          enabled.update_schedule.append((next_step, next_delay))
          last_scheduled_step = next_step
        # Optimize the schedule by planning ahead and dropping unseen entries.
        enabled.buffer.drop_unobserved_upcoming_items(
            enabled.update_schedule, self._physics_steps_per_control_step)

  def update(self):
    if self._enabled_structure is None:
      raise RuntimeError('`reset` must be called before `after_substep`.')
    self._step_counter += 1
    for enabled in self._enabled_list:
      if (enabled.update_schedule and
          enabled.update_schedule[0][0] == self._step_counter):
        timestamp, delay = enabled.update_schedule.popleft()
        enabled.buffer.insert(
            timestamp, delay,
            enabled.observation_callable())

  def get_observation(self):
    """Gets the current observation.

    The returned observation is only valid as of the previous call
    to `reset`. In particular, it is an error to call this function before
    the first call to `reset`.

    Returns:
      A dict, or list of dicts, or tuple of dicts, of observation values.
      The returned structure corresponds to the structure of the `observables`
      that was given at initialization time.

    Raises:
      RuntimeError: If this method is called before `reset` has been called.
    """
    if self._enabled_structure is None:
      raise RuntimeError('`reset` must be called before `observation`.')

    def aggregate_dict(enabled_dict):
      out_dict = type(enabled_dict)()
      for name, enabled in enabled_dict.items():
        if enabled.observable.aggregator:
          aggregated = enabled.observable.aggregator(
              enabled.buffer.read(self._step_counter))
        else:
          aggregated = enabled.buffer.read(self._step_counter)
        out_dict[name] = aggregated
      return out_dict

    if self._is_nested:
      return type(self._enabled_structure)(
          aggregate_dict(enabled_dict)
          for enabled_dict in self._enabled_structure)
    else:
      return aggregate_dict(self._enabled_structure)
