#!/usr/bin/env python3
"""DreamWaQ 地形可视化。

键盘：
  R=重置  H=高度采样点  T=遍历出生点调试  Esc=退出

用法:
    uv run scripts/view_dreamwaq.py                    # 金字塔地形
    uv run scripts/view_dreamwaq.py --flat --num-envs 1
"""
import argparse, os, sys, time
import numpy as np

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

if "--flat" in sys.argv:
    os.environ["DREAMWAQ_TERRAIN"] = "flat"
elif "--flat-stairs" in sys.argv:
    os.environ["DREAMWAQ_TERRAIN"] = "flat_stairs"
elif "--stairs" in sys.argv:
    os.environ["DREAMWAQ_TERRAIN"] = "stairs"
else:
    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
from motrix_envs.np.renderer import NpRenderer
from motrix_envs.math import quaternion
from motrixsim.render import RenderClosedError

# 地形类型名称（与 gen_dreamwaq_terrain.py 的 PROPORTIONS 对应）
PROPORTIONS = [0.1, 0.1, 0.35, 0.35, 0.1]
CUM = [sum(PROPORTIONS[:i + 1]) for i in range(len(PROPORTIONS))]
TYPE_NAMES = ["平滑斜坡", "粗糙斜坡", "下行楼梯", "上行楼梯", "离散障碍"]


