# 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.
# ============================================================================

"""Module defining the abstract robot class."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import abc
import collections
import os
import weakref

from absl import logging
from dm_control import mjcf
from dm_control.composer import define
import numpy as np
import six

_OPTION_KEYS = set(['update_interval', 'buffer_size', 'delay', 'aggregator',
                    'corruptor', 'enabled'])

_NO_ATTACHMENT_FRAME = 'No attachment frame found.'


# The component order differs from that used by the open-source `tf` package.
def _multiply_quaternions(quat1, quat2):
  """Multiplies two quaternions, expressed as [w, i, j, k]."""
  return np.array([
      [+quat1[0], -quat1[1], -quat1[2], -quat1[3]],
      [+quat1[1], +quat1[0], -quat1[3], +quat1[2]],
      [+quat1[2], +quat1[3], +quat1[0], -quat1[1]],
      [+quat1[3], -quat1[2], +quat1[1], +quat1[0]]]).dot(quat2)


class _ObservableKeys(object):
  """Helper object that implements the `observables.dict_keys` functionality."""

  def __init__(self, entity, observables):
    self._entity = entity
    self._observables = observables

  def __getattr__(self, name):
    try:
      model_identifier = self._entity.mjcf_model.full_identifier
    except AttributeError:
      raise ValueError('cannot retrieve the full identifier of mjcf_model')
    return os.path.join(model_identifier, name)

  def __dir__(self):
    out = set(self._observables.keys())
    out.update(dir(super(_ObservableKeys, self)))
    return list(out)


class Observables(object):
  """Base-class for Entity observables.

  Subclasses should declare getter methods annotated with @define.observable
  decorator and returning an observable object.
  """

  def __init__(self, entity):
    self._entity = weakref.proxy(entity)

    self._observables = collections.OrderedDict()
    self._keys_helper = _ObservableKeys(self._entity, self._observables)

    # Ensure consistent ordering.
    for attr_name in sorted(dir(type(self))):
      type_attr = getattr(type(self), attr_name)
      if isinstance(type_attr, define.observable):
        self._observables[attr_name] = getattr(self, attr_name)

  @property
  def dict_keys(self):
    return self._keys_helper

  def as_dict(self, fully_qualified=True):
    """Returns an OrderedDict of observables belonging to this Entity.

    The returned observables will include any added using the _add_observable
    method, as well as any generated by a method decorated with the
    @define.observable annotation.

    Args:
      fully_qualified: (bool) Whether the dict keys should be prefixed with the
        parent entity's full model identifier.
    """

    if fully_qualified:
      # We need to make sure that this property doesn't raise an AttributeError,
      # otherwise __getattr__ is executed and we get a very funky error.
      try:
        model_identifier = self._entity.mjcf_model.full_identifier
      except AttributeError:
        raise ValueError('cannot retrieve the full identifier of mjcf_model')

      return collections.OrderedDict(
          [(os.path.join(model_identifier, name), observable)
           for name, observable in six.iteritems(self._observables)])
    else:
      # Return a copy to prevent dict being edited.
      return self._observables.copy()

  def get_observable(self, name, name_fully_qualified=False):
    """Returns the observable with the given name.

    Args:
      name: (str) The identifier of the observable.
      name_fully_qualified: (bool) Whether the provided name is prefixed by the
        model's full identifier.
    """

    if name_fully_qualified:
      try:
        model_identifier = self._entity.mjcf_model.full_identifier
      except AttributeError:
        raise ValueError('cannot retrieve the full identifier of mjcf_model')
      return self._observables[name.replace(model_identifier, '')]
    else:
      return self._observables[name]

  def set_options(self, options):
    """Configure Observables with an options dict.

    Args:
      options: A dict of dicts of configuration options keyed on
        observable names, or a dict of configuration options, which will
        propagate those options to all observables.
    """
    if options is None:
      options = {}
    elif options.keys() and set(options.keys()).issubset(_OPTION_KEYS):
      options = dict([(key, options) for key in self._observables.keys()])

    for obs_key, obs_options in six.iteritems(options):
      try:
        obs = self._observables[obs_key]
      except KeyError:
        raise KeyError('No observable with name {!r}'.format(obs_key))
      obs.configure(**obs_options)

  def enable_all(self):
    """Enable all observables of this entity."""
    for obs in self._observables.values():
      obs.enabled = True

  def disable_all(self):
    """Disable all observables of this entity."""
    for obs in self._observables.values():
      obs.enabled = False

  def add_observable(self, name, observable, enabled=True):
    self._observables[name] = observable
    self._observables[name].enabled = enabled


@six.add_metaclass(abc.ABCMeta)
class FreePropObservableMixin(object):
  """Enforce observables of a free-moving object."""

  @abc.abstractproperty
  def position(self):
    pass

  @abc.abstractproperty
  def orientation(self):
    pass

  @abc.abstractproperty
  def linear_velocity(self):
    pass

  @abc.abstractproperty
  def angular_velocity(self):
    pass


@six.add_metaclass(abc.ABCMeta)
class Entity(object):
  """The abstract base class for an entity in a Composer environment."""

  def __init__(self, *args, **kwargs):
    """Entity constructor.

    Subclasses should not override this method, instead implement a _build
    method.

    Args:
      *args: Arguments passed through to the _build method.
      **kwargs: Keyword arguments. Passed through to the _build method, apart
        from the following.
          `observable_options`: A dictionary of Observable
            configuration options.
    """
    self._post_init_hooks = []

    self._parent = None
    self._attached = []

    try:
      observable_options = kwargs.pop('observable_options')
    except KeyError:
      observable_options = None

    self._build(*args, **kwargs)
    self._observables = self._build_observables()
    self._observables.set_options(observable_options)

  @abc.abstractmethod
  def _build(self, *args, **kwargs):
    """Entity initialization method to be overridden by subclasses."""
    raise NotImplementedError

  def _build_observables(self):
    """Entity observables initialization method.

    Returns:
      An object subclassing the Observables class.
    """
    return Observables(self)

  def iter_entities(self):
    """An iterator that recursively iterates through all attached entities.

    Yields:
      The first value yielded is this Entity itself. The following Entities are
      then yielded recursively in a depth-first fashion, following the order in
      which the Entities are attached.
    """
    yield self
    for attached_entity in self._attached:
      for attached_entity_of_attached_entity in attached_entity.iter_entities():
        yield attached_entity_of_attached_entity

  @property
  def observables(self):
    """The observables defined by this entity."""
    return self._observables

  def initialize_episode_mjcf(self, random_state):
    """Callback executed when the MJCF model is modified between episodes."""
    pass

  def after_compile(self, physics, random_state):
    """Callback executed after the Mujoco Physics is recompiled."""
    pass

  def initialize_episode(self, physics, random_state):
    """Callback executed during episode initialization."""
    pass

  def before_step(self, physics, random_state):
    """Callback executed before an agent control step."""
    pass

  def before_substep(self, physics, random_state):
    """Callback executed before a simulation step."""
    pass

  def after_substep(self, physics, random_state):
    """A callback which is executed after a simulation step."""
    pass

  def after_step(self, physics, random_state):
    """Callback executed after an agent control step."""
    pass

  @abc.abstractproperty
  def mjcf_model(self):
    raise NotImplementedError

  def attach(self, entity, attach_site=None):
    """Attaches an `Entity` without any additional degrees of freedom.

    Args:
      entity: The `Entity` to attach.
      attach_site: (optional) The site to which to attach the entity's model. If
          not set, defaults to self.attachment_site.

    Returns:
      The frame of the attached model.
    """

    if attach_site is None:
      attach_site = self.attachment_site

    frame = attach_site.attach(entity.mjcf_model)
    self._attached.append(entity)
    entity._parent = weakref.ref(self)  # pylint: disable=protected-access
    return frame

  def detach(self):
    """Detaches this entity if it has previously been attached."""
    if self._parent is not None:
      parent = self._parent()
      if parent:  # Weakref might dereference to None during garbage collection.
        self.mjcf_model.detach()
        parent._attached.remove(self)  # pylint: disable=protected-access
      self._parent = None
    else:
      raise RuntimeError('Cannot detach an entity that is not attached.')

  @property
  def parent(self):
    """Returns the `Entity` to which this entity is attached, or `None`."""
    return self._parent() if self._parent else None

  @property
  def attachment_site(self):
    return self._mjcf_root

  def get_pose(self, physics):
    """Get the position and orientation of this entity relative to its parent.

    Note that the semantics differ slightly depending on whether or not the
    entity has a free joint:

    * If it has a free joint the position and orientation are always given in
      global coordinates.
    * If the entity is fixed or attached with a different joint type then the
      position and orientation are given relative to the parent frame.

    For entities that are either attached directly to the worldbody, or to other
    entities that are positioned at the global origin (e.g. the arena) the
    global and relative poses are equivalent.

    Args:
      physics: An instance of `mjcf.Physics`.

    Returns:
      A 2-tuple where the first entry is a (3,) numpy array representing the
      position and the second is a (4,) numpy array representing orientation as
      a quaternion.

    Raises:
      RuntimeError: If the entity is not attached.
    """
    root_joint = mjcf.get_frame_freejoint(self.mjcf_model)
    if root_joint:
      position = physics.bind(root_joint).qpos[:3]
      quaternion = physics.bind(root_joint).qpos[3:]
    else:
      attachment_frame = mjcf.get_attachment_frame(self.mjcf_model)
      if attachment_frame is None:
        raise RuntimeError(_NO_ATTACHMENT_FRAME)
      position = physics.bind(attachment_frame).pos
      quaternion = physics.bind(attachment_frame).quat
    return position, quaternion

  def set_pose(self, physics, position=None, quaternion=None):
    """Sets position and/or orientation of this entity relative to its parent.

    If the entity is attached with a free joint, this method will set the
    respective DoFs of the joint. If the entity is either fixed or attached with
    a different joint type, this method will update the position and/or
    quaternion of the attachment frame.

    Note that the semantics differ slightly between the two cases: the DoFs of a
    free body are specified in global coordinates, whereas the position of a
    non-free body is specified in relative coordinates with respect to the
    parent frame. However, for entities that are either attached directly to the
    worldbody, or to other entities that are positioned at the global origin
    (e.g. the arena), there is no difference between the two cases.

    Args:
      physics: An instance of `mjcf.Physics`.
      position: (optional) A NumPy array of size 3.
      quaternion: (optional) A NumPy array of size 4.

    Raises:
      RuntimeError: If the entity is not attached.
    """
    root_joint = mjcf.get_frame_freejoint(self.mjcf_model)
    if root_joint:
      if position is not None:
        physics.bind(root_joint).qpos[:3] = position
      if quaternion is not None:
        physics.bind(root_joint).qpos[3:] = quaternion
    else:
      attachment_frame = mjcf.get_attachment_frame(self.mjcf_model)
      if attachment_frame is None:
        raise RuntimeError(_NO_ATTACHMENT_FRAME)
      if position is not None:
        physics.bind(attachment_frame).pos = position
      if quaternion is not None:
        normalised_quaternion = quaternion / np.linalg.norm(quaternion)
        physics.bind(attachment_frame).quat = normalised_quaternion

  def shift_pose(self, physics, position=None, quaternion=None):
    """Shifts the position and/or orientation from its current configuration.

    This is a convenience function that performs the same operation as
    `set_pose`, but where the specified `position` is added to the current
    position, and the specified `quaternion` is premultiplied to the current
    quaternion.

    Args:
      physics: An instance of `mjcf.Physics`.
      position: (optional) A NumPy array of size 3.
      quaternion: (optional) A NumPy array of size 4.

    Raises:
      RuntimeError: If the entity is not attached.
    """
    current_position, current_quaternion = self.get_pose(physics)
    new_position, new_quaternion = None, None
    if position is not None:
      new_position = current_position + position
    if quaternion is not None:
      new_quaternion = _multiply_quaternions(quaternion, current_quaternion)
    self.set_pose(physics, new_position, new_quaternion)

  def set_velocity(self, physics, velocity=None, angular_velocity=None):
    """Sets the linear velocity and/or angular velocity of this free entity.

    If the entity is attached with a free joint, this method will set the
    respective DoFs of the joint. Otherwise a warning is logged.

    Args:
      physics: An instance of `mjcf.Physics`.
      velocity: (optional) A NumPy array of size 3 specifying the
        linear velocity.
      angular_velocity: (optional) A NumPy array of size 3 specifying the
        angular velocity
    """
    root_joint = mjcf.get_frame_freejoint(self.mjcf_model)
    if root_joint:
      if velocity is not None:
        physics.bind(root_joint).qvel[:3] = velocity
      if angular_velocity is not None:
        physics.bind(root_joint).qvel[3:] = angular_velocity
    else:
      logging.warning('Cannot set velocity on Entity with no free joint.')


class ModelWrapperEntity(Entity):
  """An entity class that wraps an MJCF model without any additional logic."""

  def _build(self, mjcf_model):
    self._mjcf_model = mjcf_model

  @property
  def mjcf_model(self):
    return self._mjcf_model
