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

"""A 2x4 Duplo brick."""

import collections
import os

from dm_control import composer
from dm_control import mjcf
from dm_control.composer import define
from dm_control.composer.observation import observable
import numpy as np

_DUPLO_XML_PATH = os.path.join(os.path.dirname(__file__), 'duplo2x4.xml')

# Stud radii are drawn from a uniform distribution. The `variation` argument
# scales the minimum and maximum whilst keeping the lower quartile constant.
_StudSize = collections.namedtuple(
    '_StudSize', ['minimum', 'lower_quartile', 'maximum'])
_StudParams = collections.namedtuple('_StudParams', ['easy_align', 'flanges'])

_STUD_SIZE_PARAMS = {
    _StudParams(easy_align=False, flanges=False):
        _StudSize(minimum=0.004685, lower_quartile=0.004781, maximum=0.004898),
    _StudParams(easy_align=False, flanges=True):
        _StudSize(minimum=0.004609, lower_quartile=0.004647, maximum=0.004716),
    _StudParams(easy_align=True, flanges=False):
        _StudSize(minimum=0.004754, lower_quartile=0.004844, maximum=0.004953),
    _StudParams(easy_align=True, flanges=True):
        _StudSize(minimum=0.004695, lower_quartile=0.004717, maximum=0.004765)
}

_COLOR_NOT_BETWEEN_0_AND_1 = (
    'All values in `color` must be between 0 and 1, got {!r}.')


class Duplo(composer.Entity):
  """A 2x4 Duplo brick."""

  def _build(self, easy_align=False, flanges=True, variation=0.0,
             color=(1., 0., 0.)):
    """Initializes a new `Duplo` instance.

    Args:
      easy_align: If True, the studs on the top of the brick will be capsules
        rather than cylinders. This makes alignment easier.
      flanges: Whether to use flanges on the bottom of the brick. These make the
        dynamics more expensive, but allow the bricks to be clicked together in
        partially overlapping configurations.
      variation: A float that controls the amount of variation in stud size (and
        therefore separation force). A value of 1.0 results in a distribution of
        separation forces that approximately matches the empirical distribution
        measured for real Duplo bricks. A value of 0.0 yields a deterministic
        separation force approximately equal to the mode of the empirical
        distribution.
      color: An optional tuple of (R, G, B) values specifying the color of the
        Duplo brick. These should be floats between 0 and 1. The default is red.

    Raises:
      ValueError: If `color` contains any value that is not between 0 and 1.
    """
    self._mjcf_root = mjcf.from_path(_DUPLO_XML_PATH)

    stud = self._mjcf_root.default.find('default', 'stud')
    if easy_align:
      # Make cylindrical studs invisible and disable contacts.
      stud.geom.group = 3
      stud.geom.contype = 9
      stud.geom.conaffinity = 8
      # Make capsule studs visible and enable contacts.
      stud_cap = self._mjcf_root.default.find('default', 'stud-capsule')
      stud_cap.geom.group = 0
      stud_cap.geom.contype = 0
      stud_cap.geom.conaffinity = 4
      self._active_stud_dclass = stud_cap
    else:
      self._active_stud_dclass = stud

    if flanges:
      flange_dclass = self._mjcf_root.default.find('default', 'flange')
      flange_dclass.geom.contype = 4  # Enable contact with flanges.

    stud_size = _STUD_SIZE_PARAMS[(easy_align, flanges)]
    offset = (1 - variation) * stud_size.lower_quartile
    self._lower = offset + variation * stud_size.minimum
    self._upper = offset + variation * stud_size.maximum

    self._studs = np.ndarray((2, 4), dtype=object)
    self._holes = np.ndarray((2, 4), dtype=object)

    for row in range(2):
      for column in range(4):
        self._studs[row, column] = self._mjcf_root.find(
            'site', 'stud_{}{}'.format(row, column))
        self._holes[row, column] = self._mjcf_root.find(
            'site', 'hole_{}{}'.format(row, column))

    if not all(0 <= value <= 1 for value in color):
      raise ValueError(_COLOR_NOT_BETWEEN_0_AND_1.format(color))
    self._mjcf_root.default.geom.rgba[:3] = color

  def initialize_episode_mjcf(self, random_state):
    """Randomizes the stud radius (and therefore the separation force)."""
    radius = random_state.uniform(self._lower, self._upper)
    self._active_stud_dclass.geom.size[0] = radius

  def _build_observables(self):
    return DuploObservables(self)

  @property
  def studs(self):
    """A (2, 4) numpy array of `mjcf.Elements` corresponding to stud sites."""
    return self._studs

  @property
  def holes(self):
    """A (2, 4) numpy array of `mjcf.Elements` corresponding to hole sites."""
    return self._holes

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


class DuploObservables(composer.Observables, composer.FreePropObservableMixin):
  """Observables for the `Duplo` prop."""

  @define.observable
  def position(self):
    return observable.MJCFFeature(
        'sensordata',
        self._entity.mjcf_model.find('sensor', 'position'))

  @define.observable
  def orientation(self):
    return observable.MJCFFeature(
        'sensordata',
        self._entity.mjcf_model.find('sensor', 'orientation'))

  @define.observable
  def linear_velocity(self):
    return observable.MJCFFeature(
        'sensordata',
        self._entity.mjcf_model.find('sensor', 'linear_velocity'))

  @define.observable
  def angular_velocity(self):
    return observable.MJCFFeature(
        'sensordata',
        self._entity.mjcf_model.find('sensor', 'angular_velocity'))

  @define.observable
  def force(self):
    return observable.MJCFFeature(
        'sensordata',
        self._entity.mjcf_model.find('sensor', 'force'))
