'''
Robotic arm kinematics solver for realman-75

Files:
    rm75_kinematics.py
    rm75_kine/rm75_kine_qp.py
    rm75_kine/rm75_kine_rm.py

How to use:
    Example in main.py:
        python main.py
'''

from kine_ctrl.rm75_kine.rm75_kine_qp import KinematicsSolver as kine_qp
from kine_ctrl.rm75_kine.rm75_kine_rm import rm75_kine_api as kine_rm


class rm75_kinematics():
    def __init__(self,urdf_path="", mesh_dir="", tcps=None,tools_in_ee=None,min_j=0.0, max_j=0.0):
        self.robot_kine_qp = kine_qp(urdf_path=urdf_path, mesh_dir=mesh_dir)
        self.robot_kine_qp.add_tool_frames(tools_in_ee)
        self.robot_kine_qp.cfg_j_limit(min_j=min_j, max_j=max_j, rad_flag=True)

        #   ----------      rm75 official algorithm         -----------
        self.robot_kine_rm = kine_rm()
        self.robot_kine_rm.add_tool_frames(tools_in_ee)
        self.robot_kine_rm.cfg_j_limit(min_j=min_j, max_j=max_j, rad_flag=True)

        self.ik_sts = False
        self.joint_solved = [0.0] * 7

    def get_ik_result(self, target_position, target_rpy, initial_guess, tool):
        '''
        Try both the RM official solver and the QP solver; on success, store the result.
        return: True if either solver succeeded, joint_solved in rad.
        '''
        ret_rm, q_out = self.robot_kine_rm.inverse_kinematics(target_position, target_rpy, initial_guess, tool)
        if ret_rm != 0:
            ret_rm, q_out = self.robot_kine_qp.inverse_kinematics(target_position=target_position, target_rpy=target_rpy, initial_guess=initial_guess, tool=tool, max_iter=300)

        if ret_rm == 0:
            self.joint_solved = q_out
            self.ik_sts = True
        else:
            self.ik_sts = False

        return self.ik_sts, self.joint_solved

    def get_fk_result(self, joint_angles, tool):
        '''
        Get the forward kinematics result for given joint angles and tool.
        :param joint_angles:   list of joint values, in rad
        :param return:      [x,y,z,rx,ry,rz], m & rad
        return: [x, y, z, rx, ry, rz] in meters and radians.
        '''
        fk_result = self.robot_kine_rm.forward_kinematics(joint_angles, flag=1, tool=tool)
        return fk_result

    def get_self_collision(self, joint_angles):
        '''
        Check for self-collision given joint angles.
        :param joint_angles:   list of joint values, in rad
        :return: True if self-collision is detected, False otherwise.
        '''
        collision_detected = self.robot_kine_qp.collision_detect(joint_angles)
        return collision_detected