# TODO(kevin): Make this work for all hands.
# TODO(kevin): In the future, we'd like to solve for finger orientation as well.

import copy
import dataclasses
from typing import Mapping, Optional, Sequence, Tuple

import mujoco
import numpy as np
from absl import logging
from dm_control import mjcf
from dm_robotics.geometry import geometry
from dm_robotics.geometry import mujoco_physics

from dexterity import controllers
from dexterity.models.hands import shadow_hand_e_constants as consts
from dexterity.utils import mujoco_utils

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

# Integration timestep used to convert from joint velocities to joint positions.
_INTEGRATION_TIMESTEP_SEC = 1.0

# Damping factor.
_REGULARIZATION_WEIGHT = 1e-5

# If the norm of the error divided by the magnitude of the joint position update is
# greater than this value, then the solve is ended prematurely. This helps us avoid
# getting stuck in local minima.
_PROGRESS_THRESHOLD = 20.0


@dataclasses.dataclass
class _Solution:
    """Return value of an IK solution."""

    qpos: np.ndarray
    linear_err: float


class IKSolver:
    """Inverse kinematics solver for a dexterous hand."""

    def __init__(
        self,
        model: mjcf.RootElement,
        fingers: Optional[Sequence[consts.Components]] = None,
        prefix: str = "",
    ) -> None:
        """Constructor.

        Args:
            model: The MJCF model root.
            prefix: The prefix assigned to the hand model in case it is attached to
                another entity.
        """
        self._physics = mjcf.Physics.from_mjcf_model(model)
        self._geometry_physics = mujoco_physics.wrap(self._physics)

        if fingers is None:
            self._fingers = consts.FINGERS
        else:
            self._fingers = tuple(fingers)

        # Wrist information.
        wrist_joint_names = [j.name for j in consts.JOINT_GROUP[consts.Components.WR]]
        self._wrist_controllable_joints = []
        for joint_name in wrist_joint_names:
            joint_elem = model.find(
                "joint", mujoco_utils.prefix_identifier(joint_name, prefix)
            )
            self._wrist_controllable_joints.append(joint_elem)
        assert len(self._wrist_controllable_joints) == len(wrist_joint_names)
        self._wrist_num_joints = len(wrist_joint_names)
        self._wirst_joint_bindings = self._physics.bind(self._wrist_controllable_joints)
        self._wrist_nullspace_joint_position_reference = np.zeros(
            self._wrist_num_joints
        )

        # Finger information.
        self._controllable_joints = {}
        self._elements = {}
        for finger in self._fingers:
            fingertip_name = consts.FINGER_FINGERTIP_MAPPING[finger]
            fingertip_site_name = mujoco_utils.prefix_identifier(
                f"{fingertip_name}_site", prefix
            )
            fingertip_site_elem = model.find("site", fingertip_site_name)
            assert fingertip_site_elem is not None
            self._elements[finger] = fingertip_site_elem

            joint_names = [j.name for j in consts.JOINT_GROUP[finger]]
            joint_names += wrist_joint_names
            controllable_joints = []
            for joint_name in joint_names:
                joint_elem = model.find(
                    "joint", mujoco_utils.prefix_identifier(joint_name, prefix)
                )
                controllable_joints.append(joint_elem)
            assert len(controllable_joints) == len(joint_names)
            self._controllable_joints[finger] = controllable_joints

        # Get all the joints of the hand.
        self._all_joints = []
        for joint_name in consts.JOINT_NAMES:
            joint_elem = model.find(
                "joint", mujoco_utils.prefix_identifier(joint_name, prefix)
            )
            assert joint_elem is not None
            self._all_joints.append(joint_elem)

        self._joint_bindings = {}
        self._num_joints = {}
        for finger, controllable_joints in self._controllable_joints.items():
            self._joint_bindings[finger] = self._physics.bind(controllable_joints)
            self._num_joints[finger] = len(controllable_joints)
        self._all_joints_binding = self._physics.bind(self._all_joints)

        self._nullspace_joint_position_reference = np.mean(
            self._all_joints_binding.range, axis=1
        )

        self._create_mapper()

    @property
    def fingers(self) -> Tuple[consts.Components, ...]:
        return self._fingers

    def _create_mapper(self) -> None:
        obj_types = []
        obj_names = []
        for finger in self._fingers:
            obj_types.append(mujoco_utils.get_element_type(self._elements[finger]))
            obj_names.append(self._elements[finger].full_identifier)
        params = controllers.dls.DampedLeastSquaresParameters(
            model=self._physics.model,
            object_types=obj_types,
            object_names=obj_names,
            regularization_weight=_REGULARIZATION_WEIGHT,
        )
        self._mapper = controllers.dls.DampedLeastSquaresMapper(params)

    def solve(
        self,
        target_positions: Mapping[consts.Components, np.ndarray],
        linear_tol: float = 1e-3,
        max_steps: int = 100,
        early_stop: bool = False,
        num_attempts: int = 30,
        stop_on_first_successful_attempt: bool = False,
    ) -> Optional[np.ndarray]:
        nullspace_jnt_qpos_min_err: float = np.inf
        success: bool = False
        sol_qpos: Optional[np.ndarray] = None

        # Each iteration of this loop attempts to solve the IK problem.
        for attempt in range(num_attempts):
            # Randomize the initial joint configuration so that the IK can find
            # different solutions.
            if attempt == 0:
                self._all_joints_binding.qpos[
                    :
                ] = self._nullspace_joint_position_reference
            else:
                self._all_joints_binding.qpos[:] = np.random.uniform(
                    self._all_joints_binding.range[:, 0],
                    self._all_joints_binding.range[:, 1],
                )

            solution = self._solve_ik(
                target_positions,
                linear_tol,
                max_steps,
                early_stop,
            )

            if solution.linear_err <= linear_tol:
                success = True

                nullspace_jnt_qpos_err = float(
                    np.linalg.norm(
                        solution.qpos - self._nullspace_joint_position_reference
                    )
                )
                if nullspace_jnt_qpos_err < nullspace_jnt_qpos_min_err:
                    nullspace_jnt_qpos_min_err = nullspace_jnt_qpos_err
                    sol_qpos = solution.qpos

            if success and stop_on_first_successful_attempt:
                break

        if not success:
            logging.warning(f"{self.__class__.__name__} failed to find a solution.")

        return sol_qpos

    def _solve_ik(
        self,
        target_positions: Mapping[consts.Components, np.ndarray],
        linear_tol: float,
        max_steps: int,
        early_stop: bool,
    ) -> _Solution:
        """Solves for a joint configuration that brings element pose to target pose."""
        cur_frames = {}
        cur_poses = {}
        previous_poses = {}
        for finger, target_position in target_positions.items():
            cur_frame = geometry.PoseStamped(pose=None, frame=self._elements[finger])
            cur_pose = cur_frame.get_world_pose(self._geometry_physics)
            cur_frames[finger] = cur_frame
            cur_poses[finger] = cur_pose
            previous_poses[finger] = copy.copy(cur_pose)

        # Each iteration of this loop attempts to reduce the error between the site's
        # position and the target position.
        for _ in range(max_steps):
            twists = []
            for finger, target_position in target_positions.items():
                twist = _compute_twist(
                    cur_poses[finger],
                    geometry.Pose(position=target_position, quaternion=None),
                    _LINEAR_VELOCITY_GAIN,
                    _INTEGRATION_TIMESTEP_SEC,
                )
                twists.append(twist)

            qdot_sol = self._compute_joint_velocities(twists)

            mujoco.mj_integratePos(
                self._physics.model.ptr,
                self._physics.data.qpos,
                qdot_sol,
                _INTEGRATION_TIMESTEP_SEC,
            )
            self._update_physics_data()

            avg_linear_err: float = 0.0
            close_enough: bool = True
            not_enough_progress: bool = False

            for finger, target_position in target_positions.items():
                # Get the distance between the current pose and the target pose.
                cur_pose = cur_frames[finger].get_world_pose(self._geometry_physics)
                linear_err = float(np.linalg.norm(target_position - cur_pose.position))
                avg_linear_err += linear_err

                # Stop if the pose is close enough to the target pose.
                if linear_err > linear_tol:
                    close_enough = False

                # Stop the solve if not enough progress is being made.
                previous_pose = previous_poses[finger]
                linear_change = np.linalg.norm(
                    cur_pose.position - previous_pose.position
                )
                if linear_err / (linear_change + 1e-10) > _PROGRESS_THRESHOLD:
                    not_enough_progress = True

                previous_poses[finger] = copy.copy(cur_pose)
                cur_poses[finger] = cur_pose

            # Average out the linear error.
            avg_linear_err /= len(target_positions)

            # Break conditions.
            if early_stop and close_enough:
                break
            if not_enough_progress:
                break

        qpos = np.array(self._all_joints_binding.qpos, copy=True)
        return _Solution(qpos=qpos, linear_err=avg_linear_err)

    def _compute_joint_velocities(
        self, cartesian_6d_target: Sequence[np.ndarray]
    ) -> np.ndarray:
        """Maps a Cartesian 6D target velocity to joint velocities."""
        return self._mapper.compute_joint_velocities(
            data=self._physics.data,
            target_velocities=cartesian_6d_target,
            nullspace_bias=None,
        )

    def _update_physics_data(self) -> None:
        """Updates the physics data following the integration of velocities."""
        # Clip joint positions.
        qpos = self._all_joints_binding.qpos
        min_range = self._all_joints_binding.range[:, 0]
        max_range = self._all_joints_binding.range[:, 1]
        qpos = np.clip(qpos, min_range, max_range)
        self._all_joints_binding.qpos[:] = qpos

        # Forward kinematics to update the pose of the tracked site.
        mujoco.mj_normalizeQuat(self._physics.model.ptr, self._physics.data.qpos)
        mujoco.mj_kinematics(self._physics.model.ptr, self._physics.data.ptr)
        mujoco.mj_comPos(self._physics.model.ptr, self._physics.data.ptr)


def _compute_twist(
    init_pose: geometry.Pose,
    final_pose: geometry.Pose,
    linear_velocity_gain: float,
    control_timestep_seconds: float,
) -> np.ndarray:
    """Returns the twist to apply to the element to reach final_pose from init_pose."""
    position_error = final_pose.position - init_pose.position
    linear = linear_velocity_gain * position_error / control_timestep_seconds
    return linear
