#!/usr/bin/env python3
"""最小 mesh 碰撞测试——验证 MotrixSim 的 OBJ mesh 是否支持碰撞。"""

import os, numpy as np

xml_dir = '/home/8x54zj-m/MotrixLab/motrix_envs/src/motrix_envs/locomotion/go1/xmls'
assets_dir = os.path.join(xml_dir, 'assets', 'tmp_test')
os.makedirs(assets_dir, exist_ok=True)
obj_path = os.path.join(assets_dir, 'test_box.obj')

# 封闭 box OBJ (1x1x0.1m)，带法线
with open(obj_path, 'w') as f:
    f.write("""# closed box
v -0.5 -0.5 0.0
v 0.5 -0.5 0.0
v 0.5 0.5 0.0
v -0.5 0.5 0.0
v -0.5 -0.5 0.1
v 0.5 -0.5 0.1
v 0.5 0.5 0.1
v -0.5 0.5 0.1
f 1 3 2
f 1 4 3
f 5 6 7
f 5 7 8
f 1 5 6
f 1 6 2
f 2 6 7
f 2 7 3
f 3 7 8
f 3 8 4
f 4 8 5
f 4 5 1
""")

# 生成测试 XML
xml_path = os.path.join(xml_dir, 'scene_test_mesh.xml')
with open(xml_path, 'w') as f:
    f.write("""<mujoco model="test mesh">
  <include file="go1_motor_actuator.xml" />
  <include file="materials.xml" />
  <statistic center="0 0 0.3" extent="1" />
  <visual>
    <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0" />
  </visual>
  <asset>
    <mesh name="test_box" file="assets/tmp_test/test_box.obj" />
  </asset>
  <worldbody>
    <light pos="0 0 2" dir="0 0 -1" directional="true" />
    <!-- 无 plane floor -- 纯 mesh 碰撞测试 -->
    <geom name="box_mesh" type="mesh" mesh="test_box" pos="0 0 0.15"
          contype="1" conaffinity="1" rgba="0.8 0.3 0.3 1" friction="0.8 0.3 0.3"/>
  </worldbody>
</mujoco>
""")

import motrixsim as mtx
model = mtx.load_model(xml_path)
print('加载成功')

data = mtx.SceneData(model, batch=[1])
data.reset(model)
body = model.get_body(0)

init_pos = model.compute_init_dof_pos().reshape(1, -1)
init_pos[0, 0:2] = 0.0   # 在 box 正上方
init_pos[0, 2] = 0.8      # 从 0.8m 自由落体 (box 顶面在 z=0.25)
data.set_dof_pos(init_pos, model)
model.forward_kinematic(data)

print('自由落体到 box mesh (顶部 z=0.25):')
for i in range(80):
    model.step(data)
    bz = body.get_pose(data)[0, 2]
    if i < 15 or i % 15 == 0:
        print(f'  步{i+1}: base_z={bz:.4f}')

bz_final = body.get_pose(data)[0, 2]
print(f'\n最终 base_z={bz_final:.4f}')
if bz_final > 0.45:
    print('✅ mesh 碰撞正常！机器人站在 box 上')
elif bz_final < 0.10:
    print('❌ mesh 碰撞不工作！机器人穿透 box 坠入深渊')
else:
    print(f'⚠ 不确定: base_z={bz_final:.4f}')