#!/usr/bin/env python3
"""
MuJoCo API 核心用法演示 - 关节信息读取与控制
"""
import mujoco
import numpy as np

# 加载模型
XML_PATH = "/home/8x54zj-m/unitree_mujoco/data/go1/xml/go1.xml"
MESH_PATH = "/home/8x54zj-m/unitree_mujoco/data/go1/meshes"

with open(XML_PATH, 'r') as f:
    xml = f.read()
xml = xml.replace('meshdir="../meshes/"', f'meshdir="{MESH_PATH}"')
xml = xml.replace('<default>\n=', '<default>\n')

model = mujoco.MjModel.from_xml_string(xml)
data = mujoco.MjData(model)

# ============================================================
# 1. 获取模型基本信息
# ============================================================
print("=" * 60)
print("1. 模型基本信息")
print("=" * 60)
print(f"  nq  (自由度数量): {model.nq}")      # 位置自由度
print(f"  nv  (速度自由度):  {model.nv}")      # 速度自由度
print(f"  nu  (控制数量):    {model.nu}")      # actuator 数量
print(f"  njnt (关节数量):  {model.njnt}")    # 关节数量

# ============================================================
# 2. 读取关节信息
# ============================================================
print("\n" + "=" * 60)
print("2. 关节信息")
print("=" * 60)

# 关节索引 0 是 world 或 freejoint, 从 1 开始是实际关节
print("\n关节列表 (用 mj_id2name):")
for i in range(model.njnt):
    name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i)
    print(f"  [{i:2d}] {name}")

# ============================================================
# 3. 状态向量详解
# ============================================================
print("\n" + "=" * 60)
print("3. 状态向量")
print("=" * 60)

print(f"""
MuJoCo 状态向量布局 (Go1 是自由关节机器人):

qpos (位置, {model.nq}维):
  qpos[0:3]   = 全局位置 (x, y, z)
  qpos[3:7]   = 全局四元数 (qx, qy, qz, qw)
  qpos[7:19]  = 12 个关节位置 (hip, thigh, calf)
  qpos[19:]   = 额外数据

qvel (速度, {model.nv}维):
  qvel[0:3]   = 全局线速度
  qvel[3:6]   = 全局角速度
  qvel[6:18]  = 12 个关节速度
""")

print("当前状态:")
print(f"  qpos = {data.qpos}")
print(f"  qvel = {data.qvel}")

# ============================================================
# 4. 通过关节名读取特定关节数据
# ============================================================
print("\n" + "=" * 60)
print("4. 通过关节名读取数据")
print("=" * 60)

# 方法: mj_name2id 获取关节 ID, 然后用 jnt_qposadr/jnt_dofadr 找到索引
def get_joint_state(model, data, joint_name):
    """获取指定关节的位置和速度

    注意: MuJoCo 中:
    - qpos 位置索引通过 jnt_qposadr 获取
    - qvel 速度索引通过 jnt_dofadr 获取 (不是 jnt_qveladr!)
    """
    joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, joint_name)
    if joint_id < 0:
        return None, None
    qposadr = model.jnt_qposadr[joint_id]
    dofadr = model.jnt_dofadr[joint_id]  # 注意: 是 dofadr 不是 qveladr
    # 读取位置 (每个关节 1 个 qpos 值, freejoint 是 7 个)
    nq = model.jnt_type[joint_id] == mujoco.mjtJoint.mjJNT_FREE and 7 or 1
    pos = data.qpos[qposadr:qposadr+nq]
    # 读取速度 (每个关节 1 个 dof, freejoint 没有速度直接索引)
    vel = data.qvel[dofadr] if dofadr >= 0 else 0.0
    return pos, vel

# 示例: 读取前腿关节
print("\n读取前右腿关节:")
for name in ["FR_hip_joint", "FR_thigh_joint", "FR_calf_joint"]:
    pos, vel = get_joint_state(model, data, name)
    pos_str = f"{pos:.4f}" if np.isscalar(pos) else str(pos)
    print(f"  {name:20s}: pos={pos_str}, vel={vel:.4f}")

# ============================================================
# 5. 控制输出
# ============================================================
print("\n" + "=" * 60)
print("5. 控制输出 (ctrl)")
print("=" * 60)

print(f"\n控制向量长度: {model.nu}")
print(f"控制范围 (从 XML 读取):\n{model.actuator_ctrlrange}")

