from Robotic_Arm.rm_robot_interface import *
import numpy as np
import math

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_75 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.cfg_j_limit()

        self.work_frames = {
            'work': rm_frame_t(frame_name="work", pose=(0.0, 0.0, 0.0, 0.0, 0, 0.0), payload=1, x=0, y=0, z=0),
        }

        self.tool_name = "no_tool"
        self.work_name = "work"

    def cfg_j_limit(self, min_j=None, max_j=None, rad_flag = True):
        if max_j is None:
            max_j = np.array([3.14159, 2.2689, 3.14159, 2.3562, 3.14159, 2.234, 3.14159])
        if min_j is None:
            min_j = np.array([ -3.14159, -2.2689, -3.14159, -2.3562, -3.14159, -2.234, -3.14159  ])

        max_j = np.array(max_j)
        min_j = np.array(min_j)
        if rad_flag:
            self.robot_kine_rm.rm_algo_set_joint_max_limit((max_j * 180 / math.pi).tolist())
            self.robot_kine_rm.rm_algo_set_joint_min_limit((min_j * 180 / math.pi).tolist())
        else:
            self.robot_kine_rm.rm_algo_set_joint_max_limit(max_j.tolist())
            self.robot_kine_rm.rm_algo_set_joint_min_limit(min_j.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 quaternion_to_euler(self, q):
        """
        Convert quaternion to Euler angles (roll, pitch, yaw)

        Args:
            qx, qy, qz, qw: quaternion components

        Returns:
            tuple: (roll, pitch, yaw) in radians
        """
        # Roll (x-axis rotation)
        sinr_cosp = 2.0 * (q[3] * q[0] + q[1] * q[2])
        cosr_cosp = 1.0 - 2.0 * (q[0] * q[0] + q[1] * q[1])
        roll = np.arctan2(sinr_cosp, cosr_cosp)

        # Pitch (y-axis rotation)
        sinp = 2.0 * (q[3] * q[1] - q[2] * q[0])
        if abs(sinp) >= 1:
            pitch = np.copysign(np.pi / 2, sinp)  # Use 90 degrees if out of range
        else:
            pitch = np.arcsin(sinp)

        # Yaw (z-axis rotation)
        siny_cosp = 2.0 * (q[3] * q[2] + q[0] * q[1])
        cosy_cosp = 1.0 - 2.0 * (q[1] * q[1] + q[2] * q[2])
        yaw = np.arctan2(siny_cosp, cosy_cosp)

        return [roll, pitch, yaw]

    def add_tool_frames(self, dict_frames):
        self.tool_frames = {}
        for tool_name in dict_frames:
            tool_attr = dict_frames[tool_name]
            position = tool_attr[0][0:3]
            rotationXYZ = self.quaternion_to_euler(tool_attr[0][3:7])
            f = rm_frame_t(frame_name=tool_name, pose=(position[0], position[1], position[2], rotationXYZ[0], rotationXYZ[1], rotationXYZ[2]), payload=1, x=0, y=0, z=0)

            self.tool_frames.update({tool_name:f})

    def forward_kinematics(self, joint_angles, flag  = 1  , tool="omnipic", work="work"):
        '''
        :param joint_angles:   list of joint values, in rad
        :param flag:       0: return list [x,y,z,w,x,y,z]. 1: return list [x,y,z,rx,ry,rz]
        :param return:      [x,y,z,rx,ry,rz], m & rad
        '''
        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_s*180/math.pi for q_s in joint_angles]  , flag=flag)

    def inverse_kinematics(self, target_position, target_rpy=None, initial_guess=None, tool="omnipic", work="work"):
        '''
        :param target_position:     list of position values, m
        :param target_rpy:          list of rpy values, rad
        :param initial_guess:       initial guess of angles, rad
        :param tool:                tool name, refer to self.tool_frames
        :param work:                work name, refer to self.work_frames

        return  ret:                 state of ik calculation, 0:success, -2: out of workspace
                [q_]:               the ik calculated angles for joints, rad
        '''
        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 = list(target_position) + list(target_rpy)

        if initial_guess is not None:
            q_ref = [ 180/math.pi * ig for ig in  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)
        # print(f'the arm angle is ret = {ret}, and phi = {phi}')
        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)
        pose_fk = self.robot_kine_rm.rm_algo_forward_kinematics(joint=q_out, flag=1)
        pose_dis = cal_pose_deviation(pose_fk, target)

        # print(f'target pose is {target}, fk pose is {pose_fk}, dis of poses is {pose_dis}')
        #
        # print(f'\nin the rm75_kine_rm, l133, inverse_kinematics, q_ref = {q_ref}, target = {target}  phi = {phi}, q_out = {q_out}, ret = {ret}\n\n')
        if ret < 0:
            return ret, [ q/180*math.pi for q in q_out]
        elif pose_dis < 0.01:
            return ret, [ q/180*math.pi for q in q_out]
        else:
            return -10, [ q/180*math.pi for q in q_out]

def cal_pose_deviation(pose1, pose2):
    d_fk_p1 = np.array(pose1) - np.array(pose2)
    for j in [3, 4, 5]:
        while d_fk_p1[j] > math.pi:
            d_fk_p1[j] -= 2 * math.pi
        while d_fk_p1[j] < -math.pi:
            d_fk_p1[j] += 2 * math.pi
    d_fk = np.linalg.norm(d_fk_p1)
    return d_fk