#!/usr/bin/env python3
"""Go1 sim2sim MuJoCo viewer — Original 30k flat training.

Usage: python go1_sim2sim.py

Requires: mujoco, onnxruntime, pynput
Install: pip install mujoco onnxruntime pynput

Controls:
  W/S: forward/back    Q/E: strafe left/right
  A/D: rotate          Space: stop    R: reset    Esc: quit
"""
import numpy as np, mujoco, onnxruntime as ort, os, time, threading, queue
from mujoco import viewer
from pynput import keyboard

HERE = os.path.dirname(os.path.abspath(__file__))
ONNX = os.path.join(HERE, "90k_45.onnx")

# ── Parameters (original MotrixLab Go1 config) ──
NUM_OBS = 45
KP, KD = 80.0, 0.5      # KD=0.5 + MuJoCo joint_damping(0.5) = 1.0 = training kd
ACTION_SCALE = 0.05
CLIP = 23.7
DEFAULT_ANGLES = np.array([
    -0.0, 0.9, -1.8,
     0.0, 0.9, -1.8,
    -0.0, 0.9, -1.8,
     0.0, 0.9, -1.8,
], dtype=np.float32)

# ── Keyboard ──
class KB:
    def __init__(s):
        s._q = queue.Queue(); s.running = True; s.held = set()
    def _n(s, k):
        try:
            if hasattr(k, 'char') and k.char: return k.char.lower()
        except: pass
        return str(k).lower()
    def _w(s):
        while s.running:
            try:
                et, k = s._q.get(timeout=0.05)
                n = s._n(k)
                if et == 'press': s.held.add(n)
                elif et == 'release': s.held.discard(n)
            except queue.Empty: pass
    def init(s):
        s._l = keyboard.Listener(
            on_press=lambda k: s._q.put(('press', k)),
            on_release=lambda k: s._q.put(('release', k)))
        s._l.start()
        s._t = threading.Thread(target=s._w, daemon=True); s._t.start()
    def keys(s): return s.held.copy()
    def stop(s): s.running = False; s._l.stop()

