#   conda activate coppeliasim
#   env fix, in terminal:  fix_robotics_env.sh

from rm75_kine_qp import KinematicsSolver as kine_qp
from rm75_kine_rm import rm75_kine_api as kine_rm
from rm75_mjc import MuJoCoPositionController
from Robotic_Arm.rm_robot_interface import *

import time
from math import radians, degrees, pi, cos, sin
import numpy as np

def main():
    """Demonstrate pure position control"""

    # Create controller
    robot_mjk = MuJoCoPositionController()

    tool_name = "scissor"
    #  -----------      rm75 qp based kine          ------------
    robot_kine_qp = kine_qp()

    #   ----------      rm75 official algorithm         -----------
    robot_kine_rm = kine_rm()



#   --------------  for comparison   ----------------
    print(f'in the comparison part')

    if True:

        # ub = np.array([150.0, 110.0, 170.0,  130, 175.0, 125.0,  179.0])
        # lb = np.array([-150.0, -30.0, -170.0, -130, -175.0, -125.0,  -179.0])

        ub = np.array([179.0, 129.0, 179.0, 134, 179.0, 127.0, 359.0])/180*pi
        lb = -ub

        robot_kine_qp.cfg_j_limit(min_j=lb, max_j=ub, rad_flag=True)

        robot_kine_rm.cfg_j_limit(min_j=lb, max_j=ub, rad_flag=True)

        result = np.array([[0,0],[0,0]], dtype=np.int32)        # to collect ik result   qp_fk, qp_ik, rm_fk, rm_ik

        solve_sum = 0

        for i in range(10):
            print(f'\n-------------- in i = {i} ----------------')
            joint_rand = np.random.uniform(ub, lb)
            print(f'the predefined joints are {joint_rand}')

            #   -------------- fk ------------------
            fk_qp_p1 = robot_kine_qp.forward_kinematics(joint_angles=joint_rand.tolist(), tool=tool_name)

            fk_rm_p1 = robot_kine_rm.forward_kinematics(joint_angles=joint_rand.tolist(), tool=tool_name)

            d_fk = cal_pose_deviation(pose1=fk_rm_p1, pose2=fk_qp_p1)
            print(f'fk_qp_p1 = {fk_qp_p1}, fk_rm_p1 = {fk_rm_p1}, d_fk = {d_fk}\n')


            #   ----------- ik ----------------
            t_p = fk_rm_p1
            joint_rand_init = np.random.uniform(ub, lb)
            print(f'the guess is {joint_rand_init}')

            ret_qp, q = robot_kine_qp.inverse_kinematics( target_position=t_p[0:3], target_rpy=t_p[3:6], initial_guess=joint_rand_init, tool=tool_name)

            if ret_qp == 0:
                fk_qp_p2 = robot_kine_qp.forward_kinematics(q, tool=tool_name)
                d_p_ik = cal_pose_deviation(pose1=t_p, pose2=fk_qp_p2)
                print(f'-- success, in the qp ik, fk_qp_p2 = {fk_qp_p2}, d_p_ik = {d_p_ik}')
                if d_p_ik < 0.01:
                    result[0][1] += 1

                robot_mjk.send_command(q)
                robot_mjk.wait_until_reached()
                robot_mjk.print_state()
            else:
                fk_qp_p2 = robot_kine_qp.forward_kinematics(q, tool=tool_name)
                d_p_ik = cal_pose_deviation(pose1=t_p, pose2=fk_qp_p2)
                print(f'-- fail, in the qp ik, fk_qp_p2 = {fk_qp_p2}, d_p_ik = {d_p_ik},q = {q}, ret_qp = {ret_qp}')

            ret_rm, q = robot_kine_rm.inverse_kinematics(target_position=t_p[0:3], target_rpy=t_p[3:6], initial_guess=joint_rand_init, tool=tool_name)
            if ret_rm == 0:
                fk_rm_p2 = robot_kine_rm.forward_kinematics(joint_angles=q, tool=tool_name)
                d_p_ik = cal_pose_deviation(pose1=t_p, pose2=fk_rm_p2)
                print(f'== sucess, in the rm ik, fk_rm_p2 = {fk_rm_p2}, d_p_ik = {d_p_ik} ,q = {q}, ret_qp = {ret_qp}')
                if d_p_ik < 0.01:
                    result[1][1] += 1
            else:
                print(f'== fail in the rm ik, ret = {ret_rm}, q = {q}')

            if ret_qp == 0 or ret_rm == 0:
                solve_sum += 1

        print(f'result is {result}')
        print(f'solve_sum is {solve_sum}')


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] > pi:
            d_fk_p1[j] -= 2 * pi
        while d_fk_p1[j] < -pi:
            d_fk_p1[j] += 2 * pi
    d_fk = np.linalg.norm(d_fk_p1)
    return d_fk



if __name__ == "__main__":
    main()