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

"""Walkers based on an actuated jumping ball."""

import io
import os

from dm_control import composer
from dm_control import mjcf
from dm_control.composer.observation import observable
from dm_control.locomotion.walkers import legacy_base
import numpy as np
from PIL import Image

from dm_control.utils import io as resources

_ASSETS_PATH = os.path.join(os.path.dirname(__file__), 'assets', 'boxhead')
_MAX_WALKER_ID = 10
_INVALID_WALKER_ID = 'walker_id must be in [0-{}], got: {{}}.'.format(
    _MAX_WALKER_ID)


def _compensate_gravity(physics, body_elements):
  """Applies Cartesian forces to bodies in order to exactly counteract gravity.

  Note that this will also affect the output of pressure, force, or torque
  sensors within the kinematic chain leading from the worldbody to the bodies
  that are being gravity-compensated.

  Args:
    physics: An `mjcf.Physics` instance to modify.
    body_elements: An iterable of `mjcf.Element`s specifying the bodies to which
      gravity compensation will be applied.
  """
  gravity = np.hstack([physics.model.opt.gravity, [0, 0, 0]])
  bodies = physics.bind(body_elements)
  bodies.xfrc_applied = -gravity * bodies.mass[..., None]


def _alpha_blend(foreground, background):
  """Does alpha compositing of two RGBA images.

  Both inputs must be (..., 4) numpy arrays whose shapes are compatible for
  broadcasting. They are assumed to contain float RGBA values in [0, 1].

  Args:
    foreground: foreground RGBA image.
    background: background RGBA image.

  Returns:
    A numpy array of shape (..., 4) containing the blended image.
  """
  fg, bg = np.broadcast_arrays(foreground, background)
  fg_rgb = fg[..., :3]
  fg_a = fg[..., 3:]
  bg_rgb = bg[..., :3]
  bg_a = bg[..., 3:]
  out = np.empty_like(bg)
  out_a = out[..., 3:]
  out_rgb = out[..., :3]
  # https://en.wikipedia.org/wiki/Alpha_compositing#Alpha_blending
  out_a[:] = fg_a + bg_a * (1. - fg_a)
  out_rgb[:] = fg_rgb * fg_a + bg_rgb * bg_a * (1. - fg_a)
  # Avoid division by zero if foreground and background are both transparent.
  out_rgb[:] = np.where(out_a, out_rgb / out_a, out_rgb)
  return out


def _asset_png_with_background_rgba_bytes(asset_fname, background_rgba):
  """Decode PNG from asset file and add solid background."""

  # Retrieve PNG image contents as a bytestring, convert to a numpy array.
  contents = resources.GetResource(os.path.join(_ASSETS_PATH, asset_fname))
  digit_rgba = np.array(Image.open(io.BytesIO(contents)), dtype=np.double)

  # Add solid background with `background_rgba`.
  blended = 255. * _alpha_blend(digit_rgba / 255., np.asarray(background_rgba))

  # Encode composite image array to a PNG bytestring.
  img = Image.fromarray(blended.astype(np.uint8), mode='RGBA')
  buf = io.BytesIO()
  img.save(buf, format='PNG')
  png_encoding = buf.getvalue()
  buf.close()

  return png_encoding


class BoxHeadObservables(legacy_base.WalkerObservables):
  """BoxHead observables with low-res camera and modulo'd rotational joints."""

  def __init__(self, entity, camera_resolution):
    self._camera_resolution = camera_resolution
    super().__init__(entity)

  @composer.observable
  def egocentric_camera(self):
    width, height = self._camera_resolution
    return observable.MJCFCamera(self._entity.egocentric_camera,
                                 width=width, height=height)

  @property
  def proprioception(self):
    proprioception = super().proprioception
    if self._entity.observable_camera_joints:
      return proprioception + [self.camera_joints_pos, self.camera_joints_vel]
    return proprioception

  @composer.observable
  def camera_joints_pos(self):

    def _sin(value, random_state):
      del random_state
      return np.sin(value)

    def _cos(value, random_state):
      del random_state
      return np.cos(value)

    sin_rotation_joints = observable.MJCFFeature(
        'qpos', self._entity.observable_camera_joints, corruptor=_sin)

    cos_rotation_joints = observable.MJCFFeature(
        'qpos', self._entity.observable_camera_joints, corruptor=_cos)

    def _camera_joints(physics):
      return np.concatenate([
          sin_rotation_joints(physics),
          cos_rotation_joints(physics)
      ], -1)

    return observable.Generic(_camera_joints)

  @composer.observable
  def camera_joints_vel(self):
    return observable.MJCFFeature(
        'qvel', self._entity.observable_camera_joints)