def _cell_origin(row, col, border_m=5.0, cell_m=8.0, num_rows=10, num_cols=20):
    """计算 cell (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 _type_name(col):
    """根据列索引返回地形类型名称。"""
    choice = col / 20 + 0.001
    for i, cum in enumerate(CUM):
        if choice < cum:
            return TYPE_NAMES[i]
    return TYPE_NAMES[-1]


def _spawn_at(env, cx, cy, spawn_z=None):
    """在指定世界坐标 spawn 单个机器人。"""
    state = env._state
    data = state.data
    init_pos = env._init_dof_pos.copy().reshape(1, -1)
    init_pos[0, 0] = cx
    init_pos[0, 1] = cy
    # 计算地形高度
    terrain_z = float(env._sample_terrain_height(
        np.array([[cx, cy]], dtype=np.float32), radius=0.35)[0])
    if spawn_z is None:
        spawn_z = terrain_z + 0.45  # 默认 clearance
    init_pos[0, 2] = spawn_z
    data.reset(env._model)
    data.set_dof_pos(init_pos, env._model)
    env._model.forward_kinematic(data)
    # 重置 info
    state.info["commands"][0] = np.array([0.0, 0.0, 0.0], dtype=np.float32)
    state.info["steps"][0] = 0
    state.info["obs_history"][0] = 0.0
    return terrain_z, spawn_z


def main():
    p = argparse.ArgumentParser(description="DreamWaQ 地形可视化")
    p.add_argument("--num-envs", type=int, default=1)
    p.add_argument("--flat", action="store_true")
    p.add_argument("--flat-stairs", action="store_true")
    p.add_argument("--stairs", action="store_true")
    p.add_argument("--level", type=int, default=None)
    p.add_argument("--no-stand", action="store_true")
    p.add_argument("--vx", type=float, default=0.5)
    args = p.parse_args()

    env = env_registry.make("go1-dreamwaq-walk", num_envs=max(args.num_envs, 1))
    if args.level is not None:
        env._force_level = args.level

    env.init_state()
    n = env._num_envs
    cmd = np.array([args.vx, 0.0, 0.0], dtype=np.float32)

    try:
        renderer = NpRenderer(env)
    except Exception as e:
        print(f"[ERROR] 渲染器创建失败: {e}")
        renderer = None

    terrain_name = os.environ.get("DREAMWAQ_TERRAIN", "pyramid")
    print(f"[View] {n} 机器人 | 地形={terrain_name}")
    print(f"[View] R=重置 H=高度点 T=遍历出生点 Esc=退出")

    show_heights = False
    traverse_mode = False
    traverse_row = 0
    traverse_col = 0
    traverse_pending = False  # 刚切换 cell, 等待稳定
    traverse_settle = 0
    step_count = 0

    # 地形信息
    num_rows = env._num_rows
    num_cols = env._num_cols
    cell_m = env._cell_size

    def enter_traverse():
        nonlocal traverse_mode, traverse_row, traverse_col, traverse_pending, traverse_settle
        traverse_mode = True
        traverse_row = 0
        traverse_col = 0
        traverse_pending = True
        traverse_settle = 0
        print(f"\n[T] 遍历模式: {num_rows}行 × {num_cols}列")
        print(f"[T] 按 T 前进, R 退出遍历\n")

    def exit_traverse():
        nonlocal traverse_mode, traverse_pending
        traverse_mode = False
        traverse_pending = False
        env.init_state()
        print("[T] 退出遍历模式\n")

    def advance_traverse():
        nonlocal traverse_row, traverse_col, traverse_pending, traverse_settle
        traverse_col += 1
        if traverse_col >= num_cols:
            traverse_col = 0
            traverse_row += 1
        if traverse_row >= num_rows:
            print("[T] 遍历完成! 按 R 退出")
            traverse_row = num_rows - 1
            traverse_col = num_cols - 1
            return False
        traverse_pending = True
        traverse_settle = 0
        return True

    def do_traverse_spawn():
        """在当前位置 spawn 并打印信息。"""
        nonlocal traverse_settle
        cx, cy = _cell_origin(traverse_row, traverse_col)
        tname = _type_name(traverse_col)
        terrain_z, spawn_z = _spawn_at(env, cx, cy)

        # 标记 spawn 点（绿色球）
        if renderer is not None:
            g = renderer._render.gizmos
            g.draw_sphere(0.15, (np.float32(cx), np.float32(cy),
                                 np.float32(spawn_z)))

        # 让机器人稳定几步
        for _ in range(30):
            env.step(np.zeros((1, 12), dtype=np.float32))
            if renderer is not None:
                renderer.render()
                time.sleep(0.005)

        base_z = env._body.get_pose(env._state.data)[0, 2]
        contacts = env._state.info.get("contacts", np.zeros(4))
        cf = env._state.info.get("privileged_obs",
                                  np.zeros((1, 247)))[0, 45:57]
        total_cf = np.sum(np.abs(cf))

        print(f"  r{traverse_row}c{traverse_col:02d}  {tname:6s}  "
              f"origin=({cx:+.0f},{cy:+.0f})  "
              f"terrain_z={terrain_z:.3f}  spawn_z={spawn_z:.3f}  "
              f"base_z={base_z:.3f}  cf={total_cf:.0f}N  "
              f"feet={contacts.astype(int).tolist()}")

        traverse_settle = 30

    try:
        while True:
            # ── 键盘 ──
            if renderer is not None:
                try:
                    inp = renderer._render.input
                    if inp.is_key_just_pressed("r"):
                        if traverse_mode:
                            exit_traverse()
                        else:
                            env.init_state()
                            step_count = 0
                            print("[R] 重置")
                    if inp.is_key_just_pressed("h"):
                        show_heights = not show_heights
                        print(f"[H] 高度点: {'开' if show_heights else '关'}")
                    if inp.is_key_just_pressed("t"):
                        if traverse_mode:
                            advance_traverse()
                        else:
                            enter_traverse()
                except Exception:
                    pass

            # ── 遍历模式 ──
            if traverse_mode and traverse_pending:
                do_traverse_spawn()
                traverse_pending = False

            # ── 正常模式动作 ──
            if not traverse_mode or traverse_settle > 0:
                if args.no_stand:
                    env._state.info["commands"][:] = cmd
                else:
                    act = 0.3 * (np.random.rand(n, 12).astype(np.float32) - 0.5)
                    env.step(act)
                if traverse_settle > 0:
                    traverse_settle -= 1

                if step_count % 100 == 0 and step_count > 0:
                    bz = env._body.get_pose(env._state.data)[0, 2]
                    print(f"[{step_count}] base_z={bz:.3f}")

            # ── 高度点 ──
            if show_heights and renderer is not None:
                state = env._state
                pose = env._body.get_pose(state.data)
                bp = pose[0, :3]
                yaw = quaternion.get_yaw(pose[0:1, 3:7])[0]
                cos_y, sin_y = np.cos(yaw), np.sin(yaw)
                for gy in env._hy:
                    for gx in env._hx:
                        wx = bp[0] + cos_y * gx - sin_y * gy
                        wy = bp[1] + sin_y * gx + cos_y * gy
                        try:
                            wz = float(env._sample_terrain_height(
                                np.array([[wx, wy]]))[0])
                            g = renderer._render.gizmos
                            g.draw_sphere(0.02, (np.float32(wx), np.float32(wy),
                                                 np.float32(wz)))
                        except Exception:
                            pass

            if renderer is not None:
                renderer.render()
                time.sleep(0.01)

            step_count += 1

    except (KeyboardInterrupt, RenderClosedError):
        pass
    try:
        if renderer is not None:
            renderer.close()
    except Exception:
        pass
    print("[View] 结束")


if __name__ == "__main__":
    main()
