import pygame
import time
import sys
from  lerobot.teleoperators import Teleoperator
from .config import XboxConfig
"""
Xbox类用于检测和监控飞行手柄。
"""
class Xbox(Teleoperator):
    config_class = XboxConfig
    name = "xbox"
    def __init__(self, config: XboxConfig):
        self.config = config
        # 控制参数
        self.linear_step = 0.0015  # 线性移动步长(m)
        self.angular_step = 0.001  # 角度步长(rad) - 从度转换为弧度
        # 摇杆死区
        self.deadzone = 0.15
        
        # 精细控制模式
        self.fine_control_mode = False
        self.reset()
    def reset(self):
        self.joystick = None
    @property
    def action_features(self):
        return {
            "dtype": "float32",
            "shape": (7,),
            "names": {"x.vel":0, "y.vel":1, "z.vel":2, "rx.vel.vel":3, "ry.vel":4, "rz.vel":5, "gripper.vel":6},
        }

    @property
    def feedback_features(self):
        return {}

    @property
    def is_connected(self):
        return self.joystick is not None

    def connect(self, calibrate = True):
        # 初始化pygame
        pygame.init()
        pygame.joystick.init()
        """检测连接的手柄"""
        joystick_count = pygame.joystick.get_count()
        print(f"检测到 {joystick_count} 个手柄设备")
        
        if joystick_count == 0:
            print("未检测到任何手柄设备！")
            return False

        if joystick_count<self.config.index+1:
            print(f"未检测到配置中指定的索引设备！") 
            return False
        # 初始化所有检测到的手柄
        self.joystick = pygame.joystick.Joystick(self.config.index)
        print(f"已连接设备：{self.joystick.get_name()}")

    @property
    def is_calibrated(self):
        return True

    def calibrate(self):
        pass

    def configure(self):
        pass

    def get_action(self):
        pygame.event.pump()
        pose_speeds = {}
        """更新末端位姿控制"""
        # 根据控制模式调整步长
        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)
        if self.joystick.get_button(0):  # A按钮
                self.fine_control_mode = True
        if self.joystick.get_button(1):  # B按钮
                self.fine_control_mode = False
        
        # print(f"步长设置 - 线性: {current_linear_step}, 角度: {current_angular_step}")
        print(f"精细控制模式: {self.fine_control_mode}")
        
        # 方向键控制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+
        # print(f"方向键状态: up={hat_up}, down={hat_down}, left={hat_left}, right={hat_right}")
        # 计算各轴速度
        pose_speeds['x.vel'] = current_linear_step if hat_left else (-current_linear_step if hat_right else 0.0)  # X
        pose_speeds['y.vel'] = -current_linear_step if hat_up else (current_linear_step if hat_down else 0.0)  # Y

        # 左右扳机控制z轴
        right_tiger = ((self.joystick.get_axis(5)+1)/2)
        left_tiger = ((self.joystick.get_axis(4)+1)/2)
        z_mapping = self.apply_nonlinear_mapping(left_tiger - right_tiger)
        pose_speeds['z.vel'] = z_mapping * current_angular_step * 2

        #右摇杆X轴控制RX旋转
        # 设置RX旋转速度
        rx_mapping = self.apply_nonlinear_mapping(self.joystick.get_axis(2))
        # print(f"RX轴非线性映射: {x_axis} -> {rx_mapping}")
        pose_speeds['rx.vel'] = rx_mapping * current_angular_step * 2  # RX

        # 右摇杆Y轴控制RY旋转
        # 设置RY旋转速度
        ry_mapping = self.apply_nonlinear_mapping(-self.joystick.get_axis(3))
        # print(f"RY轴非线性映射: {y_axis} -> {ry_mapping}")
        pose_speeds['ry.vel'] = ry_mapping * current_angular_step * 2  # RY
    
        # 左摇杆X轴RZ旋转
        # 设置RZ旋转速度
        rz_mapping = self.apply_nonlinear_mapping(-self.joystick.get_axis(0))
        # print(f"RZ轴非线性映射: {z_axis} -> {rz_mapping}")
        pose_speeds['rz.vel'] = rz_mapping * current_angular_step * 2  # RZ

        # 左摇杆Y轴控制夹爪
        # 设置夹爪开合速度
        gripper_mapping = -self.apply_nonlinear_mapping(-self.joystick.get_axis(1))
        # print(f"gripper轴非线性映射: {self.joystick.get_axis(0)} -> {gripper_mapping}")
        pose_speeds['gripper.vel'] = gripper_mapping * current_angular_step * 2  # RY
        print(f"pose_speeds: {pose_speeds}")
        
        return pose_speeds
    
    def apply_nonlinear_mapping(self, value):
        """应用非线性映射以提高控制精度"""
        # 保持符号，但对数值应用平方映射以提高精度
        value = 0.0 if abs(value) < self.deadzone else value
        sign = 1 if value >= 0 else -1
        return sign * (abs(value) ** 2)
    def normalize_value(self, v1,v2):
        """把两个轴归一化到一个轴上"""
        return v1-v2

    def send_feedback(self, feedback):
        pass

    def disconnect(self):
        self.reset()
        pygame.joystick.quit()
        pygame.quit()