# 遍历所有 actuator
print("\nActuator 列表:")
for i in range(model.nu):
    name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, i)
    # 获取 actuator 关联的 joint
    trn_joint = model.actuator_trnid[i, 0]
    joint_name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, trn_joint) if trn_joint >= 0 else "none"
    ctrl_range = model.actuator_ctrlrange[i]
    print(f"  [{i:2d}] {name:15s} -> {joint_name:15s} range=[{ctrl_range[0]:.2f}, {ctrl_range[1]:.2f}]")

# ============================================================
# 6. 控制输出方法
# ============================================================
print("\n" + "=" * 60)
print("6. 控制输出方法")
print("=" * 60)

print("""
# 方法1: 直接设置所有控制
data.ctrl[:] = 0.0

# 方法2: 设置单个控制
data.ctrl[0] = 1.0

# 方法3: 用索引设置 (与 actuator 顺序对应)
data.ctrl[np.arange(12)] = target_positions

# 方法4: 先计算再设置
target_torque = kp * (desired - actual) + kd * (desired_vel - actual_vel)
data.ctrl[:] = target_torque
""")

# 示例: 设置目标关节位置
desired_joint_pos = np.zeros(12)
for i in range(12):
    desired_joint_pos[i] = np.random.uniform(-0.5, 0.5)

print(f"\n示例: 设置目标关节位置: {desired_joint_pos}")
data.ctrl[:] = desired_joint_pos
print(f"已写入 ctrl: {data.ctrl}")

# ============================================================
# 7. 仿真循环
# ============================================================
print("\n" + "=" * 60)
print("7. 仿真循环示例")
print("=" * 60)

print("""
# 基本仿真循环:
for _ in range(100):
    # 1. 读取当前状态
    current_pos = data.qpos[7:19]  # 关节位置
    current_vel = data.qvel[6:18]  # 关节速度

    # 2. 计算控制量 (这里用简单的 PD 控制作为示例)
    desired_pos = np.array([...])  # 目标位置
    error = desired_pos - current_pos
    data.ctrl[:] = kp * error - kd * current_vel

    # 3. 执行一步仿真
    mujoco.mj_step(model, data)

    # 4. (可选) 渲染
    # renderer.update_scene(data)
    # img = renderer.render()
""")

# 实际运行 50 步
print("\n运行 50 步仿真 (零控制):")
data.ctrl[:] = 0.0
for step in range(50):
    mujoco.mj_step(model, data)
    if step % 10 == 0:
        print(f"  Step {step:3d}: trunk_z={data.qpos[2]:.4f}m, "
              f"FR_hip={data.qpos[7]:.4f}, FR_thigh={data.qpos[8]:.4f}")

print("\n机器人倒下是因为零控制 + 重力!")
print("真实 RL 控制需要在每步根据策略计算正确的 ctrl 值。")

# ============================================================
# 8. 重要 API 速查
# ============================================================
print("\n" + "=" * 60)
print("8. 重要 API 速查")
print("=" * 60)
print("""
┌─────────────────────────────────────────────────────────────┐
│ 读取模型信息                                                │
├─────────────────────────────────────────────────────────────┤
│ model.nq                 # 自由度数                          │
│ model.nu                 # 控制数                            │
│ model.njnt              # 关节数                             │
│ mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i)   │
│ mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name) │
├─────────────────────────────────────────────────────────────┤
│ 读取状态                                                    │
├─────────────────────────────────────────────────────────────┤
│ data.qpos              # 位置向量                           │
│ data.qvel              # 速度向量                           │
│ data.qacc              # 加速度向量                         │
│ data.ctrl              # 控制向量 (要设置的)                │
├─────────────────────────────────────────────────────────────┤
│ 关节层面                                                    │
├─────────────────────────────────────────────────────────────┤
│ joint_id = mujoco.mj_name2id(model, ..., joint_name)       │
│ qpos_adr = model.jnt_qposadr[joint_id]                    │
│ qvel_adr = model.jnt_qveladr[joint_id]                    │
│ joint_pos = data.qpos[qpos_adr]                           │
│ joint_vel = data.qvel[qvel_adr]                           │
├─────────────────────────────────────────────────────────────┤
│ 仿真                                                        │
├─────────────────────────────────────────────────────────────┤
│ mujoco.mj_step(model, data)   # 执行一步                   │
│ mujoco.mj_stepN(model, data, N)  # 执行 N 步               │
│ mujoco.mj_resetData(model, data)  # 重置状态               │
└─────────────────────────────────────────────────────────────┘
""")