# -*- coding: utf-8 -*-
"""
@file: servoj_demo.py
@brief: 实时模式 - servoJ功能demo
@copyright: Copyright (C) 2025 ROKAE (Beijing) Technology Co., LTD. All Rights Reserved.
Information in this file is the intellectual property of Rokae Technology Co., Ltd,
And may contains trade secrets that must be stored and viewed confidentially.
"""

import setup_path
import platform
import math
import time

# 根据操作系统导入相应的模块
if platform.system() == "Windows":
    from Release.windows import xCoreSDK_python
    from Release.windows.xCoreSDK_python import RtSupportedFields
    from Release.windows.xCoreSDK_python import RtControllerMode
elif platform.system() == "Linux":
    from Release.linux import xCoreSDK_python
    from Release.linux.xCoreSDK_python import RtSupportedFields
    from Release.linux.xCoreSDK_python import RtControllerMode
else:
    raise ImportError("Unsupported operating system")
from log import print_log, print_separator
from datetime import timedelta

M_PI = math.pi

# 用户指令下发周期(s)
PLAN_PERIOD = 0.02


def busy_wait(milliseconds: int):
    """延迟发送，与指令周期对齐"""
    end = time.perf_counter() + milliseconds / 1000.0
    while time.perf_counter() < end:
        pass


def servoj_op(robot: xCoreSDK_python.xMateRobot, ec: dict):
    """
    实时模式 servoJ 示例：周期性 sendCommand 下发关节位置。
    对应 C++ 示例 servoj_demo.cpp。
    """
    print_separator("servoj_op", length=110)
    robot.setOperateMode(xCoreSDK_python.OperateMode.automatic, ec)
    print_log("setOperateMode", ec)
    # 必做：启用实时模式
    robot.setMotionControlMode(xCoreSDK_python.MotionControlMode.RtCommandMode, ec)
    print_log("setMotionControlMode", ec)
    robot.setPowerState(True, ec)
    print_log("setPowerState", ec)

    rtCon: xCoreSDK_python.motioncontrolRT.PyRTmotioncontrol6 = (
        robot.getRtMotionController()
    )
    # 设置要接收的数据
    robot.startReceiveRobotState(
        timedelta(milliseconds=1), [RtSupportedFields.jointPos_m]
    )

    q_drag = [0, M_PI / 6, -M_PI / 2, 0, -M_PI / 3, 0]
    q2_drag = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]

    while robot.updateRobotState(timedelta(milliseconds=0)):
        pass

    jntPos = list(robot.jointPos(ec))[:6]
    # 运行至拖拽位（MoveJ 仅接受 6 轴关节角）
    rtCon.MoveJ(0.3, jntPos, q_drag)
    jntPos = list(robot.jointPos(ec))[:6]

    # 必做：启用 servoj 功能
    # 周期下发不稳定时，可使用前瞻时间和增益调节，默认前瞻时间为一个指令周期，增益为0
    rtCon.setServoJoint(PLAN_PERIOD, PLAN_PERIOD, 0, ec)
    print_log("setServoJoint", ec)
    jntPos = list(robot.jointPos(ec))[:6]
    # 必做：设置运动模式，注意在启用 servoj 功能后
    rtCon.startMove(RtControllerMode.jointPosition)

    start = time.perf_counter()
    while True:
        robot.updateRobotState(timedelta(milliseconds=1))
        elapsed_seconds = time.perf_counter() - start
        delta_angle = M_PI / 30.0 * (1 - math.cos(M_PI / 2.5 * elapsed_seconds))
        cmd = xCoreSDK_python.JointPosition()
        cmd.joints = [
            jntPos[0] + delta_angle,
            jntPos[1] + delta_angle,
            jntPos[2] - delta_angle,
            jntPos[3] + delta_angle,
            jntPos[4] - delta_angle,
            jntPos[5] + delta_angle,
        ]
        # 必做：发送计算出的关节位置
        rtCon.sendCommand(cmd)
        # 必做：间隔指令周期发送
        busy_wait(int(PLAN_PERIOD * 1000))

        if elapsed_seconds > 60:
            cmd.setFinished()
            rtCon.sendCommand(cmd)
            # 必做：关闭 servoj 功能
            rtCon.stopServoJoint()
            break

    while robot.updateRobotState(timedelta(milliseconds=0)):
        pass
    rtCon.MoveJ(0.3, list(robot.jointPos(ec))[:6], q2_drag)
    print("控制结束")

    robot.setMotionControlMode(
        xCoreSDK_python.MotionControlMode.NrtCommandMode, ec
    )
    print_log("setMotionControlMode", ec)
    robot.setOperateMode(xCoreSDK_python.OperateMode.manual, ec)
    print_log("setOperateMode", ec)


if __name__ == "__main__":
    try:
        remote_ip = "192.168.0.160"
        local_ip = "192.168.0.2"
        robot = xCoreSDK_python.xMateRobot(remote_ip, local_ip)
        ec = {}
        servoj_op(robot, ec)
    except Exception as e:
        print(f"An error occurred: {e}")
