import time
import math
import numpy as np
import mujoco
import mujoco.viewer
import sys
import threading
import select

# 加载模型
model_path = 'BrainCoLeftHandURDF/model/BrainCoLeftHandURDF.xml'
m = mujoco.MjModel.from_xml_path(model_path)
d = mujoco.MjData(m)

# 获取驱动器名称和索引的映射
actuator_names = [mujoco.mj_id2name(m, mujoco.mjtObj.mjOBJ_ACTUATOR, i) 
                  for i in range(m.nu)]
print("可用的驱动器:", actuator_names)

# 获取关节名称和索引的映射
joint_names = [mujoco.mj_id2name(m, mujoco.mjtObj.mjOBJ_JOINT, i) 
               for i in range(m.nv)]
print("可用的关节:", joint_names)

# 控制模式枚举
class ControlMode:
    SINE_WAVE = 0      # 正弦波控制
    GRASP = 1          # 抓取手势
    POINT = 2          # 指向手势
    FIST = 3           # 握拳
    OPEN = 4           # 张开
    CUSTOM = 5         # 自定义控制

# 控制参数
control_mode = ControlMode.SINE_WAVE
time_counter = 0.0

def get_actuator_id(name):
    """根据名称获取驱动器ID"""
    try:
        return mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_ACTUATOR, name)
    except:
        return -1

def set_control_sine_wave(d, t):
    """正弦波控制模式 - 所有手指做正弦运动"""
    # thumbCMC
    ctrl_id = get_actuator_id('thumbCMC')
    if ctrl_id >= 0:
        d.ctrl[ctrl_id] = 0.5 + 0.5 * math.sin(t * 0.5)
    
    # thumbMCP
    ctrl_id = get_actuator_id('thumbMCP')
    if ctrl_id >= 0:
        d.ctrl[ctrl_id] = 0.3 + 0.3 * math.sin(t * 0.5 + math.pi/4)
    
    # indexMCP
    ctrl_id = get_actuator_id('indexMCP')
    if ctrl_id >= 0:
        d.ctrl[ctrl_id] = 0.4 + 0.4 * math.sin(t * 0.6)
    
    # middleMCP
    ctrl_id = get_actuator_id('middleMCP')
    if ctrl_id >= 0:
        d.ctrl[ctrl_id] = 0.4 + 0.4 * math.sin(t * 0.6 + math.pi/6)
    
    # ringMCP
    ctrl_id = get_actuator_id('ringMCP')
    if ctrl_id >= 0:
        d.ctrl[ctrl_id] = 0.4 + 0.4 * math.sin(t * 0.6 + math.pi/3)
    
    # littleMCP
    ctrl_id = get_actuator_id('littleMCP')
    if ctrl_id >= 0:
        d.ctrl[ctrl_id] = 0.3 + 0.3 * math.sin(t * 0.6 + math.pi/2)

def set_control_grasp(d):
    """抓取手势 - 所有手指弯曲"""
    # thumbCMC - 拇指外展
    ctrl_id = get_actuator_id('thumbCMC')
    if ctrl_id >= 0:
        d.ctrl[ctrl_id] = 0.8
    
    # thumbMCP - 拇指弯曲
    ctrl_id = get_actuator_id('thumbMCP')
    if ctrl_id >= 0:
        d.ctrl[ctrl_id] = 0.7
    
    # 其他手指弯曲
    for name in ['indexMCP', 'middleMCP', 'ringMCP', 'littleMCP']:
        ctrl_id = get_actuator_id(name)
        if ctrl_id >= 0:
            d.ctrl[ctrl_id] = 0.8

