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

# 加载模型
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
mode_changed = False  # 标志，表示模式已更改

def get_mode_name(mode):
    """获取模式名称"""
    mode_names = {
        ControlMode.SINE_WAVE: "正弦波控制",
        ControlMode.GRASP: "抓取手势",
        ControlMode.POINT: "指向手势",
        ControlMode.FIST: "握拳",
        ControlMode.OPEN: "张开",
        ControlMode.CUSTOM: "自定义控制"
    }
    return mode_names.get(mode, "未知模式")

def process_input(input_str):
    """处理用户输入"""
    global control_mode, mode_changed
    
    input_str = input_str.strip()
    
    if input_str.lower() == 'q':
        print("\n退出程序...")
        sys.exit(0)
    elif input_str.isdigit():
        mode_num = int(input_str)
        if 0 <= mode_num <= 5:
            control_mode = mode_num
            mode_changed = True
            print(f"\n>>> [模式已切换] {get_mode_name(control_mode)} <<<")
        else:
            print(f"\n无效的模式编号: {mode_num}，请输入 0-5")
    elif input_str:
        print(f"\n无效的输入: '{input_str}'，请输入数字 0-5 或 'q'")

def select_control_mode():
    """交互式选择控制模式"""
    print("\n" + "="*50)
    print("=== 手部控制程序 ===")
    print("="*50)
    print("\n请选择控制模式:")
    print("  0: 正弦波控制 - 所有手指做正弦运动")
    print("  1: 抓取手势 - 所有手指弯曲")
    print("  2: 指向手势 - 只有食指伸直")
    print("  3: 握拳 - 所有手指完全弯曲")
    print("  4: 张开 - 所有手指伸直")
    print("  5: 自定义控制 - 手动设置每个驱动器")
    print("\n提示: 在程序运行过程中，可以按数字键 0-5 切换模式")
    print("按 'q' 或 'Q' 退出程序\n")
    
    while True:
        try:
            choice = input("请输入模式编号 (0-5，默认0): ").strip()
            if choice == "":
                return ControlMode.SINE_WAVE
            choice = int(choice)
            if 0 <= choice <= 5:
                return choice
            else:
                print("无效的选择，请输入 0-5 之间的数字")
        except ValueError:
            print("无效的输入，请输入数字")
        except KeyboardInterrupt:
            print("\n程序已取消")
            sys.exit(0)

def keyboard_listener():
    """键盘输入监听线程 - 在后台持续监听键盘输入"""
    global control_mode, mode_changed
    
    import select
    
    # 尝试使用 termios 设置终端为原始模式（Linux/Mac）
    try:
        import termios
        import tty
        fd = sys.stdin.fileno()
        old_settings = termios.tcgetattr(fd)
        tty.setcbreak(fd)  # 设置为 cbreak 模式，可以立即读取字符
        termios_available = True
    except (ImportError, OSError, AttributeError):
        # Windows 或不支持 termios 的系统
        termios_available = False
        old_settings = None
        fd = None
    
    print("\n键盘监听已启动！")
    print("在查看器运行时，可以在终端输入数字 0-5 切换模式，输入 'q' 退出\n")
    
    buffer = ""
    
    while True:
        try:
            if termios_available:
                # 使用 select 非阻塞检查是否有输入
                if select.select([sys.stdin], [], [], 0.1)[0]:
                    char = sys.stdin.read(1)
                    
                    if char == '\x1b':  # ESC 键
                        break
                    elif char == '\n' or char == '\r':  # 回车
                        if buffer:
                            process_input(buffer.strip())
                            buffer = ""
                    elif char == '\x7f' or char == '\b':  # 退格
                        if buffer:
                            buffer = buffer[:-1]
                    elif char.lower() == 'q':
                        print("\n退出程序...")
                        if old_settings:
                            termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
                        sys.exit(0)
                    else:
                        buffer += char
                        # 如果是数字，立即处理
                        if char.isdigit():
                            process_input(char)
                            buffer = ""
            else:
                # 如果没有 termios，尝试简单的非阻塞读取
                try:
                    if select.select([sys.stdin], [], [], 0.5)[0]:
                        line = sys.stdin.readline().strip()
                        if line:
                            process_input(line)
                except:
                    time.sleep(0.5)
        except KeyboardInterrupt:
            break
        except Exception as e:
            # 如果出错，继续运行
            time.sleep(0.1)
    
    # 恢复终端设置
    if termios_available and old_settings and fd is not None:
        try:
            termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
        except:
            pass

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] 范围

# 启动前选择控制模式
control_mode = select_control_mode()
print(f"\n已选择模式: {get_mode_name(control_mode)}")
print("正在启动查看器...\n")

# 启动键盘监听线程
keyboard_thread = threading.Thread(target=keyboard_listener, daemon=True)
keyboard_thread.start()

# 主控制循环
with mujoco.viewer.launch_passive(m, d) as viewer:
    print(f"当前控制模式: {get_mode_name(control_mode)}")
    print("查看器已启动，手部模型正在运行...\n")
    
    start_time = time.time()
    last_mode_change = time.time()
    last_mode_display = control_mode
    
    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)
        
        # 如果模式改变，显示提示
        if mode_changed and control_mode != last_mode_display:
            print(f"[模式已切换] {get_mode_name(control_mode)}")
            last_mode_display = control_mode
            mode_changed = False
        
        # 更新查看器
        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)

