#!/usr/bin/env python3
# -*-coding:utf8-*-
from typing import Optional
import time
from piper_sdk import *
import pygame
import threading
from typing import Dict


def enable_fun(piper:C_PiperInterface_V2):
    '''
    使能机械臂并检测使能状态,尝试5s,如果使能超时则退出程序
    '''
    enable_flag = False
    # 设置超时时间（秒）
    timeout = 5
    # 记录进入循环前的时间
    start_time = time.time()
    elapsed_time_flag = False
    while not (enable_flag):
        elapsed_time = time.time() - start_time
        print("--------------------")
        enable_flag = piper.GetArmLowSpdInfoMsgs().motor_1.foc_status.driver_enable_status and \
            piper.GetArmLowSpdInfoMsgs().motor_2.foc_status.driver_enable_status and \
            piper.GetArmLowSpdInfoMsgs().motor_3.foc_status.driver_enable_status and \
            piper.GetArmLowSpdInfoMsgs().motor_4.foc_status.driver_enable_status and \
            piper.GetArmLowSpdInfoMsgs().motor_5.foc_status.driver_enable_status and \
            piper.GetArmLowSpdInfoMsgs().motor_6.foc_status.driver_enable_status
        print("使能状态:",enable_flag)
        piper.EnableArm(7)
        piper.GripperCtrl(0,1000,0x01, 0)
        print("--------------------")
        # 检查是否超过超时时间
        if elapsed_time > timeout:
            print("超时....")
            elapsed_time_flag = True
            enable_flag = True
            break
        time.sleep(1)
        pass
    if(elapsed_time_flag):
        print("程序自动使能超时,退出程序")
        exit(0)


