#!/usr/bin/env python3
"""地形碰撞诊断工具 — 系统性检测 hfield PNG 地形的物理交互质量。

检测项：
  1. 脚部穿透（foot z < terrain z）
  2. 接触力异常（为 0 但应该接触 / 过大）
  3. NaN/Inf 传播
  4. 物理爆炸（关节速度飙升）
  5. 机器人卡住（速度接近 0 但有命令）
  6. 躯干沉入地形（base z < terrain z）

用法:
    uv run scripts/diag_terrain_collision.py                     # 全量检测
    uv run scripts/diag_terrain_collision.py --level 5           # 仅 level 5
    uv run scripts/diag_terrain_collision.py --type stairs       # 仅楼梯地形
    uv run scripts/diag_terrain_collision.py --render            # 渲染可视化
"""
import argparse, os, sys, time
import numpy as np

os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
os.environ.setdefault("JAX_PLATFORMS", "cpu")
os.environ.setdefault("DREAMWAQ_TERRAIN", "pyramid")

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

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

# ═══════════════════════════════════════════════════════════════════════
# 配置
# ═══════════════════════════════════════════════════════════════════════

NUM_ROWS = 10
NUM_COLS = 20
CELL_M = 8.0
BORDER_M = 5.0
SETTLE_STEPS = 100     # 自由落体稳定步数（200Hz）
TEST_STEPS = 400       # 测试步数
STEPS_PER_CELL = 5     # 每个 cell 测试的 spawn 位置数

# 地形类型名（与 gen_dreamwaq_terrain.py PROPORTIONS 对应）
TYPE_NAMES = ["平滑斜坡", "粗糙斜坡", "下行楼梯", "上行楼梯", "离散障碍"]
CUM = [0.1, 0.2, 0.55, 0.9, 1.0]


def _cell_origin(row, col):
    half_x = BORDER_M + NUM_COLS * CELL_M / 2.0
    half_y = BORDER_M + NUM_ROWS * CELL_M / 2.0
    cx = -half_x + BORDER_M + col * CELL_M + CELL_M / 2
    cy = half_y - BORDER_M - row * CELL_M - CELL_M / 2
    return cx, cy


def _terrain_type(col):
    choice = col / NUM_COLS + 0.001
    for i, c in enumerate(CUM):
        if choice < c:
            return i
    return len(TYPE_NAMES) - 1


# ═══════════════════════════════════════════════════════════════════════
# 诊断核心
# ═══════════════════════════════════════════════════════════════════════

