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

from rm75_kine_qp import KinematicsSolver as kine_ctrl
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 demo_position_control():
    """Demonstrate pure position control"""

    urdf_path = "/home/zl/Downloads/urdf_rm75/RM75-B.urdf"



    # Create controller
    robot_mjk = MuJoCoPositionController(urdf_path, smoothness=0.05, enable_viewer=True)
    robot_mjk.start()
    time.sleep(1)

    print("\n[Test 1] Move joint 1 to 45 degrees")
    robot_mjk.send_command([0.785, 0, 0, 0, 0, 0, 0])
    robot_mjk.wait_until_reached()
    robot_mjk.print_state()
    time.sleep(0.5)

    print("\n[Test 2] Move joint 2 to -30 degrees")
    robot_mjk.send_command([0, -0.524, 0, 0, 0, 0, 0])
    robot_mjk.wait_until_reached()
    robot_mjk.print_state()
    time.sleep(0.5)

    print("\n[Test 3] Move multiple joints simultaneously")
    robot_mjk.send_command([0.5, -0.4, 0.3, 0.2, 0.1, 0, 0])
    robot_mjk.wait_until_reached()
    robot_mjk.print_state()
    time.sleep(0.5)

    print("\n[Test 4] Return home")
    robot_mjk.send_command([0, 0, 0, 0, 0, 0, 0])
    robot_mjk.wait_until_reached()
    robot_mjk.print_state()





    robot_kine = kine_ctrl()

    # Test 1: Forward Kinematics
    print("\n1. Forward Kinematics Test")
    print("-" * 40)

    tool_name = "scissor"
    joint_angles_zero = [0.1] * 7
    fk_result = robot_kine.forward_kinematics(joint_angles_zero, tool=tool_name)

    # Test 2: Inverse Kinematics with more reachable target
    print("\n2. Inverse Kinematics Test")
    print("-" * 40)

    # Try a simpler target first
    target_pos = [0.3, 0.2, 0.4]  # More reachable position
    target_rpy = [0.0, 0.0, radians(45)]  # Simpler orientation

    print(f"Target: ({target_pos[0]:.3f}, {target_pos[1]:.3f}, {target_pos[2]:.3f}) m")

    init_joints = [0.2] * 7
    time0 = time.time()
    for ii in range(100):
        joint_solution, success, error = robot_kine.inverse_kinematics(
            target_pos, target_rpy=target_rpy, initial_guess=init_joints,
            max_iter=500, debug=False, tool=tool_name
        )
    time1 = time.time()
    print(f"Time: {time1 - time0}")

    if success:
        print(f"✓ Solution found! Error: {error:.6f} m")
        for i, angle in enumerate(joint_solution):
            print(f"  Joint {i + 1}: {degrees(angle):7.2f}°")

        # Verify
        fk_verify = robot_kine.forward_kinematics(joint_solution,tool=tool_name)
        print(
            f"  Position: ({fk_verify['position'][0]:.3f}, {fk_verify['position'][1]:.3f}, {fk_verify['position'][2]:.3f}) m")
    else:
        print("✗ IK failed to find a solution!")

    # Test 3: Jacobian
    print("\n3. Jacobian Matrix")
    print("-" * 40)

    J = robot_kine.compute_jacobian(joint_angles_zero, tool=tool_name)
    print(f"Jacobian shape: {J.shape}")
    for i in range(min(3, J.shape[0])):
        row_str = "  ".join([f"{J[i, j]:7.3f}" for j in range(7)])
        print(f"  Row {i + 1}: {row_str}")

    # Test 4: Trajectory Planning with reachable positions
    print("\n4. Cartesian Trajectory Planning")
    print("-" * 40)

    start_pos = [0.3, 0.0, 0.4]  # Start position
    end_pos = [0.3, 0.0, 0.55]  # End position (smaller movement)

    fk0 = robot_kine.forward_kinematics([0.1] * 7, tool=tool_name)

    trajectory = robot_kine.plan_cartesian_trajectory(
        start_pos,
        end_pos,
        start_rpy=fk0['rpy'],
        end_rpy=[
            fk0['rpy'][0] + radians(10),
            fk0['rpy'][1],
            fk0['rpy'][2]
        ],
        num_steps=10,
        tool=tool_name
    )

    if trajectory:
        print(f"\n✓ Generated {len(trajectory)} waypoints")

    if success:
        print("✓ Inverse kinematics working (with simplified target)")
    else:
        print("⚠ Inverse kinematics may need tuning - try different targets")


    print("\n" + "=" * 60)
    print(f'test subchain Jacobian, for future obstacle avoidance')
    frame_names = [
        "link_2",
        "link_4",
        "link_7"
    ]
    Js_sub = robot_kine.get_subchain_jacobian(
        joint_angles=joint_angles_zero,
        frame_names=frame_names
    )
    print(f'Js_sub: {Js_sub}')


    #   ----------      rm75 official algorithm         -----------
    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
    robot_kine_rm = Algo(arm_model, force_type)
    frame = rm_frame_t("work", [0.0, 0.0, 0.0, 0.0, 0, 0.0])
    robot_kine_rm.rm_algo_set_workframe(frame)
    print(robot_kine_rm.rm_algo_get_curr_workframe())
    frame = rm_frame_t("work", [0.0, 0.0, 0.0, 0.0, 0, 0.0])
    robot_kine_rm.rm_algo_set_toolframe(frame)
    print(robot_kine_rm.rm_algo_get_curr_toolframe())
    joint_max_limit = np.array([
            3.14159, 2.2689, 3.14159, 2.3562, 3.14159, 2.234, 3.14159
        ])*57
    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
    robot_kine_rm.rm_algo_set_joint_min_limit(joint_max_limit.tolist())

    q_ref = [0.0, 110.0, 20.0, 40.0, 30.0, 180.0, 20.0]
    ret, phi = robot_kine_rm.rm_algo_calculate_arm_angle_from_config_rm75(q_ref)
    params = rm_inverse_kinematics_params_t([0.0, 110.0, 20.0, 40.0, 30.0, 180.0, 20.0],
                                            [0.3, 0.0, 0.3, 3.14, 0.0, 3.14], 1)
    ret, q_out = robot_kine_rm.rm_algo_inverse_kinematics_rm75_for_arm_angle(params, phi)
    print(f"rm_algo_inverse_kinematics_rm75_for_arm_angle ret: {ret} q_out: {q_out}")




    try:
        while robot_mjk.viewer and robot_mjk.viewer.is_running():
            time.sleep(0.1)
    except KeyboardInterrupt:
        pass
    robot_mjk.stop()



def main():

    demo_position_control()



if __name__ == "__main__":
    main()