#!/usr/bin/env python3
"""
    MOTRIXLAB_UNTRIEE_GO1_SIM2SIM
    source /opt/mujoco/venv/bin/activate
    cd /opt/unitree_mujoco
    python demo/go1_sim2sim_mujoco.py

"""

import numpy as np
import mujoco
from mujoco import viewer
import os
import threading
import signal
import queue
import argparse
import time

g_exit_requested = False

def signal_handler(signum, frame):
    global g_exit_requested
    g_exit_requested = True

signal.signal(signal.SIGINT, signal_handler)

# ============================================================
# 配置
# ============================================================
_PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_ONNX_PATH = os.path.join(_PROJECT_DIR, "exports_go1_flat", "policy.onnx")
MOTRIX_XML_DIR = os.path.join(_PROJECT_DIR, "motrix_envs", "src", "motrix_envs", "locomotion", "go1", "xmls")
XML_PATH = f"{MOTRIX_XML_DIR}/go1_motor_actuator.xml"

TERRAIN = "none"

# ============================================================
# MotrixLab 参数 (来自 cfg.py)
# ============================================================
NUM_OBS = 45   # 去掉线速度观测 (原来是48)
NUM_ACTIONS = 12
OBS_SCALES = {'lin_vel': 2.0, 'ang_vel': 0.25, 'dof_pos': 1.0, 'dof_vel': 0.05}
ACTION_SCALE = 0.05
KP = 80.0
KD = 1.0
CLIP_ACTIONS = 23.7
CLIP_OBSERVATIONS = 100.0
MAX_LIN_VEL_X = 1.0
MAX_LIN_VEL_Y = 1.0
MAX_ANG_VEL = 1.0   # 匹配训练时的角速度命令范围 [-1.0, 1.0]

# ============================================================
# 关节名称和顺序
# ============================================================
POLICY_JOINT_NAMES = [
    "FR_hip", "FR_thigh", "FR_calf",
    "FL_hip", "FL_thigh", "FL_calf",
    "RR_hip", "RR_thigh", "RR_calf",
    "RL_hip", "RL_thigh", "RL_calf",
]

DEFAULT_JOINT_ANGLES = np.array([
    -0.0,  0.9, -1.8,   # FR_hip, FR_thigh, FR_calf
     0.0,  0.9, -1.8,   # FL_hip, FL_thigh, FL_calf
    -0.0,  0.9, -1.8,   # RR_hip, RR_thigh, RR_calf
     0.0,  0.9, -1.8,   # RL_hip, RL_thigh, RL_calf
], dtype=np.float32)

MUJOCO_TO_POLICY = np.arange(12, dtype=np.int64)
POLICY_TO_MUJOCO = np.arange(12, dtype=np.int64)

# ============================================================
# 键盘输入
# ============================================================
from pynput import keyboard

class KeyboardReader:
    def __init__(self):
        self._event_queue = queue.Queue()
        self.running = True
        self.shared_keys_held = set()
        self.shared_one_shot = set()
        self._reader_thread = None
        self._listener = None

    def _normalize_key(self, key):
        try:
            if hasattr(key, 'char') and key.char is not None:
                return key.char.lower()
        except:
            pass
        key_str = str(key)
        if key_str == 'Key.esc':
            return 'escape'
        elif key_str == 'Key.space':
            return 'space'
        elif key_str.startswith('Key.'):
            return key_str.lower()
        return key_str.lower()

    def _reader_worker(self):
        while self.running:
            try:
                event_type, key = self._event_queue.get(timeout=0.05)
                k = self._normalize_key(key)
                if event_type == 'press':
                    self.shared_keys_held.add(k)
                    self.shared_one_shot.discard(k)
                elif event_type == 'release':
                    self.shared_keys_held.discard(k)
                    self.shared_one_shot.discard(k)
            except queue.Empty:
                pass

    def init(self):
        def on_press(key):
            self._event_queue.put(('press', key))
        def on_release(key):
            self._event_queue.put(('release', key))
        try:
            self._listener = keyboard.Listener(on_press=on_press, on_release=on_release)
            self._listener.start()
            self._reader_thread = threading.Thread(target=self._reader_worker, daemon=True)
            self._reader_thread.start()
            print("[INFO] 键盘监听已启动")
        except Exception as e:
            print(f"[WARN] 无法初始化键盘监听: {e}")

    def is_key_pressed(self, key):
        k = self._normalize_key(key) if isinstance(key, str) else self._normalize_key(key)
        if k not in self.shared_keys_held or k in self.shared_one_shot:
            return False
        self.shared_one_shot.add(k)
        return True

    def is_key_held(self, key):
        k = self._normalize_key(key) if isinstance(key, str) else self._normalize_key(key)
        return k in self.shared_keys_held

    def restore(self):
        self.running = False
        if self._listener:
            self._listener.stop()