class CollisionDiag:
    def __init__(self, env, render=False):
        self.env = env
        self.render = render
        self.renderer = None
        if render:
            from motrix_envs.np.renderer import NpRenderer
            self.renderer = NpRenderer(env)

        # 统计数据
        self.stats = {
            "penetrations": [],       # (row, col, ttype, foot_z - terrain_z)
            "contact_anomalies": [],  # (row, col, ttype, max_contact_force)
            "nan_events": [],         # (row, col, ttype, step)
            "explosions": [],         # (row, col, ttype, max_dof_vel)
            "stuck_events": [],       # (row, col, ttype, speed)
            "base_sink": [],          # (row, col, ttype, base_z - terrain_z)
            "ok_count": 0,
            "total_tested": 0,
        }

    def _get_terrain_z(self, x, y):
        """查询 (x, y) 处的地形高度。"""
        hm = self.env._hm_cache
        if hm is None:
            return 0.0
        nr, nc = hm.shape
        xmin = self.env._h_xmin
        ymax = self.env._h_ymax
        w = self.env._h_xmax - xmin
        h = ymax - self.env._h_ymin
        z_base = self.env._h_zbase
        col = int(np.clip((x - xmin) / max(w, 1e-6) * (nc - 1), 0, nc - 1))
        row = int(np.clip((ymax - y) / max(h, 1e-6) * (nr - 1), 0, nr - 1))
        return float(hm[row, col]) + z_base

    def _get_foot_positions(self, state):
        """获取所有足部在世界坐标系中的位置 (N, 4, 3)。"""
        data = state.data
        n = data.shape[0]
        feet = np.zeros((n, 4, 3), dtype=np.float32)
        foot_names = ["FR_foot", "FL_foot", "RR_foot", "RL_foot"]
        for i, name in enumerate(foot_names):
            gidx = self.env._model.get_geom_index(name)
            if gidx is not None:
                pose = self.env._model.geoms[gidx].get_pose(data)
                feet[:, i, :] = pose[:, :3]
        return feet

    def test_cell(self, row, col, level):
        """在单个 cell 上测试多个 spawn 位置。"""
        cx, cy = _cell_origin(row, col)
        ttype = _terrain_type(col)
        tname = TYPE_NAMES[ttype]

        # 在 cell 内取几个 spawn 点（中央 + 四角附近）
        offsets = [
            (0, 0),
            (CELL_M * 0.3, 0), (-CELL_M * 0.3, 0),
            (0, CELL_M * 0.3), (0, -CELL_M * 0.3),
        ][:STEPS_PER_CELL]

        for ox, oy in offsets:
            sx, sy = cx + ox, cy + oy
            terrain_z = float(self.env._sample_terrain_height(
                np.array([[sx, sy]], dtype=np.float32), radius=0.35)[0])

            self.stats["total_tested"] += 1
            self._test_single(sx, sy, terrain_z, row, col, level, ttype, tname)

    def _test_single(self, sx, sy, terrain_z, row, col, level, ttype, tname):
        """单个 spawn 点的完整测试。"""
        env = self.env
        data = env._state.data
        n = 1

        # Spawn
        init_pos = env._init_dof_pos.copy().reshape(1, -1)
        init_pos[0, 0] = sx
        init_pos[0, 1] = sy
        spawn_z = terrain_z + 0.45  # 标准 clearance
        init_pos[0, 2] = spawn_z
        yaw = np.random.uniform(-np.pi, np.pi)
        init_pos[0, 3:7] = [0, 0, np.sin(yaw / 2), np.cos(yaw / 2)]

        data.reset(env._model)
        data.set_dof_pos(init_pos, env._model)
        env._model.forward_kinematic(data)

        # 给一个随机动作防止完全静止（小幅度 PD 控制）
        action = 0.1 * (np.random.randn(n, 12).astype(np.float32))

        has_nan = False
        max_dof_vel = 0.0
        max_contact_force = 0.0
        min_foot_z = spawn_z
        min_base_z = spawn_z
        total_speed = 0.0

        for step in range(SETTLE_STEPS + TEST_STEPS):
            # Apply PD
            actions_scaled = action * env.cfg.control_config.action_scale
            target = actions_scaled + env.default_angles.reshape(1, -1)
            lo = env._model.joint_limits[0].reshape(1, -1)
            hi = env._model.joint_limits[1].reshape(1, -1)
            target = np.clip(target, lo, hi)
            torques = env.kps * (target - env.get_dof_pos(data)) - env.kds * env.get_dof_vel(data)
            data.actuator_ctrls = np.clip(torques, -80.0, 80.0)

            env._model.step(data)

            # ── 检测 ──

            # 1. NaN 检测
            dv = np.array(data.dof_vel)
            dp = np.array(data.dof_pos)
            if np.any(~np.isfinite(dv)) or np.any(~np.isfinite(dp)):
                self.stats["nan_events"].append((row, col, ttype, step))
                has_nan = True
                break

            # 2. 速度爆炸
            cur_max_vel = float(np.max(np.abs(dv[:, 6:])))
            max_dof_vel = max(max_dof_vel, cur_max_vel)
            if cur_max_vel > 100.0:
                self.stats["explosions"].append((row, col, ttype, cur_max_vel))
                break

            # 3. 接触力
            # Read contact forces via MotrixSim sensor API
            cf_vals = []
            for foot in ["FR", "FL", "RR", "RL"]:
                try:
                    v = env._model.get_sensor_value(f"{foot}_foot_contact", data)
                    cf_vals.append(v[0])
                except Exception:
                    cf_vals.append(np.zeros(3, dtype=np.float32))
            cf_raw = np.concatenate(cf_vals)
            cur_cf = float(np.max(np.abs(cf_raw)))
            max_contact_force = max(max_contact_force, cur_cf)

            # 4. 足部穿透（settle 之后测）
            if step >= SETTLE_STEPS:
                feet = self._get_foot_positions(env._state)
                for fi in range(4):
                    fz = float(feet[0, fi, 2])
                    tz = self._get_terrain_z(float(feet[0, fi, 0]), float(feet[0, fi, 1]))
                    penetration = tz - fz  # 正 = 穿透
                    min_foot_z = min(min_foot_z, fz - tz)  # 相对地形的高度
                    if penetration > 0.02:  # > 2cm 穿透
                        self.stats["penetrations"].append(
                            (row, col, ttype, float(penetration)))

            # 5. 躯干沉入
            base_pose = env._body.get_pose(data)
            bz = float(base_pose[0, 2])
            tz = self._get_terrain_z(float(base_pose[0, 0]), float(base_pose[0, 1]))
            min_base_z = min(min_base_z, bz - tz)
            if bz < tz + 0.10 and step >= SETTLE_STEPS:  # 躯干离地面 < 10cm
                self.stats["base_sink"].append((row, col, ttype, float(bz - tz)))

            # 6. 速度（卡住检测）
            if step >= SETTLE_STEPS:
                linvel = env.get_local_linvel(data)[0]
                total_speed += float(np.linalg.norm(linvel[:2]))

            if self.renderer and step % 20 == 0:
                self.renderer.render()
                time.sleep(0.005)

        # ── 统计 ──
        if not has_nan and max_dof_vel < 100.0:
            avg_speed = total_speed / max(TEST_STEPS, 1)
            if avg_speed < 0.01 and step == SETTLE_STEPS + TEST_STEPS - 1:
                self.stats["stuck_events"].append((row, col, ttype, avg_speed))

        if not has_nan and max_dof_vel < 50.0 and min_base_z > 0.05:
            self.stats["ok_count"] += 1

        # 记录接触力异常（完全无接触力但足部应该着地）
        if max_contact_force < 1e-3 and step == SETTLE_STEPS + TEST_STEPS - 1:
            feet = self._get_foot_positions(env._state)
            for fi in range(4):
                fz = float(feet[0, fi, 2])
                tz = self._get_terrain_z(float(feet[0, fi, 0]), float(feet[0, fi, 1]))
                if abs(fz - tz) < 0.05:  # 足部接近地面
                    self.stats["contact_anomalies"].append(
                        (row, col, ttype, 0.0))  # 0 = 无接触力
                    break

    def run(self, target_level=None, target_type=None):
        """遍历所有地形 cell 测试。"""
        total = 0
        start_t = time.time()

        for row in range(NUM_ROWS):
            if target_level is not None and row != target_level:
                continue
            for col in range(NUM_COLS):
                ttype = _terrain_type(col)
                if target_type is not None and ttype != target_type:
                    continue
                total += 1

        print(f"[Diag] 测试 {NUM_ROWS}×{NUM_COLS} 网格, "
              f"每 cell {STEPS_PER_CELL} 个 spawn 点 = {total * STEPS_PER_CELL} 次")
        print(f"[Diag] Settle={SETTLE_STEPS}步 Test={TEST_STEPS}步")

        count = 0
        for row in range(NUM_ROWS):
            if target_level is not None and row != target_level:
                continue
            for col in range(NUM_COLS):
                ttype = _terrain_type(col)
                if target_type is not None and ttype != target_type:
                    continue
                self.test_cell(row, col, row)
                count += 1
                if count % 10 == 0:
                    elapsed = time.time() - start_t
                    print(f"  [{count}/{total}] {elapsed:.1f}s  "
                          f"OK={self.stats['ok_count']}/{self.stats['total_tested']}  "
                          f"穿透={len(self.stats['penetrations'])}  "
                          f"NaN={len(self.stats['nan_events'])}  "
                          f"爆炸={len(self.stats['explosions'])}")

        elapsed = time.time() - start_t
        self._report(elapsed)

    def _report(self, elapsed):
        s = self.stats
        t = s["total_tested"]
        print("\n" + "=" * 70)
        print(f"地形碰撞诊断报告 ({elapsed:.1f}s, {t} 次测试)")
        print("=" * 70)

        # 总体统计
        ok_rate = s["ok_count"] / max(t, 1) * 100
        print(f"\n  ✅ 通过: {s['ok_count']}/{t} ({ok_rate:.1f}%)")

        # 穿透
        if s["penetrations"]:
            pens = s["penetrations"]
            pvals = [p[3] for p in pens]
            print(f"\n  🔴 足部穿透: {len(pens)} 次")
            print(f"     最大穿透: {max(pvals):.3f}m  平均: {np.mean(pvals):.3f}m")

            # 按地形类型分组
            by_type = {}
            for p in pens:
                tn = TYPE_NAMES[p[2]]
                by_type.setdefault(tn, []).append(p[3])
            for tn, vals in sorted(by_type.items()):
                print(f"     {tn}: {len(vals)}次  max={max(vals):.3f}m  avg={np.mean(vals):.3f}m")

            # 按 level 分组
            by_level = {}
            for p in pens:
                by_level.setdefault(p[0], []).append(p[3])
            print(f"     按 level: ", end="")
            for lv in sorted(by_level.keys()):
                vals = by_level[lv]
                print(f"L{lv}({len(vals)}x max={max(vals):.3f}) ", end="")
            print()
        else:
            print(f"\n  🟢 足部穿透: 0 次 ✓")

        # NaN
        if s["nan_events"]:
            print(f"\n  🔴 NaN/Inf: {len(s['nan_events'])} 次")
            for n in s["nan_events"][:5]:
                print(f"     L{n[0]} C{n[1]:02d} {TYPE_NAMES[n[2]]} @step {n[3]}")
        else:
            print(f"  🟢 NaN/Inf: 0 次 ✓")

        # 物理爆炸
        if s["explosions"]:
            print(f"\n  🟡 速度爆炸 (>{100}rad/s): {len(s['explosions'])} 次")
            for e in s["explosions"][:5]:
                print(f"     L{e[0]} C{e[1]:02d} {TYPE_NAMES[e[2]]} max_vel={e[3]:.1f}")
        else:
            print(f"  🟢 速度爆炸: 0 次 ✓")

        # 接触力异常
        if s["contact_anomalies"]:
            print(f"\n  🟡 接触力异常 (无接触但足部近地): {len(s['contact_anomalies'])} 次")
        else:
            print(f"  🟢 接触力异常: 0 次 ✓")

        # 卡住
        if s["stuck_events"]:
            print(f"\n  🟡 卡住 (速度<0.01m/s): {len(s['stuck_events'])} 次")
            by_type = {}
            for e in s["stuck_events"]:
                tn = TYPE_NAMES[e[2]]
                by_type.setdefault(tn, []).append(1)
            for tn, vals in sorted(by_type.items()):
                print(f"     {tn}: {len(vals)}次")
        else:
            print(f"  🟢 卡住: 0 次 ✓")

        # 躯干沉入
        if s["base_sink"]:
            sinks = s["base_sink"]
            print(f"\n  🟡 躯干贴地 (<10cm): {len(sinks)} 次")
            vals = [s[3] for s in sinks]
            print(f"     最差: {min(vals):.3f}m  平均: {np.mean(vals):.3f}m")
        else:
            print(f"  🟢 躯干贴地: 0 次 ✓")

        # 按地形类型汇总
        print(f"\n── 按地形类型汇总 ──")
        for ti, tn in enumerate(TYPE_NAMES):
            pens_t = sum(1 for p in s["penetrations"] if p[2] == ti)
            nans_t = sum(1 for n in s["nan_events"] if n[2] == ti)
            exps_t = sum(1 for e in s["explosions"] if e[2] == ti)
            icon = "🔴" if (pens_t + nans_t + exps_t) > 0 else "🟢"
            print(f"  {icon} {tn}: 穿透{pens_t} NaN{nans_t} 爆炸{exps_t}")

        # 结论
        print(f"\n── 全局结论 ──")
        if ok_rate > 95 and len(s["nan_events"]) == 0 and len(s["penetrations"]) < 5:
            print(f"  ✅ 仿真环境正常 — hfield 地形碰撞基本可靠")
        elif ok_rate > 80:
            print(f"  ⚠ 仿真环境存在一些问题 — 部分地形类型/level 有问题")
        else:
            print(f"  🔴 仿真环境问题严重 — 需要排查 MotrixSim hfield 实现")

        if len(s["penetrations"]) > 10:
            print(f"  💡 穿透较多 → hfield 碰撞面可能偏软或 contact margin 过大")
            print(f"     可尝试增大 geom margin 或检查 MotrixSim 碰撞参数")


