#!/usr/bin/env python3
# -*-coding:utf8-*-
# 注意demo无法直接运行，需要pip安装sdk后才能运行
from typing import (
    Optional,
)
import time
from piper_sdk import *

import pygame
import threading
import time
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 EndPoseController:
    def __init__(self):
        pygame.init()
        pygame.joystick.init()
        if pygame.joystick.get_count() == 0:
            raise Exception("未检测到手柄")
        self.joystick = pygame.joystick.Joystick(0)
        self.joystick.init()
        
        # 末端初始姿态 [X, Y, Z, RX, RY, RZ]
        # {'X': 56127, 'Y': 0, 'Z': 213266, 'RX': 0, 'RY': 84999, 'RZ': 0, 'gripper': -70} # 右臂
        self.pose = [0.056127, 0, 0.213266, 0, 84.999, 0]
        self.speeds = [0.0] * 6
        
        # 夹爪状态和速度
        self.gripper = 0.0
        self.gripper_speed = 0.0

        # 末端位姿限制
        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.003  # 线性移动步长(m)，降低以提高精度
        self.angular_step = 0.1    # 角度步长(deg)，提高以便更好控制
        
        # 摇杆死区
        self.deadzone = 0.15
        
        # 控制模式
        self.fine_control_mode = False  # 精细控制模式
        
        self.running = True
        self.thread = threading.Thread(target=self.update_pose)
        self.thread.start()
        
        print("末端位姿控制器已启动")
        print("使用方向键上下左右控制XY，右摇杆控制Z")
        print("L1/R1控制RX，△/□控制RY，右摇杆Y轴控制RZ")
        print("按下L3(左摇杆按下)切换精细/普通控制模式")
    
    def update_pose(self):
        while self.running:
            try:
                pygame.event.pump()
            except Exception as e:
                print(f"控制器错误: {e}")
                self.stop()
                continue
            
            # 检查精细控制模式切换
            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)  # 防止多次触发
            
            # 根据控制模式调整步长
            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和RZ
            right_y = -self.joystick.get_axis(4)  # Z控制，取反使向上为正
            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
            
            # 设置XY速度（方向键控制）
            # self.speeds[0] = current_linear_step if hat_right else (-current_linear_step if hat_left else 0.0)  # X
            # self.speeds[1] = current_linear_step if hat_up else (-current_linear_step if hat_down else 0.0)     # Y

            self.speeds[0] = current_linear_step if hat_up else (-current_linear_step if hat_down else 0.0)  # X
            self.speeds[1] = current_linear_step if hat_left else (-current_linear_step if hat_right else 0.0)     # Y
            
            # 设置Z速度（右摇杆Y轴控制）
            self.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.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.speeds[4] = (current_angular_step if triangle else (-current_angular_step if square else 0.0))
            
            # 左摇杆Y轴控制RZ旋转
            self.speeds[5] = self._apply_nonlinear_mapping(left_y) * current_angular_step * 2  # RZ，增加系数使旋转更明显
            
            # 夹爪控制（圈/叉）
            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)
            
            # 更新末端位姿
            for i in range(6):
                self.pose[i] += self.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])
            
            # 更新夹爪
            self.gripper += self.gripper_speed
            self.gripper = max(0.0, min(0.08, self.gripper))
            
            time.sleep(0.02)
    
    def _normalize_angle(self, angle):
        """将角度归一化到[-180, 180]范围内"""
        while angle > 180:
            angle -= 360
        while angle < -180:
            angle += 360
        return angle
    
    def _apply_nonlinear_mapping(self, value):
        """应用非线性映射以提高控制精度"""
        # 保持符号，但对数值应用平方映射以提高精度
        sign = 1 if value >= 0 else -1
        return sign * (abs(value) ** 2)
    
    def get_action(self) -> Dict:
        """获取当前末端位姿和夹爪状态"""
        return {
            'X': self.pose[0]*1000,
            'Y': self.pose[1]*1000,
            'Z': self.pose[2]*1000,
            'RX': self.pose[3],
            'RY': self.pose[4],
            'RZ': self.pose[5],
            'gripper': self.gripper*1000
        }
    
    def stop(self):
        """停止控制器"""
        self.running = False
        if self.thread.is_alive():
            self.thread.join()
        pygame.quit()
        print("末端位姿控制器已退出")
    
    def reset(self):
        """重置到初始位姿"""
        self.pose = [0.056127, 0, 0.213266, 0, 84.999, 0]  # 使用您提供的初始位姿
        self.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)
    factor = 1000
    position = [
                55.0, \
                0.0, \
                206.0, \
                0, \
                85.0, \
                0, \
                0]
    # position = [0.0, \
    #             0.0, \
    #             80.0, \
    #             0, \
    #             203.386, \
    #             0, \
    #             0.8]
    count = 0
    teleop = EndPoseController()
    while True:
        print(piper.GetArmEndPoseMsgs())
        # print(piper.GetArmStatus())
        import time
        # count  = count + 1
        # # print(count)
        # if(count == 0):
        #     print("1-----------")
        #     position = [
        #         55.0, \
        #         0.0, \
        #         206.0, \
        #         0, \
        #         85.0, \
        #         0, \
        #         0]
        # elif(count == 200):
        #     print("2-----------")
        #     position = [
        #         55.0, \
        #         0.0, \
        #         260.0, \
        #         0, \
        #         85.0, \
        #         0, \
        #         0]
        # elif(count == 400):
        #     print("1-----------")
        #     position = [
        #         55.0, \
        #         0.0, \
        #         280.0, \
        #         0, \
        #         85.0, \
        #         0, \
        #         0]
        #     count = 0
        position = teleop.get_action()
        position = list(position.values())
        X = round(position[0]*factor)
        Y = round(position[1]*factor)
        Z = round(position[2]*factor)
        RX = round(position[3]*factor)
        RY = round(position[4]*factor)
        RZ = round(position[5]*factor)
        joint_6 = round(position[6]*factor)
        print(X,Y,Z,RX,RY,RZ, joint_6)
        # piper.MotionCtrl_1()
        piper.MotionCtrl_2(0x01, 0x00, 100, 0x00)
        piper.EndPoseCtrl(X,Y,Z,RX,RY,RZ)
        piper.GripperCtrl(abs(joint_6), 1000, 0x01, 0)
        
        time.sleep(0.1)
        pass