# 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 entity class."""

import abc
import collections
import os
import weakref

from absl import logging
from dm_control import mjcf
from dm_control.composer import define
from dm_control.mujoco.wrapper import mjbindings
import numpy as np

_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):
  result = np.empty_like(quat1)
  mjbindings.mjlib.mju_mulQuat(result, quat1, quat2)
  return result


def _rotate_vector(vec, quat):
  """Rotates a vector by the given quaternion."""
  result = np.empty_like(vec)
  mjbindings.mjlib.mju_rotVecQuat(result, vec, quat)
  return result


class _ObservableKeys:
  """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 as exc:
      raise ValueError(
          'cannot retrieve the full identifier of mjcf_model') from exc
    return os.path.join(model_identifier, name)

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


class Observables:
  """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 as exc:
        raise ValueError(
            'Cannot retrieve the full identifier of mjcf_model.') from exc

      return collections.OrderedDict(
          [(os.path.join(model_identifier, name), observable)
           for name, observable in self._observables.items()])
    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 as exc:
        raise ValueError(
            'Cannot retrieve the full identifier of mjcf_model.') from exc
      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 options.items():
      try:
        obs = self._observables[obs_key]
      except KeyError as exc:
        raise KeyError('No observable with name {!r}'.format(obs_key)) from exc
      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


class FreePropObservableMixin(metaclass=abc.ABCMeta):
  """Enforce observables of a free-moving object."""

  @property
  @abc.abstractmethod
  def position(self):
    pass

  @property
  @abc.abstractmethod
  def orientation(self):
    pass

  @property
  @abc.abstractmethod
  def linear_velocity(self):
    pass

  @property
  @abc.abstractmethod
  def angular_velocity(self):
    pass


class Entity(metaclass=abc.ABCMeta):
  """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, exclude_self=False):
    """An iterator that recursively iterates through all attached entities.

    Args:
      exclude_self: (optional) Whether to exclude this `Entity` itself from the
        iterator.

    Yields:
      If `exclude_self` is `False`, 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.
    """
    if not exclude_self:
      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

  @property
  @abc.abstractmethod
  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()  # pylint: disable=not-callable
      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  # pylint: disable=not-callable

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

  @property
  def root_body(self):
    if self.parent:
      return mjcf.get_attachment_frame(self.mjcf_model)
    else:
      return self.mjcf_model.worldbody

  def global_vector_to_local_frame(self, physics, vec_in_world_frame):
    """Linearly transforms a world-frame vector into entity's local frame.

    Note that this function does not perform an affine transformation of the
    vector. In other words, the input vector is assumed to be specified with
    respect to the same origin as this entity's local frame. This function
    can also be applied to matrices whose innermost dimensions are either 2 or
    3. In this case, a matrix with the same leading dimensions is returned
    where the innermost vectors are replaced by their values computed in the
    local frame.

    Args:
      physics: An `mjcf.Physics` instance.
      vec_in_world_frame: A NumPy array with last dimension of shape (2,) or
      (3,) that represents a vector quantity in the world frame.

    Returns:
      The same quantity as `vec_in_world_frame` but reexpressed in this
      entity's local frame. The returned np.array has the same shape as
      np.asarray(vec_in_world_frame).

    Raises:
      ValueError: if `vec_in_world_frame` does not have shape ending with (2,)
        or (3,).
    """
    vec_in_world_frame = np.asarray(vec_in_world_frame)

    xmat = np.reshape(physics.bind(self.root_body).xmat, (3, 3))
    # The ordering of the np.dot is such that the transformation holds for any
    # matrix whose final dimensions are (2,) or (3,).
    if vec_in_world_frame.shape[-1] == 2:
      return np.dot(vec_in_world_frame, xmat[:2, :2])
    elif vec_in_world_frame.shape[-1] == 3:
      return np.dot(vec_in_world_frame, xmat)
    else:
      raise ValueError('`vec_in_world_frame` should have shape with final '
                       'dimension 2 or 3: got {}'.format(
                           vec_in_world_frame.shape))

  def global_xmat_to_local_frame(self, physics, xmat):
    """Transforms another entity's `xmat` into this entity's local frame.

    This function takes another entity's (E) xmat, which is an SO(3) matrix
    from E's frame to the world frame, and turns it to a matrix that transforms
    from E's frame into this entity's local frame.

    Args:
      physics: An `mjcf.Physics` instance.
      xmat: A NumPy array of shape (3, 3) or (9,) that represents another
        entity's xmat.

    Returns:
      The `xmat` reexpressed in this entity's local frame. The returned
      np.array has the same shape as np.asarray(xmat).

    Raises:
      ValueError: if `xmat` does not have shape (3, 3) or (9,).
    """
    xmat = np.asarray(xmat)

    input_shape = xmat.shape
    if xmat.shape == (9,):
      xmat = np.reshape(xmat, (3, 3))

    self_xmat = np.reshape(physics.bind(self.root_body).xmat, (3, 3))
    if xmat.shape == (3, 3):
      return np.reshape(np.dot(self_xmat.T, xmat), input_shape)
    else:
      raise ValueError('`xmat` should have shape (3, 3) or (9,): got {}'.format(
          xmat.shape))

  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:
        normalised_quaternion = quaternion / np.linalg.norm(quaternion)
        physics.bind(root_joint).qpos[3:] = normalised_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,
                 rotate_velocity=False):
    """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.
      rotate_velocity: (optional) A bool, whether to shift the current linear
        velocity along with the pose. This will rotate the current linear
        velocity, which is expressed relative to the world frame. The angular
        velocity, which is expressed relative to the local frame is left
        unchanged.

    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:
      quaternion = np.asarray(quaternion, dtype=np.float64)
      new_quaternion = _multiply_quaternions(quaternion, current_quaternion)
      root_joint = mjcf.get_frame_freejoint(self.mjcf_model)
      if root_joint and rotate_velocity:
        # Rotate the linear velocity. The angular velocity (qvel[3:])
        # is left unchanged, as it is expressed in the local frame.
        # When rotatating the body frame the angular velocity already
        # tracks the rotation but the linear velocity does not.
        velocity = physics.bind(root_joint).qvel[:3]
        rotated_velocity = _rotate_vector(velocity, quaternion)
        self.set_velocity(physics, rotated_velocity)
    self.set_pose(physics, new_position, new_quaternion)

  def get_velocity(self, physics):
    """Gets the linear and angular velocity of this free entity.

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

    Returns:
      A 2-tuple where the first entry is a (3,) numpy array representing the
      linear velocity and the second is a (3,) numpy array representing the
      angular velocity.

    """
    root_joint = mjcf.get_frame_freejoint(self.mjcf_model)
    if root_joint:
      velocity = physics.bind(root_joint).qvel[:3]
      angular_velocity = physics.bind(root_joint).qvel[3:]
      return velocity, angular_velocity
    else:
      raise ValueError('get_velocity cannot be used on a non-free entity')

  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.')

  def configure_joints(self, physics, position):
    """Configures this entity's internal joints.

    The default implementation of this method simply sets the `qpos` of all
    joints in this entity to the values specified in the `position` argument.
    Entity subclasses with actuated joints may override this method to achieve a
    stable reconfiguration of joint positions, for example the control signal
    of position actuators may be changed to match the new joint positions.

    Args:
      physics: An instance of `mjcf.Physics`.
      position: The desired position of this entity's joints.
    """
    joints = self.mjcf_model.find_all('joint', exclude_attachments=True)
    physics.bind(joints).qpos = position


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