# Copyright 2020 DeepMind Technologies Limited.
#
# 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.

"""IK solver for initialization of robot arms."""

import copy
from typing import List, NamedTuple, Optional, Sequence, Union

from absl import logging
from dm_control import mjcf
from dm_control.mujoco.wrapper import mjbindings
from dm_control.mujoco.wrapper.mjbindings.enums import mjtObj
from dm_robotics.controllers import cartesian_6d_to_joint_velocity_mapper
from dm_robotics.geometry import geometry
from dm_robotics.geometry import mujoco_physics
from dm_robotics.transformations import transformations as tr
import numpy as np

# Default value for the nullspace gain parameter.
_NULLSPACE_GAIN = 0.4

# Gain for the linear and angular twist computation, these values should always
# be between 0 and 1. 0 corresponds to not move and 1 corresponds to move to the
# target in a single integration timestep.
_LINEAR_VELOCITY_GAIN = 0.95
_ANGULAR_VELOCITY_GAIN = 0.95

# Integration timestep used when solving the IK.
_INTEGRATION_TIMESTEP_SEC = 1.0


# At each step of the solve, we measure how much the tracked element
# translated (linear progress) and rotated (angular progress). We compare this
# progress to the total linear and angular error and if not enough progress is
# made stop the solve before the maximum number of steps is reached.
_ERR_TO_PROGRESS_THRESHOLD = 20.0

### ---------------PARAMETERS USED FOR THE QP MAPPER: START----------------- ###
# Regularisation parameter used by the qp to compute joint velocities.
_REGULARIZATION_WEIGHT = 0.01

# Ensure that the joint limits are respected.
_ENABLE_JOINT_POSITION_LIMITS = True

# Gain that scales the joint velocities down when close to the joint limits.
_JOINT_POSITION_LIMIT_VELOCITY_SCALE = 1.0

# The minimal distance to joint limits the IK solution can have.
_MINIMUM_DISTANCE_FROM_JOINT_POSITION_LIMIT = 0.0

# Maximum number of iteration to find a joint velocity command that applies the
# desired twist to the element.
_MAX_CARTESIAN_VELOCITY_CONTROL_ITERATIONS = 300

# Number of iterations for the nullspace control problem.
_MAX_NULLSPACE_CONTROL_ITERATIONS = 300

# Maximum error allowed for the nullspace problem.
_NULLSPACE_PROJECTION_SLACK = 1e-5

# Maximum error allowed between the requested twist command and the actual one.
_SOLUTION_TOLERANCE = 1e-4

# Remove the logging when the nullspace cannot find a solution as this
# clutters the logging.
_LOG_NULLSPACE_FAILURE_WARNINGS = False
### -----------------PARAMETERS USED FOR THE QP MAPPER: END----------------- ###

_Binding = Union[mjcf.physics.Binding, mjcf.physics._EmptyBinding]  # pylint: disable=protected-access
_MjcfElement = mjcf.element._ElementImpl  # pylint: disable=protected-access


class _Solution(NamedTuple):
  """Return value of an ik solution.

  Attributes:
    qpos: The joint configuration.
    linear_err: The linear error between the target pose and desired pose.
    angular_err: The angular error between the target pose and desired pose.
  """
  qpos: np.ndarray
  linear_err: float
  angular_err: float


class IkSolver():
  """Inverse kinematics solver.

  This class computes a joint configuration that brings an element to a certain
  pose.
  """

  # The cartesian velocity controller used to solve the IK.
  _qp_mapper: cartesian_6d_to_joint_velocity_mapper.Mapper

  # Array of indices that sorts the joints in ascending order. The qp_mapper
  # returns values in joint-ID ascending order which could be different than
  # the order of the joints provided by the user.
  _joints_argsort: List[int]

  # The desired joint configuration that is set as the nullspace goal. This
  # corresponds to the mid-range of each joint. The user can override this
  # reference configuration in the `solve` method.
  _nullspace_joint_position_reference: List[float]

  def __init__(
      self,
      model: mjcf.RootElement,
      controllable_joints: List[_MjcfElement],
      element: _MjcfElement,
      nullspace_gain: float = _NULLSPACE_GAIN,
      ):
    """Constructor.

    Args:
      model: The MJCF model root.
      controllable_joints: The joints that can be controlled to achieve
        the desired target pose. Only 1 DoF joints are supported.
      element: The MJCF element that is being placed by the inverse kinematics
        solver. Only body, geoms, and sites are supported
      nullspace_gain: Scales the nullspace velocity bias. If the gain is set to
        0, there will be no nullspace optimization during the solve process.
    """
    self._physics = mjcf.Physics.from_mjcf_model(model)
    self._geometry_physics = mujoco_physics.wrap(self._physics)
    self._joints_binding = _binding(self._physics, controllable_joints)
    self._num_joints = len(controllable_joints)
    self._element = element
    self._nullspace_gain = nullspace_gain
    self._create_qp_mapper()

  def solve(self,
            ref_pose: geometry.Pose,
            linear_tol: float = 1e-3,
            angular_tol: float = 1e-3,
            max_steps: int = 100,