class UnifiedArmController:
    def __init__(self):
        # 初始化pygame和手柄
        pygame.init()
        pygame.joystick.init()
        
        # 检查是否有连接的手柄
        if pygame.joystick.get_count() == 0:
            raise Exception("未检测到手柄")
        
        # 初始化手柄
        self.joystick = pygame.joystick.Joystick(0)
        self.joystick.init()
        
        # 控制模式标志 - True为末端位姿控制，False为关节控制
        self.end_pose_mode = False
        
        # 摇杆死区
        self.deadzone = 0.15
        
        # 精细控制模式
        self.fine_control_mode = False
        
        # 初始化关节状态
        self.joints = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]  # 6个关节
        self.joint_speeds = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]  # 6个关节的速度
        
        # 关节弧度限制
        self.joint_limits = [
            (-92000 / 57324.840764, 92000 / 57324.840764),  # joint1
            ( 0 / 57324.840764, 120000 / 57324.840764),   # joint2
            (-80000 / 57324.840764, 0 / 57324.840764),   # joint3
            (-90000 / 57324.840764, 90000 / 57324.840764),  # joint4
            (-65000 / 57324.840764, 65000 / 57324.840764),  # joint5
            (-90000 / 57324.840764, 90000 / 57324.840764)   # joint6
        ]
        
        # 初始化末端姿态 [X, Y, Z, RX, RY, RZ]
        self.pose = [0.056127, 0, 0.213266, 0, 84.999, 0]
        self.pose_speeds = [0.0] * 6
        
        # 末端位姿限制
        self.pose_limits = [
            (-0.6, 0.6),   # X (m)
            (-0.6, 0.6),   # Y (m)
            (0.05, 0.6),   # Z (m) - 设置最小高度防止碰撞
            (-180, 180),   # RX (deg)
            (-180, 180),   # RY (deg)
            (-180, 180)    # RZ (deg)
        ]
        
        # 控制参数
        self.linear_step = 0.0015  # 线性移动步长(m)
        self.angular_step = 0.05   # 角度步长(deg)
        self.joint_step = 0.015    # 关节步长(rad)
        
        # 夹爪状态和速度
        self.gripper = 0.0
        self.gripper_speed = 0.0
        
        # 启动更新线程
        self.running = True
        self.thread = threading.Thread(target=self.update_controller)
        self.thread.start()
        
        print("机械臂统一控制器已启动")
        print("按下OPTIONS(Start)切换控制模式")
        print("当前模式:", "末端位姿控制" if self.end_pose_mode else "关节控制")
    
    def _apply_nonlinear_mapping(self, value):
        """应用非线性映射以提高控制精度"""
        # 保持符号，但对数值应用平方映射以提高精度
        sign = 1 if value >= 0 else -1
        return sign * (abs(value) ** 2)
    
    def _normalize_angle(self, angle):
        """将角度归一化到[-180, 180]范围内"""
        while angle > 180:
            angle -= 360
        while angle < -180:
            angle += 360
        return angle
    
    def update_controller(self):
        while self.running:
            try:
                pygame.event.pump()
            except Exception as e:
                print(f"控制器错误: {e}")
                self.stop()
                continue
            
            # 检查控制模式切换 (使用左摇杆按钮)
            if self.joystick.get_button(9):  # 左摇杆按钮
                self.end_pose_mode = not self.end_pose_mode
                print(f"切换到{'末端位姿控制' if self.end_pose_mode else '关节控制'}模式")
                time.sleep(0.3)  # 防止多次触发
            
            # 检查精细控制模式切换 (使用L3按钮)
            if self.joystick.get_button(10):  # L3按钮
                self.fine_control_mode = not self.fine_control_mode
                print(f"切换到{'精细' if self.fine_control_mode else '普通'}控制模式")
                time.sleep(0.3)  # 防止多次触发
            
            # 检查重置按钮 (7号按钮，通常是Start按钮)
            if self.joystick.get_button(7):  # Start按钮
                print("重置机械臂到0位...")
                # 重置关节和位姿
                self.joints = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
                self.joint_speeds = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
                self.gripper = 0.0
                self.gripper_speed = 0.0
                # 临时切换到关节控制模式
                self.end_pose_mode = False
                print("机械臂已重置到0位")
                time.sleep(0.3)  # 防止多次触发

            # 根据控制模式获取输入并更新状态
            if self.end_pose_mode:
                self.update_end_pose()
            else:
                self.update_joints()
            
            # 夹爪控制（圈/叉）- 两种模式下都保持一致
            circle = self.joystick.get_button(1)  # 夹爪开
            cross = self.joystick.get_button(0)   # 夹爪关
            self.gripper_speed = 0.01 if circle else (-0.01 if cross else 0.0)
            
            # 更新夹爪
            self.gripper += self.gripper_speed
            self.gripper = max(0.0, min(0.08, self.gripper))
            
            time.sleep(0.02)
    
    def update_end_pose(self):
        """更新末端位姿控制"""
        # 根据控制模式调整步长
        current_linear_step = self.linear_step * (0.1 if self.fine_control_mode else 1.0)
        current_angular_step = self.angular_step * (0.1 if self.fine_control_mode else 1.0)
        
        # 方向键控制XY
        hat = self.joystick.get_hat(0)
        hat_up = hat[1] == 1     # Y+
        hat_down = hat[1] == -1  # Y-
        hat_left = hat[0] == -1  # X-
        hat_right = hat[0] == 1  # X+
        
        # 右摇杆控制Z
        right_y = -self.joystick.get_axis(4)  # Z控制，取反使向上为正
        
        # 左摇杆控制RZ
        left_y = -self.joystick.get_axis(1)   # RZ控制，取反使向上为正
        
        # 应用死区
        right_y = 0.0 if abs(right_y) < self.deadzone else right_y
        left_y = 0.0 if abs(left_y) < self.deadzone else left_y

        self.pose_speeds[0] = current_linear_step if hat_up else (-current_linear_step if hat_down else 0.0)  # X
        self.pose_speeds[1] = current_linear_step if hat_left else (-current_linear_step if hat_right else 0.0)  # Y
        
        # 设置Z速度（右摇杆Y轴控制）
        self.pose_speeds[2] = self._apply_nonlinear_mapping(right_y) * current_linear_step  # Z
        
        # L1/R1控制RX旋转
        LB = self.joystick.get_button(4)  # RX-
        RB = self.joystick.get_button(5)  # RX+
        self.pose_speeds[3] = (-current_angular_step if LB else (current_angular_step if RB else 0.0))
        
        # △/□控制RY旋转
        triangle = self.joystick.get_button(2)  # RY+
        square = self.joystick.get_button(3)    # RY-
        self.pose_speeds[4] = (current_angular_step if triangle else (-current_angular_step if square else 0.0))
        
        # 左摇杆Y轴控制RZ旋转
        self.pose_speeds[5] = self._apply_nonlinear_mapping(left_y) * current_angular_step * 2  # RZ，增加系数使旋转更明显
        
        # 更新末端位姿
        for i in range(6):
            self.pose[i] += self.pose_speeds[i]
        
        # 位置限制
        for i in range(3):
            min_val, max_val = self.pose_limits[i]
            self.pose[i] = max(min_val, min(max_val, self.pose[i]))
        
        # 角度归一化处理
        for i in range(3, 6):
            self.pose[i] = self._normalize_angle(self.pose[i])
    
    def update_joints(self):
        """更新关节控制"""
        # 根据控制模式调整步长
        current_joint_step = self.joint_step * (0.1 if self.fine_control_mode else 1.0)
        
        # 使用类似于末端位姿控制的映射，但直接控制关节
        
        # 左摇杆控制关节1和2 (类似于末端位姿控制中的X和Y)
        left_x = -self.joystick.get_axis(0)  # 左摇杆x轴
        left_y = -self.joystick.get_axis(1)  # 左摇杆y轴
        
        # 应用死区
        left_x = 0.0 if abs(left_x) < self.deadzone else left_x
        left_y = 0.0 if abs(left_y) < self.deadzone else left_y
        
        # 右摇杆控制关节3和4
        right_x = self.joystick.get_axis(3)  # 右摇杆x轴
        right_y = self.joystick.get_axis(4)  # 右摇杆y轴
        
        # 应用死区
        right_x = 0.0 if abs(right_x) < self.deadzone else right_x
        right_y = 0.0 if abs(right_y) < self.deadzone else right_y
        
        # 方向键控制关节5和6
        hat = self.joystick.get_hat(0)
        up = hat[1] == 1
        down = hat[1] == -1
        left = hat[0] == -1
        right = hat[0] == 1
        
        # 映射输入到关节速度
        self.joint_speeds[0] = left_x * current_joint_step  # joint1速度
        self.joint_speeds[1] = left_y * current_joint_step  # joint2速度
        self.joint_speeds[2] = right_y * current_joint_step  # joint3速度
        self.joint_speeds[3] = right_x * current_joint_step  # joint4速度
        self.joint_speeds[4] = -current_joint_step if up else (current_joint_step if down else 0.0)  # joint5速度
        self.joint_speeds[5] = current_joint_step if right else (-current_joint_step if left else 0.0)  # joint6速度
        
        # 积分速度到关节位置
        for i in range(6):
            self.joints[i] += self.joint_speeds[i]
        
        # 关节范围保护
        for i in range(6):
            min_val, max_val = self.joint_limits[i]
            self.joints[i] = max(min_val, min(max_val, self.joints[i]))

    def update_state(self, ctrl_mode, end_pose, joint_state):
        if ctrl_mode == 'end_pose':
            _joint_state = [0] * 6
            _joint_state[0] = joint_state.joint_1 / 57324.840764
            _joint_state[1] = joint_state.joint_2 / 57324.840764
            _joint_state[2] = joint_state.joint_3 / 57324.840764
            _joint_state[3] = joint_state.joint_4 / 57324.840764
            _joint_state[4] = joint_state.joint_5 / 57324.840764
            _joint_state[5] = joint_state.joint_6 / 57324.840764
            self.joints = _joint_state
        else:
            _end_pose = [0] * 6
            _end_pose[0] = end_pose.X_axis / 1000 / 1000
            _end_pose[1] = end_pose.Y_axis / 1000 / 1000 
            _end_pose[2] = end_pose.Z_axis / 1000 / 1000 
            _end_pose[3] = end_pose.RX_axis / 1000 
            _end_pose[4] = end_pose.RY_axis / 1000 
            _end_pose[5] = end_pose.RZ_axis / 1000 
            self.pose = _end_pose
            
    
    def get_action(self) -> Dict:
        """获取当前控制命令"""
        if self.end_pose_mode:
            # 返回末端位姿
            return {
                'X': self.pose[0],
                'Y': self.pose[1],
                'Z': self.pose[2],
                'RX': self.pose[3],
                'RY': self.pose[4],
                'RZ': self.pose[5],
                'gripper': self.gripper
            }
        else:
            # 返回关节角度
            return {
                'joint0': self.joints[0],
                'joint1': self.joints[1],
                'joint2': self.joints[2],
                'joint3': self.joints[3],
                'joint4': self.joints[4],
                'joint5': self.joints[5],
                'gripper': self.gripper
            }
    
    def get_control_mode(self):
        """返回当前控制模式"""
        return "end_pose" if self.end_pose_mode else "joints"
    
    def stop(self):
        """停止控制器"""
        self.running = False
        if self.thread.is_alive():
            self.thread.join()
        pygame.quit()
        print("控制器已退出")
    
    def reset(self):
        """重置到初始状态"""
        self.joints = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        self.joint_speeds = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        self.pose = [0.056127, 0, 0.213266, 0, 84.999, 0]
        self.pose_speeds = [0.0] * 6
        self.gripper = 0.0
        self.gripper_speed = 0.0
        print("已重置到初始状态")


