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

class MuJoCoSimulator:
    def __init__(self, model_path):
        """初始化MuJoCo仿真器"""
        self.model = mujoco.MjModel.from_xml_path(model_path)
        self.data = mujoco.MjData(self.model)
        self.viewer = None
        self.simulation_time = 0
        
    def setup_control(self, control_mode='sine_wave', amplitude=1.0, frequency=1.0):
        """设置控制参数"""
        self.control_mode = control_mode
        self.amplitude = amplitude
        self.frequency = frequency
        
    def sine_wave_control(self, cnt, joint_index=0):
        """正弦波控制"""
        self.data.ctrl[joint_index] = self.amplitude * math.sin(2 * math.pi * self.frequency * cnt)
        
    def step_control(self, cnt, joint_index=0):
        """步进控制"""
        self.data.ctrl[joint_index] = self.amplitude * (1 if (cnt % 2) < 1 else -1)
        
    def run_simulation(self, duration=30, simulation_mode='step'):
        """
        运行仿真
        
        Parameters:
        - duration: 仿真持续时间(秒)
        - simulation_mode: 仿真模式 ('step', 'step1_step2', 'forward', 'inverse')
        """
        with mujoco.viewer.launch_passive(self.model, self.data) as viewer:
            self.viewer = viewer
            start_time = time.time()
            step_count = 0
            
            print(f"开始仿真 - 模式: {simulation_mode}, 持续时间: {duration}秒")
            
            while viewer.is_running() and (time.time() - start_time) < duration:
                step_start = time.time()
                
                # 执行不同的仿真步骤
                self._execute_simulation_step(simulation_mode, step_count)
                
                # 更新控制信号
                self._update_control_signals(step_count * 0.01)
                
                # 切换可视化选项（每2秒切换接触点显示）
                self._update_visualization()
                
                # 同步查看器
                viewer.sync()
                
                # 精确的时间控制
                self._precise_timing_control(step_start)
                
                step_count += 1
                self.simulation_time = step_count * self.model.opt.timestep
                
            print(f"仿真结束 - 总步数: {step_count}, 仿真时间: {self.simulation_time:.2f}秒")
    
    def _execute_simulation_step(self, mode, step_count):
        """执行仿真步骤"""
        if mode == 'step':
            # 标准步骤（包括正向和逆向动力学）
            mujoco.mj_step(self.model, self.data)
            
        elif mode == 'step1_step2':
            # 分离的步骤（可以在step1和step2之间插入控制）
            mujoco.mj_step1(self.model, self.data)
            # 在这里可以插入更复杂的控制逻辑
            mujoco.mj_step2(self.model, self.data)
            
        elif mode == 'forward':
            # 只执行正向动力学
            mujoco.mj_forward(self.model, self.data)
            # 打印状态信息（每100步打印一次）
            if step_count % 100 == 0:
                print(f"步数 {step_count}: qpos[0]={self.data.qpos[0]:.3f}, "
                      f"qvel[0]={self.data.qvel[0]:.3f}")
                      
        elif mode == 'inverse':
            # 逆向动力学计算
            self.data.qacc[0] = math.sin(step_count * 0.01)
            mujoco.mj_inverse(self.model, self.data)
            if step_count % 100 == 0:
                print(f"步数 {step_count}: qfrc_inverse[0]={self.data.qfrc_inverse[0]:.3f}")
    
    def _update_control_signals(self, cnt):
        """更新控制信号"""
        if self.control_mode == 'sine_wave':
            self.sine_wave_control(cnt, joint_index=1)
        elif self.control_mode == 'step':
            self.step_control(cnt, joint_index=1)
        # 可以添加更多控制模式
    
    def _update_visualization(self):
        """更新可视化选项"""
        if self.viewer:
            with self.viewer.lock():
                # 每2秒切换接触点显示
                self.viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = int(self.data.time % 2)
                # 可以添加其他可视化选项
                # self.viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_JOINT] = True
    
    def _precise_timing_control(self, step_start):
        """精确的时间控制"""
        time_until_next_step = self.model.opt.timestep - (time.time() - step_start)
        if time_until_next_step > 0:
            time.sleep(time_until_next_step)
    
    def get_simulation_info(self):
        """获取仿真信息"""
        return {
            'simulation_time': self.simulation_time,
            'joint_positions': self.data.qpos.copy(),
            'joint_velocities': self.data.qvel.copy(),
            'sensor_data': self.data.sensordata.copy() if hasattr(self.data, 'sensordata') else None
        }

def main():
    # 创建仿真器实例
    simulator = MuJoCoSimulator('BrainCoLeftHandURDF/model/BrainCoLeftHandURDF.xml')
    
    # 设置控制参数
    simulator.setup_control(control_mode='sine_wave', amplitude=0.5, frequency=0.5)
    
    # 运行仿真（可以选择不同的模式）
    print("可用仿真模式: 'step', 'step1_step2', 'forward', 'inverse'")
    
    # 运行标准步骤仿真
    simulator.run_simulation(duration=10, simulation_mode='step')
    
    # 获取仿真结果
    info = simulator.get_simulation_info()
    print(f"最终关节位置: {info['joint_positions']}")

if __name__ == "__main__":
    main()