# ============================================================
# Sensor 读取 
# ============================================================
def get_sensor(model, data, name):
    sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, name)
    if sid < 0:
        return None
    adr = model.sensor_adr[sid]
    dim = model.sensor_dim[sid]
    return data.sensordata[adr:adr + dim].copy()


def compute_observations_motrix(model, data, commands, last_actions):
    """计算 45 维观测 (去掉局部线速度，策略仅靠命令+关节信息+陀螺仪来推理)

    布局 (45 dims):
      [0:3]   陀螺仪 (ang_vel * 0.25)
      [3:6]   重力向量 (躯干坐标系，无缩放)
      [6:18]  关节位置偏差 (dof_pos * 1.0)
      [18:30] 关节速度 (dof_vel * 0.05)
      [30:42] 上一步动作 (原始值)
      [42:45] 命令 [vx*2.0, vy*2.0, wz*0.25]
    """
    obs = np.zeros(NUM_OBS, dtype=np.float32)

    # 陀螺仪
    gyro = get_sensor(model, data, "gyro")
    if gyro is not None:
        obs[0:3] = gyro * OBS_SCALES['ang_vel']
    else:
        obs[0:3] = data.qvel[3:6] * OBS_SCALES['ang_vel']

    # 重力向量 (躯干坐标系)
    base_rot = data.xmat[1].reshape(3, 3)
    gravity_world = np.array([0., 0., -1.], dtype=np.float64)
    local_gravity = base_rot.T @ gravity_world
    obs[3:6] = local_gravity.astype(np.float32)

    # 关节位置偏差
    joint_pos = data.qpos[7:19]
    dof_pos_rel = (joint_pos - DEFAULT_JOINT_ANGLES) * OBS_SCALES['dof_pos']
    obs[6:18] = dof_pos_rel

    # 关节速度
    joint_vel = data.qvel[6:18]
    obs[18:30] = joint_vel * OBS_SCALES['dof_vel']

    # 上一步动作
    obs[30:42] = last_actions

    # 命令
    obs[42:45] = commands * np.array([OBS_SCALES['lin_vel'], OBS_SCALES['lin_vel'], OBS_SCALES['ang_vel']], dtype=np.float32)

    # 限幅
    obs = np.clip(obs, -CLIP_OBSERVATIONS, CLIP_OBSERVATIONS)
    return obs


