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

# 期望的俯仰角度（弧度）
x = 60  # 腿部岔开角度
a = 90 - (x / 2)
DESIRED_DEGREE = 10  # 输入一个角度,正值为俯，负值为仰
DESIRED_PITCH = np.radians(DESIRED_DEGREE) * (a / 90)  # 转换为弧度并调整比例

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)  # 确保所有状态完全重置
    
    # 初始姿态设置 - 确保设置所有相关关节
    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)
    
    # 创建查看器
    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
        
        # 初始化控制信号
        data.ctrl[hip_left] = 0
        data.ctrl[leg_left] = 0
        data.ctrl[hip_right] = 0
        data.ctrl[leg_right] = 0
        
        # 状态控制变量
        start_time = time.time()
        current_phase = "STABLE"  # STABLE, MOVING, DONE
        transition_start_time = 0
        transition_duration = 1.0  # 过渡时间为1秒
        current_hip_target = 0  # 当前的目标角度值
        
        # 初始稳定步骤 - 让物理引擎稳定
        for _ in range(100):
            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":
                # 稳定阶段 - 保持初始姿态1秒
                if elapsed_time >= 1.0:
                    current_phase = "MOVING"
                    transition_start_time = time.time()
                    print("开始向目标角度过渡...")
            
            elif current_phase == "MOVING":
                # 过渡阶段 - 在1秒内平滑过渡到目标角度
                transition_elapsed = time.time() - transition_start_time
                if transition_elapsed >= transition_duration:
                    current_hip_target = DESIRED_PITCH
                    current_phase = "DONE"
                    print(f"达到目标俯仰角度: {DESIRED_DEGREE}度")
                else:
                    # 线性插值过渡
                    fraction = transition_elapsed / transition_duration
                    current_hip_target = DESIRED_PITCH * fraction
            
            # 更新控制信号
            data.ctrl[hip_left] = current_hip_target
            data.ctrl[hip_right] = -current_hip_target
            data.ctrl[leg_left] = 0
            data.ctrl[leg_right] = 0
            
            # 执行物理模拟
            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)

if __name__ == "__main__":
    main()