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

# 目标Roll倾斜角度（单位：度）
TARGET_ROLL_DEGREE = 7.0  # 目标倾斜7度
ROLL_TRANSITION_DURATION = 1.5  # 倾斜过渡时间（秒）
STABLE_DURATION = 1.0  # 稳定阶段持续时间（秒）

def main():
    # 加载模型和创建模拟
    model = mujoco.MjModel.from_xml_path("/home/mrji/mujoco_learning/robot3/urdf/robot3.xml")  
    data = mujoco.MjData(model)
    
    # 完整重置机器人状态（第一次重置）
    mujoco.mj_resetData(model, data)
    
    # 设置初始关节角度为0
    data.qpos[model.joint("left_hip_pitch_joint").id] = 0
    data.qpos[model.joint("right_hip_pitch_joint").id] = 0
    data.qpos[model.joint("left_leg_pitch_joint").id] = 0
    data.qpos[model.joint("right_leg_pitch_joint").id] = 0
    
    # 重置所有速度为0
    data.qvel[:] = 0
    
    # 第二次重置模型数据（规避加载问题）
    mujoco.mj_resetData(model, data)
    
    # 再次设置初始关节角度为0
    data.qpos[model.joint("left_hip_pitch_joint").id] = 0
    data.qpos[model.joint("right_hip_pitch_joint").id] = 0
    data.qpos[model.joint("left_leg_pitch_joint").id] = 0
    data.qpos[model.joint("right_leg_pitch_joint").id] = 0
    
    # 创建查看器
    with mujoco.viewer.launch_passive(model, data) as viewer:
        # 设置摄像头位置
        viewer.cam.azimuth = 180  # 面向机器人正面
        viewer.cam.elevation = -20  # 稍微俯视
        viewer.cam.distance = 2.0  # 距离
        viewer.cam.lookat[:] = [0.0, 0.0, 0.65]  # 看向机器人中心
        
        # 获取执行器ID
        hip_left = model.actuator("1_pos").id
        leg_left = model.actuator("2_pos").id
        hip_right = model.actuator("3_pos").id
        leg_right = model.actuator("4_pos").id
        
        # 设置控制信号初始值为0
        data.ctrl[hip_left] = 0
        data.ctrl[leg_left] = 0
        data.ctrl[hip_right] = 0
        data.ctrl[leg_right] = 0

        # 修改：在最开始的时候额外重置一次
        mujoco.mj_resetData(model, data)
        
        # 状态控制变量
        start_time = time.time()
        current_phase = "STABLE"  # STABLE, MOVING, DONE
        transition_start_time = 0
        
        # 控制参数
        hip_adjustment = 0.1  # 髋关节调整量（根据实验确定）
        leg_adjustment = 0.15  # 腿关节调整量（根据实验确定）
        initial_roll = 0.0  # 初始Roll角度
        current_roll = 0.0  # 当前Roll角度
        
        # 初始稳定步骤 - 让物理引擎稳定
        for _ in range(200):
            mujoco.mj_step(model, data)
            viewer.sync()
            time.sleep(0.001)
        
        # 模拟循环
        while viewer.is_running():
            step_start = time.time()
            elapsed_time = time.time() - start_time
            
            # 状态机控制
            if current_phase == "STABLE":
                # 稳定阶段 - 保持初始姿态
                if elapsed_time >= STABLE_DURATION:
                    # 记录初始姿态（Roll角度）
                    quat = data.sensor("orientation").data.copy()
                    initial_roll = get_roll_from_quat(quat)
                    print(f"稳定阶段结束，初始Roll角度: {math.degrees(initial_roll):.2f}度")
                    
                    current_phase = "MOVING"
                    transition_start_time = time.time()
                    print("开始Roll倾斜过渡...")
            
            elif current_phase == "MOVING":
                # 过渡阶段 - 平滑过渡到目标倾斜
                transition_elapsed = time.time() - transition_start_time
                
                if transition_elapsed >= ROLL_TRANSITION_DURATION:
                    # 达到目标倾斜
                    data.ctrl[hip_left] = 0 - hip_adjustment
                    data.ctrl[leg_left] = 0 - leg_adjustment
                    data.ctrl[hip_right] = 0 + hip_adjustment
                    data.ctrl[leg_right] = 0 + leg_adjustment
                    
                    current_phase = "DONE"
                    print("达到目标倾斜状态")
                else:
                    # 线性插值过渡
                    fraction = transition_elapsed / ROLL_TRANSITION_DURATION
                    data.ctrl[hip_left] = 0 - hip_adjustment * fraction
                    data.ctrl[leg_left] = 0 - leg_adjustment * fraction
                    data.ctrl[hip_right] = 0 + hip_adjustment * fraction
                    data.ctrl[leg_right] = 0 + leg_adjustment * fraction
            
            elif current_phase == "DONE":
                # 倾斜完成阶段 - 保持当前状态
                # 计算并输出当前Roll角度
                quat = data.sensor("orientation").data.copy()
                current_roll = get_roll_from_quat(quat)
                roll_angle_deg = math.degrees(current_roll - initial_roll)
                print(f"当前Roll倾斜角度: {roll_angle_deg:.2f}度")
                
                # 停留一段时间后退出
                if time.time() - transition_start_time > ROLL_TRANSITION_DURATION + 2.0:
                    print(f"最终Roll倾斜角度: {roll_angle_deg:.2f}度 (目标: {TARGET_ROLL_DEGREE}度)")
                    break
            
            # 执行物理模拟
            mujoco.mj_step(model, data)
            
            # 同步查看器
            viewer.sync()
            
            # 控制步长（约1000Hz）
            time_until_next_step = 0.001 - (time.time() - step_start)
            if time_until_next_step > 0:
                time.sleep(time_until_next_step)

def get_roll_from_quat(quat):
    """从四元数中提取Roll角度（弧度）"""
    # 四元数转换为旋转矩阵
    rot_matrix = np.zeros(9)
    mujoco.mju_quat2Mat(rot_matrix, quat)
    
    # 从旋转矩阵中提取Roll角度（绕X轴的旋转）
    roll = math.atan2(rot_matrix[7], rot_matrix[8])
    return roll

if __name__ == "__main__":
    main()