from Robotic_Arm.rm_robot_interface import *
import numpy as np

class rm75_kine_api():
    def __init__(self):
        #   ----------      rm75 official algorithm         -----------
        print(f'-------  the realman official kinematic initialising -------')
        arm_model = rm_robot_arm_model_e.RM_MODEL_RM_75_E  # RM_65 Robotic arm
        force_type = rm_force_type_e.RM_MODEL_RM_B_E  # Standard version
        # Initialize the robotic arm model and sensor type in the algorithm
        self.robot_kine_rm = Algo(arm_model, force_type)

        self.tool_frames = {
            'ee': rm_frame_t(frame_name="ee", pose=(0.0, 0.0, 0.0, 0.0, 0, 0.0), payload=1, x=0, y=0, z=0),
            'scissor': rm_frame_t(frame_name="scissor", pose=(0.0, 0.0, 0.144, 0.0, 0, 0.0), payload=1, x=0, y=0, z=72),
            'camera': rm_frame_t(frame_name="camera", pose=(0.05, 0.02, 0.10, -1.57, 0, -1.57), payload=1, x=0, y=0, z=72)
        }
        self.work_frames = {
            'work': rm_frame_t(frame_name="ee", pose=(0.0, 0.0, 0.0, 0.0, 0, 0.0), payload=1, x=0, y=0, z=0),
        }

        self.tool_name = "ee"
        self.work_name = "work"

    def cfg_limit(self):
        joint_max_limit = np.array([
            3.14159, 2.2689, 3.14159, 2.3562, 3.14159, 2.234, 3.14159
        ]) * 57
        self.robot_kine_rm.rm_algo_set_joint_max_limit(joint_max_limit.tolist())
        joint_min_limit = np.array([
            -3.14159, -2.2689, -3.14159, -2.3562, -3.14159, -2.234, -3.14159
        ]) * 57
        self.robot_kine_rm.rm_algo_set_joint_min_limit(joint_min_limit.tolist())

    def cfg_work_frame(self , frame_name):
        self.robot_kine_rm.rm_algo_set_workframe(self.work_frames[frame_name])

    def get_work_frame(self):
        return self.robot_kine_rm.rm_algo_get_curr_workframe()

    def cfg_tool_frame(self, frame_name ):
        self.robot_kine_rm.rm_algo_set_toolframe(self.tool_frames[frame_name])

    def get_tool_frame(self):
        return self.robot_kine_rm.rm_algo_get_curr_toolframe()

    def forward_kinematics(self, q, flag  = 1  , tool="ee", work="work"):
        '''
        :param q:   list of joint values, in degree
        flag:       0: return list [x,y,z,w,x,y,z]. 1: return list [x,y,z,rx,ry,rz]
        '''
        if tool != self.tool_name:
            self.tool_name = tool
            self.cfg_tool_frame(tool)
        if work != self.work_name:
            self.work_name = work
            self.cfg_work_frame(work)

        return self.robot_kine_rm.rm_algo_forward_kinematics(joint=q, flag=flag)

    def inverse_kinematics(self, target_position, target_rpy=None, initial_guess=None, tool="ee", work="work"):
        if tool != self.tool_name:
            self.tool_name = tool
            self.cfg_tool_frame(tool)
        if work != self.work_name:
            self.work_name = work
            self.cfg_work_frame(work)

        target = target_position + target_rpy

        if initial_guess is not None:
            q_ref = initial_guess
        else:
            q_ref = [0.0, 110.0, 20.0, 40.0, 30.0, 180.0, 20.0]
        ret, phi = self.robot_kine_rm.rm_algo_calculate_arm_angle_from_config_rm75(q_ref)
        params = rm_inverse_kinematics_params_t(q_ref,
                                                target, 1)
        ret, q_out = self.robot_kine_rm.rm_algo_inverse_kinematics_rm75_for_arm_angle(params, phi)
        return ret, q_out