#!/usr/bin/env python3
"""测试 MotrixSim hfield 的 z_scale 上限——自己动手测，不信文档。"""
import os, sys, numpy as np, time

os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
os.environ.setdefault("JAX_PLATFORMS", "cpu")
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import motrix_envs.locomotion.go1.dreamwaq  # noqa
from motrix_envs import registry as env_registry

NUM_ENVS = 256
TEST_STEPS = 100
os.environ["DREAMWAQ_TERRAIN"] = "flat"  # 用 flat 场景，手动覆盖 hfield 参数


def _build_custom_hfield(x_radius, y_radius, z_scale, z_base=0.001):
    """构建自定义 hfield 描述字符串，用于覆盖 XML 中的 hfield 参数。"""
    import tempfile, cv2
    # 生成一个纯斜坡的 hfield PNG 用于测试
    nx, ny = 200, 200
    canvas = np.zeros((ny, nx), dtype=np.uint16)
    # 从左上到右下的斜坡：高度从 0 到 z_scale
    for i in range(ny):
        for j in range(nx):
            # 对角线斜坡，最高点在右下角
            h = int((i + j) / (nx + ny) * 65535)
            canvas[i, j] = h
    # 中间 1/3 区域做平台（平坦）
    cx, cy = nx // 2, ny // 2
    p = nx // 6
    canvas[cy - p:cy + p, cx - p:cx + p] = 0

    out_d = os.path.join(os.path.dirname(__file__), "..",
                         "motrix_envs", "src", "motrix_envs",
                         "locomotion", "go1", "xmls", "assets")
    os.makedirs(out_d, exist_ok=True)
    png_path = os.path.join(out_d, "zscale_test.png")
    cv2.imwrite(png_path, canvas)

    return png_path, (x_radius, y_radius, z_scale, z_base)


def test_z_scale(z_scale, n_envs=NUM_ENVS, n_steps=TEST_STEPS):
    """测试给定 z_scale 下的物理稳定性。"""
    import tempfile, cv2

    # 生成测试 hfield
    out_d = os.path.join(os.path.dirname(__file__), "..",
                         "motrix_envs", "src", "motrix_envs",
                         "locomotion", "go1", "xmls", "assets")
    os.makedirs(out_d, exist_ok=True)

    # 简单斜坡地形
    nx, ny = 40, 40  # 小尺寸快速生成
    canvas = np.zeros((ny, nx), dtype=np.uint16)
    for i in range(ny):
        for j in range(nx):
            h = int((i / ny) * 65535)  # y 方向斜坡
            canvas[i, j] = h
    # 中央平台
    p = nx // 6
    canvas[ny // 2 - p:ny // 2 + p, nx // 2 - p:nx // 2 + p] = 0

    png_path = os.path.join(out_d, "zscale_test.png")
    cv2.imwrite(png_path, canvas)

    total_x = nx * 0.1  # HS=0.1, 粗略
    total_y = ny * 0.1

    # 构建 scene XML
    xml_path = os.path.join(os.path.dirname(__file__), "..",
                            "motrix_envs", "src", "motrix_envs",
                            "locomotion", "go1", "xmls", "scene_zscale_test.xml")
    xml = f"""<mujoco model="zscale test">
  <include file="go1_motor_actuator.xml" />
  <include file="materials.xml" />
  <statistic center="0 0 0.3" extent="2" />
  <visual>
    <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0" />
    <rgba haze="0.15 0.25 0.35 1" />
    <global azimuth="120" elevation="-20" />
  </visual>
  <asset>
    <hfield name="test_hf" file="assets/zscale_test.png"
            size="{total_x/2:.1f} {total_y/2:.1f} {z_scale:.3f} 0.001" />
  </asset>
  <worldbody>
    <light pos="0 0 2" dir="0 0 -1" directional="true" />
    <geom name="floor" pos="0 0 0" type="hfield" hfield="test_hf"
          material="motphys-ground" contype="1" conaffinity="0" priority="1" friction="0.6" />
  </worldbody>
  <sensor>
    <contact name="FR_foot_contact" geom2="FR_foot" geom1="floor" data="force" num="1" />
    <contact name="FL_foot_contact" geom2="FL_foot" geom1="floor" data="force" num="1" />
    <contact name="RR_foot_contact" geom2="RR_foot" geom1="floor" data="force" num="1" />
    <contact name="RL_foot_contact" geom2="RL_foot" geom1="floor" data="force" num="1" />
  </sensor>
</mujoco>"""
    with open(xml_path, "w") as f:
        f.write(xml)

    import motrixsim as mtx
    t0 = time.time()

    try:
        model = mtx.load_model(xml_path)
        data = mtx.SceneData(model, batch=[n_envs])
        data.reset(model)
        body = model.get_body(0)

        # 随机初始位置（在平台上）
        init_pos = model.compute_init_dof_pos()
        init_pos = np.tile(init_pos, (n_envs, 1))
        init_pos[:, 0] += np.random.uniform(-0.5, 0.5, n_envs)
        init_pos[:, 1] += np.random.uniform(-0.5, 0.5, n_envs)
        init_pos[:, 2] = 0.5  # 从 0.5m 掉落
        data.set_dof_pos(init_pos.astype(np.float32), model)
        model.forward_kinematic(data)

        heights = np.zeros((n_steps, n_envs), dtype=np.float32)
        fall_count = 0
        for step in range(n_steps):
            model.step(data)
            h = body.get_pose(data)[:, 2]
            heights[step] = h
            # 摔倒检测：base_z < 0.15 (趴了)
            fall_count += np.sum(h < 0.15)

        mean_h = float(np.mean(heights[-20:]))  # 最后 20 步平均
        total_falls = fall_count
        dt = time.time() - t0
        return mean_h, total_falls, dt
    except Exception as e:
        return None, str(e), 0


if __name__ == "__main__":
    print(f"{'z_scale':>8s}  {'mean_base_z':>12s}  {'falls':>8s}  {'time':>8s}  verdict")
    print("-" * 65)

    for zs in [0.3, 0.5, 0.54, 0.8, 1.0, 1.5, 2.0, 3.0]:
        mean_h, falls, dt = test_z_scale(zs, n_envs=64, n_steps=50)
        if mean_h is None:
            print(f"{zs:8.3f}  {'ERROR':>12s}  {str(falls)[:20]:>8s}")
            continue
        ok = "✅ 稳定" if mean_h > 0.25 and falls < 10 else "⚠ 不稳" if mean_h > 0.15 else "❌ 崩溃"
        print(f"{zs:8.3f}  {mean_h:12.4f}  {falls:8d}  {dt:7.1f}s  {ok}")