#!/usr/bin/env python3
"""
sim2sim test for WTW deploy_wtw_pro_sdk.py using MuJoCo Go1 XML.
Parameters verified against go1_walk_these_ways_inference.py reference.

Usage:
  conda activate free_dog_sdk
  mjpython sim2sim_wtw_test.py

Controls: W/S=前后 Q/E=左右 A/D=旋转 Space=停 R=重置 1-4=步态 Esc=退出
"""

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

HERE = os.path.dirname(os.path.abspath(__file__))

# ─── Paths ───
XML_PATH = os.path.join(HERE, "..", "sim2sim_mujoco_example", "data", "go1", "xml", "go1.xml")
MESH_DIR = os.path.join(HERE, "..", "sim2sim_mujoco_example", "data", "go1", "meshes")
BODY_JIT = os.path.join(HERE, "body_latest.jit")
ADAPT_JIT = os.path.join(HERE, "adaptation_module_latest.jit")

# ─── WTW constants (verified against reference) ───
NUM_OBS = 70
NUM_ACTIONS = 12
NUM_COMMANDS = 15
NUM_OBS_HISTORY = 30
OBS_BUFFER_SIZE = 2100

ACTION_SCALE = 0.25
HIP_SCALE_REDUCTION = 0.5
CLIP_ACTIONS = 10.0
CLIP_OBS = 100.0

# Joint orders:
#   MuJoCo/SDK: [FR_hip,FR_thigh,FR_calf, FL_hip,FL_thigh,FL_calf, RR_hip,RR_thigh,RR_calf, RL_hip,RL_thigh,RL_calf]
#   WTW/Deploy: [FL_hip,FL_thigh,FL_calf, FR_hip,FR_thigh,FR_calf, RL_hip,RL_thigh,RL_calf, RR_hip,RR_thigh,RR_calf]
DEPLOY_TO_MUJOCO = np.array([3,4,5, 0,1,2, 9,10,11, 6,7,8], dtype=np.int64)

# Default angles in WTW order (from reference)
DEFAULT_WTW = np.array([
     0.1,  0.8, -1.5,   # FL
    -0.1,  0.8, -1.5,   # FR
     0.1,  1.0, -1.5,   # RL
    -0.1,  1.0, -1.5,   # RR
], dtype=np.float32)
DEFAULT_MUJOCO = np.array([
    -0.1,  0.8, -1.5,   # FR
     0.1,  0.8, -1.5,   # FL
    -0.1,  1.0, -1.5,   # RR
     0.1,  1.0, -1.5,   # RL
], dtype=np.float32)

# Commands scale (verified)
COMMANDS_SCALE = np.array([
    2.0, 2.0, 0.25,    # vx, vy, wz
    2.0,                 # body_height
    1, 1, 1, 1, 1,      # freq, phase, offset, bound, duration
    0.15,                # footswing_height
    0.3, 0.3,           # body_pitch, body_roll
    1.0, 1.0,           # stance_width, stance_length
    1.0,                 # aux_reward
], dtype=np.float32)[:15]

OBS_SCALES = {"dof_pos":1.0, "dof_vel":0.05}

# PD gains (from reference)
KP = 20.0; KD = 0.1    # matches reference (XML passive damping=1.0, total≈1.1)

# Friction (training had zero floor friction)
FLOOR_FRICTION = [0.0, 0.0, 0.0]
BODY_FRICTION = [0.6, 0.3, 0.3]

# Gait presets
GAITS = {
    '1': ('Trot',  0.5, 0.0, 0.0),
    '2': ('Pace',  0.0, 0.0, 0.5),
    '3': ('Bound', 0.0, 0.5, 0.0),
    '4': ('Pronk', 0.0, 0.0, 0.0),
}

EXIT = False
def _sig(s, f): global EXIT; EXIT = True
signal.signal(signal.SIGINT, _sig)


class Keyboard:
    def __init__(self):
        self.held = set(); self._l = None
    def _p(self, k):
        try: self.held.add(k.char.lower())
        except: self.held.add(str(k))
    def _r(self, k):
        try: self.held.discard(k.char.lower())
        except: self.held.discard(str(k))
    def init(self):
        from pynput import keyboard
        self._l = keyboard.Listener(on_press=self._p, on_release=self._r); self._l.start()
    def keys(self): return self.held.copy()
    def stop(self):
        if self._l: self._l.stop()


