#!/usr/bin/env python3
"""Compatibility adapter for the original experimental import path.

New code should import RM75Kinematics and RM75IkSolver from ``rm75_ik``.
"""

from math import pi
from pathlib import Path

import numpy as np
import pinocchio as pin

from rm75_ik import IkOptions, JointLimits, RM75IkSolver, RM75Kinematics


class KinematicsSolver:
    def __init__(self, urdf_path="urdf_rm75/RM75-B.urdf", mesh_dir=None):
        del mesh_dir
        selected_path = Path(urdf_path)
        if not selected_path.is_file() and not selected_path.is_absolute():
            selected_path = Path(__file__).resolve().parent / selected_path
        self._urdf_path = selected_path
        self._limits = None
        self._tools = {}
        self._rebuild()

    def _rebuild(self):
        self.kinematics = RM75Kinematics(self._urdf_path, self._limits)
        self.solver = RM75IkSolver(self.kinematics)
        self.model = self.kinematics.model
        self.data = self.kinematics.data

    def add_tool_frames(self, frames):
        for name, attributes in frames.items():
            pose = np.asarray(attributes[0], dtype=float)
            if pose.shape != (7,):
                raise ValueError(f"tool {name!r} pose must have seven values")
            quaternion = pin.Quaternion(pose[6], pose[3], pose[4], pose[5])
            quaternion.normalize()
            self._tools[name] = pin.SE3(quaternion.matrix(), pose[:3])

    def cfg_j_limit(self, min_j=None, max_j=None, rad_flag=True):
        if min_j is None:
            min_j = [-pi, -2.2689, -pi, -2.3562, -pi, -2.234, -2 * pi]
        if max_j is None:
            max_j = [pi, 2.2689, pi, 2.3562, pi, 2.234, 2 * pi]
        lower = np.asarray(min_j, dtype=float)
        upper = np.asarray(max_j, dtype=float)
        if not rad_flag:
            lower = np.deg2rad(lower)
            upper = np.deg2rad(upper)
        self._limits = JointLimits("legacy", lower, upper)
        self._rebuild()

    def _tool(self, name):
        try:
            return self._tools[name]
        except KeyError as exc:
            raise ValueError(f"unknown tool frame: {name!r}") from exc

    def forward_kinematics(self, joint_angles, tool="no_tool"):
        pose = self.kinematics.forward(np.asarray(joint_angles), self._tool(tool))
        return np.concatenate(
            [pose.translation.copy(), pin.rpy.matrixToRpy(pose.rotation)]
        )

    def inverse_kinematics(
        self,
        target_position,
        target_rpy=None,
        target_quat=None,
        initial_guess=None,
        max_iter=500,
        tolerance=1e-3,
        debug=False,
        tool="no_tool",
    ):
        del debug
        if target_quat is not None:
            values = np.asarray(target_quat, dtype=float)
            quaternion = pin.Quaternion(values[3], values[0], values[1], values[2])
            rotation = quaternion.matrix()
        elif target_rpy is not None:
            rotation = pin.rpy.rpyToMatrix(*target_rpy)
        else:
            rotation = np.eye(3)
        tool_pose = self._tool(tool)
        target_tool = pin.SE3(rotation, np.asarray(target_position, dtype=float))
        target_flange = target_tool * tool_pose.inverse()
        seed = np.zeros(7) if initial_guess is None else np.asarray(initial_guess, dtype=float)
        result = self.solver.solve(
            target_flange,
            seed,
            IkOptions(
                position_tolerance_m=tolerance,
                orientation_tolerance_rad=tolerance,
                max_iterations=max_iter,
            ),
        )
        return (0, result.q.tolist()) if result.success else (-1, [])

    def compute_jacobian(self, joint_angles, tool="no_tool"):
        del tool
        return self.kinematics.jacobian(np.asarray(joint_angles, dtype=float))