class BoxHead(legacy_base.Walker):
  """A rollable and jumpable ball with a head."""

  def _build(self,
             name='walker',
             marker_rgba=None,
             camera_control=False,
             camera_resolution=(28, 28),
             roll_gear=-60,
             steer_gear=55,
             walker_id=None,
             initializer=None):
    """Build a BoxHead.

    Args:
      name: name of the walker.
      marker_rgba: RGBA value set to walker.marker_geoms to distinguish between
        walkers (in multi-agent setting).
      camera_control: If `True`, the walker exposes two additional actuated
        degrees of freedom to control the egocentric camera height and tilt.
      camera_resolution: egocentric camera rendering resolution.
      roll_gear: gear determining forward acceleration.
      steer_gear: gear determining steering (spinning) torque.
      walker_id: (Optional) An integer in [0-10], this number will be shown on
        the walker's head. Defaults to `None` which does not show any number.
      initializer: (Optional) A `WalkerInitializer` object.

    Raises:
      ValueError: if received invalid walker_id.
    """
    super()._build(initializer=initializer)
    xml_path = os.path.join(_ASSETS_PATH, 'boxhead.xml')
    self._mjcf_root = mjcf.from_xml_string(resources.GetResource(xml_path, 'r'))
    if name:
      self._mjcf_root.model = name

    if walker_id is not None and not 0 <= walker_id <= _MAX_WALKER_ID:
      raise ValueError(_INVALID_WALKER_ID.format(walker_id))

    self._walker_id = walker_id
    if walker_id is not None:
      png_bytes = _asset_png_with_background_rgba_bytes(
          'digits/%02d.png' % walker_id, marker_rgba)
      head_texture = self._mjcf_root.asset.add(
          'texture',
          name='head_texture',
          type='2d',
          file=mjcf.Asset(png_bytes, '.png'))
      head_material = self._mjcf_root.asset.add(
          'material', name='head_material', texture=head_texture)
      self._mjcf_root.find('geom', 'head').material = head_material
      self._mjcf_root.find('geom', 'head').rgba = None

      self._mjcf_root.find('geom', 'top_down_cam_box').material = head_material
      self._mjcf_root.find('geom', 'top_down_cam_box').rgba = None

    self._body_texture = self._mjcf_root.asset.add(
        'texture',
        name='ball_body',
        type='cube',
        builtin='checker',
        rgb1=marker_rgba[:-1] if marker_rgba else '.4 .4 .4',
        rgb2='.8 .8 .8',
        width='100',
        height='100')
    self._body_material = self._mjcf_root.asset.add(
        'material', name='ball_body', texture=self._body_texture)
    self._mjcf_root.find('geom', 'shell').material = self._body_material

    # Set corresponding marker color if specified.
    if marker_rgba is not None:
      for geom in self.marker_geoms:
        geom.set_attributes(rgba=marker_rgba)

    self._root_joints = None
    self._camera_control = camera_control
    self._camera_resolution = camera_resolution
    if not camera_control:
      for name in ('camera_pitch', 'camera_yaw'):
        self._mjcf_root.find('actuator', name).remove()
        self._mjcf_root.find('joint', name).remove()
    self._roll_gear = roll_gear
    self._steer_gear = steer_gear
    self._mjcf_root.find('actuator', 'roll').gear[0] = self._roll_gear
    self._mjcf_root.find('actuator', 'steer').gear[0] = self._steer_gear

    # Initialize previous action.
    self._prev_action = np.zeros(shape=self.action_spec.shape,
                                 dtype=self.action_spec.dtype)

  def _build_observables(self):
    return BoxHeadObservables(self, camera_resolution=self._camera_resolution)

  @property
  def marker_geoms(self):
    geoms = [
        self._mjcf_root.find('geom', 'arm_l'),
        self._mjcf_root.find('geom', 'arm_r'),
        self._mjcf_root.find('geom', 'eye_l'),
        self._mjcf_root.find('geom', 'eye_r'),
    ]
    if self._walker_id is None:
      geoms.append(self._mjcf_root.find('geom', 'head'))
    return geoms

  def create_root_joints(self, attachment_frame):
    root_class = self._mjcf_root.find('default', 'root')
    root_x = attachment_frame.add(
        'joint', name='root_x', type='slide', axis=[1, 0, 0], dclass=root_class)
    root_y = attachment_frame.add(
        'joint', name='root_y', type='slide', axis=[0, 1, 0], dclass=root_class)
    root_z = attachment_frame.add(
        'joint', name='root_z', type='slide', axis=[0, 0, 1], dclass=root_class)
    self._root_joints = [root_x, root_y, root_z]

  def set_pose(self, physics, position=None, quaternion=None):
    if position is not None:
      if self._root_joints is not None:
        physics.bind(self._root_joints).qpos = position
      else:
        super().set_pose(physics, position, quaternion=None)
    physics.bind(self._mjcf_root.find_all('joint')).qpos = 0.
    if quaternion is not None:
      # This walker can only rotate along the z-axis, so we extract only that
      # component from the quaternion.
      z_angle = np.arctan2(
          2 * (quaternion[0] * quaternion[3] + quaternion[1] * quaternion[2]),
          1 - 2 * (quaternion[2] ** 2 + quaternion[3] ** 2))
      physics.bind(self._mjcf_root.find('joint', 'steer')).qpos = z_angle

  def set_velocity(self, physics, velocity=None, angular_velocity=None):
    if velocity is not None:
      if self._root_joints is not None:
        physics.bind(self._root_joints).qvel = velocity

    if angular_velocity is not None:
      # This walker can only rotate along the z-axis, so we extract only that
      # component from the angular_velocity.
      steer_joint = self._mjcf_root.find('joint', 'steer')
      if isinstance(angular_velocity, float):
        z_velocity = angular_velocity
      else:
        z_velocity = angular_velocity[2]
      physics.bind(steer_joint).qvel = z_velocity

  def initialize_episode(self, physics, random_state):
    if self._camera_control:
      _compensate_gravity(physics,
                          self._mjcf_root.find('body', 'egocentric_camera'))
    self._prev_action = np.zeros(shape=self.action_spec.shape,
                                 dtype=self.action_spec.dtype)

  def apply_action(self, physics, action, random_state):
    super().apply_action(physics, action, random_state)

    # Updates previous action.
    self._prev_action[:] = action

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

  @composer.cached_property
  def actuators(self):
    return self._mjcf_root.find_all('actuator')

  @composer.cached_property
  def root_body(self):
    return self._mjcf_root.find('body', 'head_body')

  @composer.cached_property
  def end_effectors(self):
    return (self._mjcf_root.find('body', 'head_body'),)

  @composer.cached_property
  def observable_joints(self):
    return (self._mjcf_root.find('joint', 'kick'),)

  @composer.cached_property
  def observable_camera_joints(self):
    if self._camera_control:
      return (
          self._mjcf_root.find('joint', 'camera_yaw'),
          self._mjcf_root.find('joint', 'camera_pitch'),
      )
    return ()

  @composer.cached_property
  def egocentric_camera(self):
    return self._mjcf_root.find('camera', 'egocentric')

  @composer.cached_property
  def ground_contact_geoms(self):
    return (self._mjcf_root.find('geom', 'shell'),)

  @property
  def prev_action(self):
    return self._prev_action