# ═══════════════════════════════════════════════════════════════════════
# Main
# ═══════════════════════════════════════════════════════════════════════

def main():
    p = argparse.ArgumentParser(description="DreamWaQ 地形碰撞诊断")
    p.add_argument("--level", type=int, default=None, help="仅测试指定 level (0-9)")
    p.add_argument("--type", dest="ttype", default=None,
                   choices=["slope", "rough", "stairs_down", "stairs_up", "obstacles"],
                   help="仅测试指定地形类型")
    p.add_argument("--render", action="store_true", help="渲染可视化")
    p.add_argument("--settle", type=int, default=SETTLE_STEPS)
    p.add_argument("--test", type=int, default=TEST_STEPS)
    args = p.parse_args()

    type_map = {"slope": 0, "rough": 1, "stairs_down": 2, "stairs_up": 3, "obstacles": 4}
    target_type = type_map.get(args.ttype)

    print("[Diag] 创建 DreamWaQ 环境...")
    env = env_registry.make("go1-dreamwaq-walk", num_envs=1)

    # 如果指定 level，强制执行
    if args.level is not None:
        env._force_level = args.level
        print(f"[Diag] 强制 level={args.level}")

    env.init_state()

    # 确保 terrain 初始化（触发 hfield cache）
    if env._model.num_hfields > 0:
        hf = env._model.get_hfield(0)
        env._hm_cache = hf.height_matrix
        env._h_xmin, env._h_ymin = hf.bound[0], hf.bound[1]
        env._h_xmax, env._h_ymax = hf.bound[3], hf.bound[4]
        env._h_nr, env._h_nc = hf.height_matrix.shape
        cfg_base = getattr(env.cfg, "hfield_z_base", None)
        env._h_zbase = float(cfg_base) if cfg_base is not None else float(-hf.height_matrix[0, 0])
        print(f"[Diag] HField: {hf.height_matrix.shape}  bound={hf.bound}")
        print(f"[Diag] z_base={env._h_zbase:.3f}")
    else:
        print("[Diag] WARNING: 无 hfield（flat 场景），跳过诊断")
        return

    diag = CollisionDiag(env, render=args.render)
    diag.SETTLE_STEPS = args.settle
    diag.TEST_STEPS = args.test

    try:
        diag.run(target_level=args.level, target_type=target_type)
    except KeyboardInterrupt:
        print("\n[Diag] 中断")
    finally:
        if diag.renderer:
            diag.renderer.close()
        print("[Diag] 完成")


if __name__ == "__main__":
    main()