if __name__ == "__main__":
    piper = C_PiperInterface_V2("can0")
    piper.ConnectPort()
    piper.EnableArm(7)
    enable_fun(piper=piper)
    piper.GripperCtrl(0,1000,0x01, 0)
    
    teleop = UnifiedArmController()
    factor = 1000
    
    while True:
        # 获取当前控制命令
        action = teleop.get_action()
        control_mode = teleop.get_control_mode()
        
        if control_mode == "end_pose":
            # 末端位姿控制
            position = list(action.values())
            X = round(position[0]*factor*factor)
            Y = round(position[1]*factor*factor)
            Z = round(position[2]*factor*factor)
            RX = round(position[3]*factor)
            RY = round(position[4]*factor)
            RZ = round(position[5]*factor)
            joint_6 = round(position[6]*factor*factor)
            
            piper.MotionCtrl_2(0x01, 0x00, 100, 0x00)
            piper.EndPoseCtrl(X, Y, Z, RX, RY, RZ)
            piper.GripperCtrl(abs(joint_6), 1000, 0x01, 0)

            new_end_pose = piper.GetArmEndPoseMsgs().end_pose
            new_joint_state = piper.GetArmJointMsgs().joint_state
            teleop.update_state(control_mode, new_end_pose, new_joint_state)
            print("控制模式: 末端控制")
            print("末端位置", new_end_pose)
            print("关节位置:", new_joint_state)


        else:
            # 关节控制
            joints = list(action.values())
            # 将关节角度转换为适合发送的格式
            joint0 = round(joints[0] * 57324.840764)  # 转换为机械臂期望的单位
            joint1 = round(joints[1] * 57324.840764)
            joint2 = round(joints[2] * 57324.840764)
            joint3 = round(joints[3] * 57324.840764)
            joint4 = round(joints[4] * 57324.840764)
            joint5 = round(joints[5] * 57324.840764)
            gripper = round(joints[6] * 1000 * 1000)
            
            # 发送关节控制命令
            piper.MotionCtrl_2(0x01, 0x01, 100, 0x00)
            piper.JointCtrl(joint0, joint1, joint2, joint3, joint4, joint5)
            piper.GripperCtrl(abs(gripper), 1000, 0x01, 0)

            new_end_pose = piper.GetArmEndPoseMsgs().end_pose
            new_joint_state = piper.GetArmJointMsgs().joint_state
            teleop.update_state(control_mode, new_end_pose, new_joint_state)
            
            print("控制模式: 关节控制")
            print("末端位置", new_end_pose)
            print("关节位置:", new_joint_state)
            
        time.sleep(0.1)