#!/usr/bin/env python3
"""Deploy Go1 AMP model in MuJoCo."""
import mujoco, numpy as np, os, torch, time, sys
from mujoco import viewer

# Use AMP's rsl_rl (NOT MotrixLab's)
sys.path.insert(0, '/home/8x54zj-m/amp_go2/rsl_rl')
sys.path.insert(0, '/home/8x54zj-m/amp_go2/legged_gym')
# Force reload of rsl_rl to pick up AMP version
for mod in list(sys.modules.keys()):
    if 'rsl_rl' in mod:
        del sys.modules[mod]

XML_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                       'motrix_envs/src/motrix_envs/locomotion/go1/xmls')
CKPT = '/home/8x54zj-m/amp_go2/legged_gym/logs/go1_amp/Jul05_01-15-41_flat/model_6000.pt'

import argparse
p = argparse.ArgumentParser()
p.add_argument("--terrain", default="flat", choices=["flat", "flat_stairs"])
args = p.parse_args()

scene = "scene_flat_stairs.xml" if args.terrain == "flat_stairs" else "go1_motor_actuator.xml"
os.chdir(XML_DIR)
m = mujoco.MjModel.from_xml_string(open(scene).read())
d = mujoco.MjData(m)

# Spawn on level 1 platform if using flat_stairs
spawn_x, spawn_y, spawn_z = 0.0, 0.0, 0.45
if args.terrain == "flat_stairs":
    col = np.random.randint(0, 4)
    spawn_x = -12.0 + col * 8.0
    spawn_y = -4.0  # level 1 center
    # Sample terrain height at spawn
    floor_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "floor")
    if floor_id >= 0 and m.geom_type[floor_id] == mujoco.mjtGeom.mjGEOM_HFIELD:
        hf = m.geom_dataid[floor_id]
        nr, nc = int(m.hfield_nrow[hf]), int(m.hfield_ncol[hf])
        sx, sy, zt, sb = m.hfield_size[hf]
        hd = m.hfield_data[m.hfield_adr[hf]:m.hfield_adr[hf]+nr*nc].reshape(nr, nc)
        gp = m.geom_pos[floor_id]
        col_idx = int(np.clip(((spawn_x-gp[0])/sx*0.5+0.5)*(nc-1), 0, nc-1))
        row_idx = int(np.clip(((spawn_y-gp[1])/sy*0.5+0.5)*(nr-1), 0, nr-1))
        spawn_z = float(gp[2] + sb + hd[row_idx, col_idx] * zt) + 0.45
    print(f"[Terrain] flat_stairs level=1 spawn=({spawn_x:.1f},{spawn_y:.1f},{spawn_z:.2f})")

ckpt = torch.load(CKPT, map_location='cpu', weights_only=False)
normalizer = ckpt['amp_normalizer']

actor = torch.nn.Sequential(
    torch.nn.Linear(45, 512), torch.nn.ELU(),
    torch.nn.Linear(512, 256), torch.nn.ELU(),
    torch.nn.Linear(256, 128), torch.nn.ELU(),
    torch.nn.Linear(128, 12)
)
sd = {k.replace('actor.',''): v for k,v in ckpt['model_state_dict'].items() if k.startswith('actor.')}
actor.load_state_dict(sd)
actor.eval()

DEFAULT = np.array([-0.1,0.8,-1.5, 0.1,0.8,-1.5, -0.1,1.0,-1.5, 0.1,1.0,-1.5], dtype=np.float32)
LIN_VEL_SCALE, ANG_VEL_SCALE = 2.0, 0.25
DOF_POS_SCALE, DOF_VEL_SCALE = 1.0, 0.05
CMD_SCALE = np.array([2.0, 2.0, 0.25], dtype=np.float32)
KP, KD, ACT_SCALE = 20.0, 0.5, 0.25

def get_gravity(q):
    qw,qx,qy,qz = q[3],q[0],q[1],q[2]
    return np.array([2*(-qz*qx+qw*qy), -2*(qz*qy+qw*qx), 1-2*(qw*qw+qz*qz)])

def compute_obs(d, last_action, cmd):
    obs = np.zeros(45, dtype=np.float32)
    obs[0:3] = d.qvel[3:6] * ANG_VEL_SCALE
    obs[3:6] = get_gravity(d.qpos[3:7])
    obs[6:9] = cmd * CMD_SCALE
    obs[9:21] = (d.qpos[7:19] - DEFAULT) * DOF_POS_SCALE
    obs[21:33] = d.qvel[6:18] * DOF_VEL_SCALE
    obs[33:45] = last_action
    return obs

d.qpos[7:19] = DEFAULT; d.qpos[0:3] = [spawn_x, spawn_y, spawn_z]; d.qvel[:] = 0
mujoco.mj_forward(m, d)

view = viewer.launch_passive(m, d)
cmd = np.array([0.8, 0.0, 0.0], dtype=np.float32)
last_action = np.zeros(12, dtype=np.float32)

print("AMP Go1 | W/S前后 Q/E左右 A/D旋转 Space停 R重置 Esc退出")

while view.is_running():
    if view.is_running():
        obs = compute_obs(d, last_action, cmd)
        obs_t = torch.from_numpy(obs).unsqueeze(0).float()
        obs_norm = obs_t.clone()
        obs_norm[:, :43] = normalizer.normalize(obs_t[:, :43])
        with torch.no_grad():
            act = actor(obs_norm)[0].numpy()
        act = np.clip(act, -18, 18)
        last_action = act.copy()
        target = DEFAULT + act * ACT_SCALE
        d.ctrl[:] = KP * (target - d.qpos[7:19]) - KD * d.qvel[6:18]
    mujoco.mj_step(m, d)
    view.sync()
    time.sleep(0.001)