import mujoco
import numpy as np
import time
from mujoco import viewer

def main():
    # 加载模型
    model = mujoco.MjModel.from_xml_path('/home/ray/ROS2Workspace/softwaresystem/RobotOS1/urdf/robot3.xml')
    data = mujoco.MjData(model)
    
    # 获取执行器索引
    actuator_names = [
        "left_wheel2_vel", "left_wheel1_vel", 
        "right_wheel2_vel", "right_wheel1_vel"
    ]
    actuator_ids = [mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, name) for name in actuator_names]
    
    # 设置自转参数 (8秒一圈)
    rotation_time = 2.2  # 可调范围5-10秒
    
    # 机器人结构参数（需要根据实际模型调整）
    effective_wheel_radius = 0.12  # 轮子有效半径（米）
    track_width = 0.7  # 左右轮中心距（米）
    
    # 计算轮子角速度差
    angular_velocity = 2 * np.pi / rotation_time  # 机器人旋转角速度(rad/s)
    wheel_angular_difference = angular_velocity * track_width / (2 * effective_wheel_radius)
    
    print(f"目标旋转时间: {rotation_time}秒/圈")
    print(f"轮子角速度差: {wheel_angular_difference:.2f} rad/s")
    
    # 使用内置查看器
    with viewer.launch_passive(model, data) as viewer_handle:
        start_time = time.time()
        rotation_count = 0
        last_yaw = 0
        last_print_time = start_time
        
        print("开始模拟...按 Ctrl+C 停止")
        
        # 设置相机
        viewer_handle.cam.distance = 3.0
        viewer_handle.cam.elevation = -20
        viewer_handle.cam.azimuth = 45
        
        # 初始稳定阶段（1秒）
        init_stable_time = 0.5
        init_end_time = start_time + init_stable_time
        
        # 渐进加速参数
        acceleration_duration = 10.0  # 加速时间
        current_speed_factor = 0.0
        max_speed_factor = 1.0
        
        # 运行模拟
        while viewer_handle.is_running():
            current_time = time.time()
            elapsed_time = current_time - start_time
            
            # 初始稳定阶段
            if current_time < init_end_time:
                # 保持静止，让机器人稳定
                data.ctrl[actuator_ids] = [0, 0, 0, 0]
            else:
                # 渐进加速
                if current_speed_factor < max_speed_factor:
                    current_speed_factor = min(max_speed_factor, 
                                             (current_time - init_end_time) / acceleration_duration)
                
                # 应用速度控制
                left_speed = current_speed_factor * wheel_angular_difference
                right_speed = -current_speed_factor * wheel_angular_difference
                
                data.ctrl[actuator_ids[0]] = left_speed   # left_wheel2
                data.ctrl[actuator_ids[1]] = left_speed   # left_wheel1
                data.ctrl[actuator_ids[2]] = right_speed  # right_wheel2
                data.ctrl[actuator_ids[3]] = right_speed  # right_wheel1
            
            # 步进模拟
            mujoco.mj_step(model, data)
            
            # 更新查看器
            viewer_handle.sync()
            
            # 获取机器人方向（使用基座方向）
            base_quat = data.qpos[3:7]  # 基座的四元数
            roll, pitch, yaw = quaternion_to_euler(base_quat)
            
            # 检测完整旋转
            if last_yaw > 3.0 and yaw < -3.0:  # 从正π到负π
                rotation_count += 1
            elif last_yaw < -3.0 and yaw > 3.0:  # 从负π到正π
                rotation_count += 1
            last_yaw = yaw
            
            # 每秒显示状态
            if current_time - last_print_time >= 1.0:
                if rotation_count > 0:
                    actual_time_per_rotation = elapsed_time / rotation_count
                    print(f"目标: {rotation_time:.1f}秒/圈 | 实际: {actual_time_per_rotation:.1f}秒/圈 | 旋转圈数: {rotation_count}")
                else:
                    print(f"旋转中...当前偏航角: {yaw:.2f} rad | 速度因子: {current_speed_factor:.2f}")
                last_print_time = current_time
            
            # 控制模拟速度
            time.sleep(model.opt.timestep)
        
        print("模拟结束")

def quaternion_to_euler(q):
    """将四元数转换为欧拉角 (roll, pitch, yaw)"""
    w, x, y, z = q
    # 绕x轴旋转 (roll)
    sinr_cosp = 2 * (w * x + y * z)
    cosr_cosp = 1 - 2 * (x * x + y * y)
    roll = np.arctan2(sinr_cosp, cosr_cosp)
    
    # 绕y轴旋转 (pitch)
    sinp = 2 * (w * y - z * x)
    if np.abs(sinp) >= 1:
        pitch = np.copysign(np.pi / 2, sinp)
    else:
        pitch = np.arcsin(sinp)
    
    # 绕z轴旋转 (yaw)
    siny_cosp = 2 * (w * z + x * y)
    cosy_cosp = 1 - 2 * (y * y + z * z)
    yaw = np.arctan2(siny_cosp, cosy_cosp)
    
    return roll, pitch, yaw

if __name__ == "__main__":
    main()