def set_control_point(d):
    """指向手势 - 只有食指伸直，其他手指弯曲"""
    # thumbCMC - 拇指弯曲
    ctrl_id = get_actuator_id('thumbCMC')
    if ctrl_id >= 0:
        d.ctrl[ctrl_id] = 0.3
    
    # thumbMCP - 拇指弯曲
    ctrl_id = get_actuator_id('thumbMCP')
    if ctrl_id >= 0:
        d.ctrl[ctrl_id] = 0.5
    
    # indexMCP - 食指伸直
    ctrl_id = get_actuator_id('indexMCP')
    if ctrl_id >= 0:
        d.ctrl[ctrl_id] = 0.1
    
    # 其他手指弯曲
    for name in ['middleMCP', 'ringMCP', 'littleMCP']:
        ctrl_id = get_actuator_id(name)
        if ctrl_id >= 0:
            d.ctrl[ctrl_id] = 0.8

def set_control_fist(d):
    """握拳 - 所有手指完全弯曲"""
    # thumbCMC
    ctrl_id = get_actuator_id('thumbCMC')
    if ctrl_id >= 0:
        d.ctrl[ctrl_id] = 1.0
    
    # thumbMCP
    ctrl_id = get_actuator_id('thumbMCP')
    if ctrl_id >= 0:
        d.ctrl[ctrl_id] = 0.9
    
    # 其他手指完全弯曲
    for name in ['indexMCP', 'middleMCP', 'ringMCP', 'littleMCP']:
        ctrl_id = get_actuator_id(name)
        if ctrl_id >= 0:
            d.ctrl[ctrl_id] = 1.0

def set_control_open(d):
    """张开 - 所有手指伸直"""
    for name in ['thumbCMC', 'thumbMCP', 'indexMCP', 'middleMCP', 'ringMCP', 'littleMCP']:
        ctrl_id = get_actuator_id(name)
        if ctrl_id >= 0:
            d.ctrl[ctrl_id] = 0.0

def set_control_custom(d, values):
    """自定义控制 - 直接设置控制值
    values: 字典，键为驱动器名称，值为控制值 (0-1)
    """
    for name, value in values.items():
        ctrl_id = get_actuator_id(name)
        if ctrl_id >= 0:
            d.ctrl[ctrl_id] = max(0.0, min(1.0, value))  # 限制在 [0, 1] 范围

# 主控制循环
with mujoco.viewer.launch_passive(m, d) as viewer:
    print("\n=== 手部控制程序 ===")
    print("控制模式:")
    print("  0: 正弦波控制 (默认)")
    print("  1: 抓取手势")
    print("  2: 指向手势")
    print("  3: 握拳")
    print("  4: 张开")
    print("  5: 自定义控制")
    print("\n提示: 可以通过修改 control_mode 变量来切换控制模式")
    print("或者直接在代码中调用相应的控制函数\n")
    
    start_time = time.time()
    last_mode_change = time.time()
    
    while viewer.is_running():
        step_start = time.time()
        current_time = time.time() - start_time
        
        # 根据控制模式设置控制值
        if control_mode == ControlMode.SINE_WAVE:
            set_control_sine_wave(d, current_time)
        elif control_mode == ControlMode.GRASP:
            set_control_grasp(d)
        elif control_mode == ControlMode.POINT:
            set_control_point(d)
        elif control_mode == ControlMode.FIST:
            set_control_fist(d)
        elif control_mode == ControlMode.OPEN:
            set_control_open(d)
        elif control_mode == ControlMode.CUSTOM:
            # 自定义控制示例：可以在这里设置自定义值
            custom_values = {
                'thumbCMC': 0.5,
                'thumbMCP': 0.4,
                'indexMCP': 0.3,
                'middleMCP': 0.3,
                'ringMCP': 0.3,
                'littleMCP': 0.3
            }
            set_control_custom(d, custom_values)
        
        # 执行物理仿真步骤
        mujoco.mj_step(m, d)
        
        # 可选：每5秒自动切换模式（演示用）
        if current_time - last_mode_change > 5.0:
            last_mode_change = current_time
            # 可以在这里添加自动切换逻辑
            # control_mode = (control_mode + 1) % 6
        
        # 更新查看器
        with viewer.lock():
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = int(d.time % 2)
        
        viewer.sync()
        
        # 时间同步
        time_until_next_step = m.opt.timestep - (time.time() - step_start)
        if time_until_next_step > 0:
            time.sleep(time_until_next_step)