def main():
    if not os.path.exists(BODY_JIT):
        print(f"[ERROR] body not found: {BODY_JIT}"); return
    if not os.path.exists(ADAPT_JIT):
        print(f"[ERROR] adapt not found: {ADAPT_JIT}"); return

    # Load MuJoCo with mesh path fix
    with open(XML_PATH) as f:
        xml = f.read()
    xml = xml.replace('meshdir="../meshes/"', f'meshdir="{MESH_DIR}"')
    model = mujoco.MjModel.from_xml_string(xml)
    data = mujoco.MjData(model)

    # Set friction
    for i in range(model.ngeom):
        model.geom_friction[i] = BODY_FRICTION

    print(f"[INFO] MuJoCo: {model.nbody} bodies, {model.nq} DoF, KP={KP}, KD(active)={KD}")

    # Load models
    body = torch.jit.load(BODY_JIT, map_location='cpu').eval()
    adapt = torch.jit.load(ADAPT_JIT, map_location='cpu').eval()
    print(f"[INFO] WTW body+adapt loaded")

    # Init
    data.qpos[0:3] = [0, 0, 0.35]
    data.qpos[3:7] = [1, 0, 0, 0]
    data.qpos[7:19] = DEFAULT_MUJOCO
    mujoco.mj_forward(model, data)

    obs_buffer = np.zeros(OBS_BUFFER_SIZE, dtype=np.float32)
    prev_action = np.zeros(12, dtype=np.float32)
    last_action = np.zeros(12, dtype=np.float32)
    gait_idx = 0.0

    step, vx, vy, wz = 0, 0.0, 0.0, 0.0
    fh = 0.15  # footswing height
    bp, br = 0.0, 0.0  # body pitch/roll
    gait_phase, gait_offset, gait_bound, gait_dur = 0.5, 0.0, 0.0, 0.5
    current_gait = '1'
    ctrl_dt = 0.02

    kb = Keyboard(); kb.init()
    view = viewer.launch_passive(model, data)

    print("[INFO] W/S=前后 Q/E=左右 A/D=旋转 1-4=步态 R=重置 Esc=退出")

    t0 = time.perf_counter()

    while view.is_running() and not EXIT:
        keys = kb.keys()
        if 'key.esc' in keys: break

        if 'r' in keys:
            data.qpos[0:3]=[0,0,0.35]; data.qpos[3:7]=[1,0,0,0]
            data.qpos[7:19]=DEFAULT_MUJOCO; data.qvel[:]=0
            obs_buffer[:]=0; prev_action[:]=0; last_action[:]=0; gait_idx=0
            mujoco.mj_forward(model, data)

        # Gait switch
        for k, (name, ph, off, bd) in GAITS.items():
            if k in keys and k != current_gait:
                current_gait = k
                gait_phase, gait_offset, gait_bound = ph, off, bd
                print(f"[INFO] Gait: {name}")

        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=3.0 if 'a' in keys else (-3.0 if 'd' in keys else 0.0)
        if ' ' in keys: vx=vy=wz=0.0

        # Inference at 50Hz (every 10 sim steps at dt=0.002)
        if step % 10 == 0:
                # Commands
                raw_cmd = np.zeros(15, dtype=np.float32)
                raw_cmd[0]=vx; raw_cmd[1]=vy; raw_cmd[2]=wz
                raw_cmd[3]=0.0; raw_cmd[4]=3.0
                raw_cmd[5]=gait_phase; raw_cmd[6]=gait_offset; raw_cmd[7]=gait_bound; raw_cmd[8]=gait_dur
                raw_cmd[9]=0.15; raw_cmd[10]=bp; raw_cmd[11]=br
                raw_cmd[12]=0.25; raw_cmd[13]=0.4
                commands = raw_cmd * COMMANDS_SCALE

                # Gait index & clock
                gait_idx += 0.02 * 3.0
                if gait_idx > 1.0: gait_idx -= 1.0
                p, o, b = gait_phase, gait_offset, gait_bound
                fi = [gait_idx+p+o+b, gait_idx+o, gait_idx+b, gait_idx+p]
                clock = np.array([np.sin(2*np.pi*f) for f in fi], dtype=np.float32)

                # Observation
                obs = np.zeros(NUM_OBS, dtype=np.float32)
                base_rot = data.xmat[1].reshape(3, 3)
                obs[0:3] = (base_rot.T @ np.array([0., 0., -1.], dtype=np.float64)).astype(np.float32)
                obs[3:18] = commands
                dof_wtw = data.qpos[7:19][DEPLOY_TO_MUJOCO]
                obs[18:30] = (dof_wtw - DEFAULT_WTW) * 1.0
                obs[30:42] = data.qvel[6:18][DEPLOY_TO_MUJOCO] * 0.05
                obs[42:54] = np.clip(prev_action, -CLIP_ACTIONS, CLIP_ACTIONS)
                obs[54:66] = np.clip(last_action, -CLIP_ACTIONS, CLIP_ACTIONS)
                obs[66:70] = clock
                obs = np.clip(obs, -CLIP_OBS, CLIP_OBS)

                # History buffer
                obs_buffer = np.concatenate([obs_buffer[NUM_OBS:], obs])

                # Inference
                obs_hist = torch.from_numpy(obs_buffer).float().unsqueeze(0)
                with torch.inference_mode():
                    latent = adapt(obs_hist)
                    action = body(torch.cat([obs_hist, latent], dim=1)).numpy().flatten()
                action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)
                last_action = prev_action.copy()
                prev_action = action.copy()

        # PD control
        action_scaled = prev_action * ACTION_SCALE
        for i in [0,3,6,9]: action_scaled[i] *= HIP_SCALE_REDUCTION
        targets_mujoco = action_scaled[DEPLOY_TO_MUJOCO] + DEFAULT_MUJOCO
        torques = KP*(targets_mujoco - data.qpos[7:19]) - KD*data.qvel[6:18]
        data.ctrl[:] = np.clip(torques, -23.7, 23.7)

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

        if step % 200 == 0:
            print(f"[STEP {step}] z={data.qpos[2]:.3f} cmd=[{vx:.1f},{vy:.1f},{wz:.1f}] "
                  f"pos=[{data.qpos[0]:.2f},{data.qpos[1]:.2f}]")

        step += 1
        expected = (step+1)*model.opt.timestep
        sleep = expected - (time.perf_counter()-t0)
        if sleep > 0: time.sleep(sleep)

    kb.stop(); view.close()
    print("[INFO] Done.")

if __name__ == "__main__":
    main()