# ── Main ──
def main():
    # The model XML is embedded below
    xml = '''<mujoco model="go1 scene">
  <compiler angle="radian" autolimits="true"/>
  <option timestep="0.005" integrator="Euler" iterations="60">
    <flag eulerdamp="disable"/>
  </option>
  <custom>
    <numeric data="30" name="max_contact_points"/>
    <numeric data="12" name="max_geom_pairs"/>
  </custom>
  <default>
    <default class="go1">
      <geom condim="1"/>
      <joint axis="0 1 0" armature="0.005" damping="0.5"/>
      <default class="abduction">
        <joint axis="1 0 0" range="-0.863 0.863" frictionloss="0.3"/>
      </default>
      <default class="hip">
        <joint range="-0.686 4.501" frictionloss="0.3"/>
      </default>
      <default class="knee">
        <joint range="-2.818 -0.888" frictionloss="1.0"/>
      </default>
    </default>
  </default>
  <asset>
    <texture name="skybox" type="skybox" builtin="gradient" rgb1="0.4 0.4 0.4" rgb2="0 0 0" width="512" height="512"/>
    <texture name="ground" type="2d" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3" markrgb="0.8 0.8 0.8" width="300" height="300"/>
    <material name="ground" texture="ground" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
  </asset>
  <worldbody>
    <light pos="0 0 1.5" dir="0 0 -1" directional="true"/>
    <geom name="floor" size="0 0 0.01" type="plane" material="ground" contype="1" conaffinity="0" priority="1" friction="0.6" condim="3"/>
    <body name="trunk" pos="0 0 0.4" childclass="go1">
      <freejoint/>
      <inertial pos="0.0223 0.002 -0.0005" quat="-0.00342088 0.705204 0.000106698 0.708996" mass="5.204" diaginertia="0.0716565 0.0630105 0.0168101"/>
      <geom name="trunk_geom" contype="0" conaffinity="0" group="2" type="box" size="0.35 0.12 0.08" rgba="0.4 0.4 0.4 1"/>
      <geom name="trunk_col" contype="1" conaffinity="1" group="3" pos="0.24 0 0" size="0.05 0.05 0.05" type="box"/>
      <site name="imu" pos="-0.01592 -0.06659 -0.00617" group="5"/>
      <!-- FR leg -->
      <body name="FR_hip" pos="0.1881 -0.04675 0">
        <joint class="abduction" name="FR_hip_joint"/>
        <geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.08" rgba="0.5 0.5 0.5 1"/>
        <body name="FR_thigh" pos="0 -0.08 0">
          <joint class="hip" name="FR_thigh_joint"/>
          <geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
          <body name="FR_calf" pos="0 0 -0.213">
            <joint class="knee" name="FR_calf_joint"/>
            <geom contype="0" conaffinity="0" group="2" type="capsule" size="0.04 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
            <geom name="FR_foot" contype="1" conaffinity="1" group="3" type="sphere" size="0.023" pos="0 0 -0.213" priority="10" condim="3"/>
          </body>
        </body>
      </body>
      <!-- FL leg -->
      <body name="FL_hip" pos="0.1881 0.04675 0">
        <joint class="abduction" name="FL_hip_joint"/>
        <geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.08" rgba="0.5 0.5 0.5 1"/>
        <body name="FL_thigh" pos="0 0.08 0">
          <joint class="hip" name="FL_thigh_joint"/>
          <geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
          <body name="FL_calf" pos="0 0 -0.213">
            <joint class="knee" name="FL_calf_joint"/>
            <geom contype="0" conaffinity="0" group="2" type="capsule" size="0.04 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
            <geom name="FL_foot" contype="1" conaffinity="1" group="3" type="sphere" size="0.023" pos="0 0 -0.213" priority="10" condim="3"/>
          </body>
        </body>
      </body>
      <!-- RR leg -->
      <body name="RR_hip" pos="-0.1881 -0.04675 0">
        <joint class="abduction" name="RR_hip_joint"/>
        <geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.08" rgba="0.5 0.5 0.5 1"/>
        <body name="RR_thigh" pos="0 -0.08 0">
          <joint class="hip" name="RR_thigh_joint"/>
          <geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
          <body name="RR_calf" pos="0 0 -0.213">
            <joint class="knee" name="RR_calf_joint"/>
            <geom contype="0" conaffinity="0" group="2" type="capsule" size="0.04 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
            <geom name="RR_foot" contype="1" conaffinity="1" group="3" type="sphere" size="0.023" pos="0 0 -0.213" priority="10" condim="3"/>
          </body>
        </body>
      </body>
      <!-- RL leg -->
      <body name="RL_hip" pos="-0.1881 0.04675 0">
        <joint class="abduction" name="RL_hip_joint"/>
        <geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.08" rgba="0.5 0.5 0.5 1"/>
        <body name="RL_thigh" pos="0 0.08 0">
          <joint class="hip" name="RL_thigh_joint"/>
          <geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
          <body name="RL_calf" pos="0 0 -0.213">
            <joint class="knee" name="RL_calf_joint"/>
            <geom contype="0" conaffinity="0" group="2" type="capsule" size="0.04 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
            <geom name="RL_foot" contype="1" conaffinity="1" group="3" type="sphere" size="0.023" pos="0 0 -0.213" priority="10" condim="3"/>
          </body>
        </body>
      </body>
    </body>
  </worldbody>
  <actuator>
    <motor class="abduction" name="FR_hip" joint="FR_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
    <motor class="hip" name="FR_thigh" joint="FR_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
    <motor class="knee" name="FR_calf" joint="FR_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
    <motor class="abduction" name="FL_hip" joint="FL_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
    <motor class="hip" name="FL_thigh" joint="FL_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
    <motor class="knee" name="FL_calf" joint="FL_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
    <motor class="abduction" name="RR_hip" joint="RR_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
    <motor class="hip" name="RR_thigh" joint="RR_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
    <motor class="knee" name="RR_calf" joint="RR_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
    <motor class="abduction" name="RL_hip" joint="RL_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
    <motor class="hip" name="RL_thigh" joint="RL_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
    <motor class="knee" name="RL_calf" joint="RL_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
  </actuator>
  <sensor>
    <gyro site="imu" name="gyro"/>
    <velocimeter site="imu" name="local_linvel"/>
  </sensor>
</mujoco>'''

    model = mujoco.MjModel.from_xml_string(xml)
    data = mujoco.MjData(model)
    data.qpos[0:3] = [0, 0, 0.42]
    data.qpos[3:7] = [1, 0, 0, 0]
    data.qpos[7:19] = DEFAULT_ANGLES
    mujoco.mj_forward(model, data)

    session = ort.InferenceSession(ONNX, providers=['CPUExecutionProvider'])
    print(f"[Go1 sim2sim] ONNX={ONNX}")
    print(f"[Go1 sim2sim] PD kp={KP} kd={KD+0.5} action_scale={ACTION_SCALE} obs={NUM_OBS}-dim")
    print(f"[Go1 sim2sim] W/S前后 Q/E左右 A/D旋转 Space停 R重置 Esc退出")

    kb = KB(); kb.init()
    view = viewer.launch_passive(model, data)
    step, vx, vy, wz = 0, 0.0, 0.0, 0.0
    last_a = np.zeros(12, dtype=np.float32)

    while view.is_running():
        keys = kb.keys()
        if 'escape' in keys: break
        if 'r' in keys:
            data.qpos[0:3] = [0, 0, 0.42]
            data.qpos[3:7] = [1, 0, 0, 0]
            data.qpos[7:19] = DEFAULT_ANGLES
            data.qvel[:] = 0; last_a[:] = 0
            mujoco.mj_forward(model, data)

        vx = 1.0 if 'w' in keys else (-1.0 if 's' in keys else 0.0)
        vy = 1.0 if 'q' in keys else (-1.0 if 'e' in keys else 0.0)
        wz = 1.0 if 'a' in keys else (-1.0 if 'd' in keys else 0.0)
        if ' ' in keys: vx = vy = wz = 0.0

        if step % 2 == 0:  # 100Hz control (MuJoCo dt=0.005)
            obs = np.zeros(NUM_OBS, dtype=np.float32)
            sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, "gyro")
            adr = model.sensor_adr[sid]
            obs[0:3] = data.sensordata[adr:adr+3] * 0.25
            R = data.xmat[1].reshape(3, 3)
            obs[3:6] = (R.T @ np.array([0., 0., -1.])).astype(np.float32)
            obs[6:18] = (data.qpos[7:19] - DEFAULT_ANGLES) * 1.0
            obs[18:30] = data.qvel[6:18] * 0.05
            obs[30:42] = last_a
            obs[42:45] = np.array([vx, vy, wz]) * np.array([2., 2., 0.25])
            obs = np.clip(obs, -100., 100.)
            action = session.run(None, {'observations': obs.reshape(1, -1).astype(np.float32)})[0][0]
            action = np.clip(action, -CLIP, CLIP)
            last_a = action.copy()

        target = DEFAULT_ANGLES + action * ACTION_SCALE
        torques = KP * (target - data.qpos[7:19]) - KD * data.qvel[6:18]
        data.ctrl[:] = np.clip(torques, -CLIP, CLIP)
        mujoco.mj_step(model, data)
        view.sync()
        step += 1
        time.sleep(0.001)

    kb.stop(); view.close()

if __name__ == "__main__":
    main()