def main():
    import re
    import onnxruntime as ort

    parser = argparse.ArgumentParser(description="MotrixLab Go1 Policy Inference in MuJoCo")
    parser.add_argument("--onnx", type=str, default=DEFAULT_ONNX_PATH)
    parser.add_argument("--terrain", type=str, default=TERRAIN, choices=["none", "rough", "stairs"])
    args = parser.parse_args()

    os.chdir(MOTRIX_XML_DIR)

    if args.terrain == "rough":
        xml_file = f"{MOTRIX_XML_DIR}/scene_rough_terrain.xml"
    elif args.terrain == "stairs":
        xml_file = f"{MOTRIX_XML_DIR}/scene_stairs_terrain.xml"
    else:
        xml_file = f"{MOTRIX_XML_DIR}/scene_motor_actuator.xml"  # flat floor

    with open(xml_file, 'r') as f:
        xml_content = f.read()

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

    print(f"[INFO] Model: {model.nbody} bodies, {model.nq} DoF, {model.nu} actuators")
    print(f"[INFO] MuJoCo timestep: {model.opt.timestep}")
    print(f"[INFO] 关节顺序 (qpos[7:19]): {[mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i) for i in range(1, 13)]}")

    for sensor_name in ["gyro", "local_linvel"]:
        sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, sensor_name)
        print(f"[SENSOR] {sensor_name}: {'存在' if sid >= 0 else '不存在'}")

    # 初始化
    data.qpos[0:3] = np.array([0.0, 0.0, 0.42])
    data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0])  # 默认四元数
    data.qpos[7:19] = DEFAULT_JOINT_ANGLES
    data.qvel[:] = 0.0
    data.ctrl[:] = 0.0
    mujoco.mj_forward(model, data)

    print(f"[INIT] qpos[2]={data.qpos[2]:.3f}")
    print(f"[INIT] qpos[7:19]={data.qpos[7:19]}")
    print(f"[INIT] DEFAULT_JOINT_ANGLES={DEFAULT_JOINT_ANGLES}")

    # 加载onnx
    session = ort.InferenceSession(args.onnx, providers=['CPUExecutionProvider'])
    print(f"[INFO] loaded")

    # 主循环
    ctrl_dt = 0.01  # 100Hz
    num_steps_per_inference = int(ctrl_dt / model.opt.timestep)
    print(f"[INFO] 每 {num_steps_per_inference} 步推理一次 ")

    step_count = 0
    inference_step = 0

    x_vel_cmd = 0.0
    y_vel_cmd = 0.0
    yaw_vel_cmd = 0.0
    commands = np.array([x_vel_cmd, y_vel_cmd, yaw_vel_cmd], dtype=np.float32)
    last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
    action = np.zeros(NUM_ACTIONS, dtype=np.float32)

    # 键盘
    keyboard_reader = KeyboardReader()
    keyboard_reader.init()

    view = viewer.launch_passive(model, data)
    print("[INFO] 已启动!")

    loop_start_time = time.time()

    while view.is_running() and not g_exit_requested:
        # 键盘命令
        if keyboard_reader.is_key_pressed(' '):
            x_vel_cmd = 0.0
            y_vel_cmd = 0.0
            yaw_vel_cmd = 0.0

        if keyboard_reader.is_key_held('w'):
            x_vel_cmd = MAX_LIN_VEL_X
        elif keyboard_reader.is_key_held('s'):
            x_vel_cmd = -MAX_LIN_VEL_X
        else:
            x_vel_cmd = 0.0

        if keyboard_reader.is_key_held('q'):
            y_vel_cmd = MAX_LIN_VEL_Y
        elif keyboard_reader.is_key_held('e'):
            y_vel_cmd = -MAX_LIN_VEL_Y
        else:
            y_vel_cmd = 0.0

        if keyboard_reader.is_key_held('a'):
            yaw_vel_cmd = MAX_ANG_VEL
        elif keyboard_reader.is_key_held('d'):
            yaw_vel_cmd = -MAX_ANG_VEL
        else:
            yaw_vel_cmd = 0.0

        if keyboard_reader.is_key_pressed('r'):
            # 重置机器人到初始位置
            data.qpos[0:3] = np.array([0.0, 0.0, 0.42])
            data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0])
            data.qpos[7:19] = DEFAULT_JOINT_ANGLES
            data.qvel[:] = 0.0
            data.ctrl[:] = 0.0
            last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
            mujoco.mj_forward(model, data)
            print("[RESET] 机器人已重置")

        if keyboard_reader.is_key_pressed('escape'):
            break

        # 推理 (每 N 步一次)
        if inference_step == 0:
            commands[0] = x_vel_cmd
            commands[1] = y_vel_cmd
            commands[2] = yaw_vel_cmd

            obs = compute_observations_motrix(model, data, commands, last_actions)

            # 推理
            action = session.run(None, {'observations': obs.reshape(1, -1).astype(np.float32)})[0][0]
            action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)
            last_actions = action.copy()

        # PD 控制
        # joint_targets = action * action_scale + default_angles
        joint_targets = DEFAULT_JOINT_ANGLES + action * ACTION_SCALE

        current_pos = data.qpos[7:19]
        current_vel = data.qvel[6:18]
        torques = KP * (joint_targets - current_pos) - KD * current_vel
        torques = np.clip(torques, -CLIP_ACTIONS, CLIP_ACTIONS)
        data.ctrl[:] = torques

        mujoco.mj_step(model, data)
        view.sync()

        expected_time = step_count * ctrl_dt
        elapsed = time.time() - loop_start_time
        sleep_time = expected_time - elapsed
        if sleep_time > 0:
            time.sleep(sleep_time)

        step_count += 1
        inference_step = (inference_step + 1) % num_steps_per_inference

        if step_count % 200 == 0:
            trunk_z = data.qpos[2]
            lin_vel = np.linalg.norm(data.qvel[0:3])
            print(f"\n========== Step {step_count} ==========")
            print(f"[CMD] x={x_vel_cmd:.2f}, y={y_vel_cmd:.2f}, yaw={yaw_vel_cmd:.2f}")
            print(f"[OBS] gyro={obs[0:3]}, grav={obs[3:6]}")
            print(f"[ACTION] raw={action[:4]}... scaled={action[:4]*ACTION_SCALE}...")
            print(f"[TARGET] {joint_targets[:4]}...")
            print(f"[TORQUE] {torques[:4]}...")
            print(f"[STATE] z={trunk_z:.3f}m, vel={lin_vel:.3f}m/s")
            print(f"==========================================\n")

    keyboard_reader.restore()
    view.close()


if __name__ == "__main__":
    main()