#!/usr/bin/env python
"""Closed-loop MicroDuck: obs.rs-раскладка 61D из MuJoCo, учитель ONNX или VLA.

Валидация конвенций через учителя: если alpha_walking в этой петле идёт —
obs-билдер и применение действий верны, и той же петлёй можно гнать VLA.

obs.rs layout:
  0..3   gyro (trunk frame, rad/s)
  3..6   projected gravity (trunk frame, unit)
  6..20  joint pos - home (14, без рта)
  20..34 joint vel (14)
  34..48 previous action (14)
  48..61 command: vx,vy,vyaw | 4 head targets | body x,y(=0) | z | roll | pitch | yaw(=0)
Действие: 14 позиционных таргетов; применение ctrl = home + action (STAND offset,
как в mjlab JointPositionAction; teacher выход в тех же единицах).
"""
import sys, json
from pathlib import Path
import numpy as np
import mujoco

SCENE='/root/microduck_rl/src/mjlab_microduck/robot/microduck/scene_walk.xml'
model=mujoco.MjModel.from_xml_path(SCENE)
data=mujoco.MjData(model)
key=model.keyframe('STAND')
HOME=key.ctrl.copy()                      # home-поза 14 суставов
CTRL_DT=1/50
SPC=max(1,int(round(CTRL_DT/model.opt.timestep)))

def quat_rotate_inv(q, v):
    w,x,y,z=q
    # R^T v
    R=np.array([[1-2*(y*y+z*z),2*(x*y-w*z),2*(x*z+w*y)],
                [2*(x*y+w*z),1-2*(x*x+z*z),2*(y*z-w*x)],
                [2*(x*z-w*y),2*(y*z+w*x),1-2*(x*x+y*y)]])
    return R.T@v

def build_obs(prev_action, command):
    q=data.qpos[3:7]                       # trunk quat (w,x,y,z)
    gyro=data.qvel[3:6].copy()             # угловая скорость в ЛОКАЛЬНОЙ раме (mujoco freejoint)
    grav=quat_rotate_inv(q, np.array([0,0,-1.0]))
    jpos=data.qpos[7:7+14]-HOME
    jvel=data.qvel[6:6+14]
    return np.concatenate([gyro,grav,jpos,jvel,prev_action,command]).astype(np.float32)

def run(policy_step, command, seconds=8, tag="run", render_cam=None, vla_hooks=None):
    import imageio.v2 as imageio
    renderer=mujoco.Renderer(model,height=320,width=320)
    cam=mujoco.MjvCamera(); cam.distance,cam.azimuth,cam.elevation=0.65,140,-12
    mujoco.mj_resetDataKeyframe(model,data,key.id); mujoco.mj_forward(model,data)
    prev=np.zeros(14,np.float32)
    frames=[]; fell_at=None
    n_ctrl=int(seconds*50)
    for i in range(n_ctrl):
        if i%5==0:  # 10 Гц перепланирование
            obs=build_obs(prev,command)
            plan=policy_step(obs, i)       # (k,14) минимум 5 шагов
        a=plan[i%5] if plan.ndim==2 else plan
        data.ctrl[:]=HOME+a
        prev=a.astype(np.float32)
        for _ in range(SPC): mujoco.mj_step(model,data)
        if data.qpos[2]<0.05 and fell_at is None: fell_at=i/50
        if i%3==0:
            cam.lookat[:]=[float(data.qpos[0]),float(data.qpos[1]),0.12]
            renderer.update_scene(data,camera=cam)
            frames.append(renderer.render().copy())
    imageio.mimwrite(f"/root/cl_{tag}.mp4",frames,fps=16,quality=7)
    print(f"{tag}: {seconds}s, высота в конце {data.qpos[2]:.3f}, "
          f"{'УПАЛ на '+str(round(fell_at,1))+'s' if fell_at else 'НЕ УПАЛ'}, "
          f"пройдено xy {np.hypot(data.qpos[0],data.qpos[1]):.2f} м", flush=True)
    return fell_at

if __name__=="__main__":
    import onnxruntime as ort
    which=sys.argv[1] if len(sys.argv)>1 else "walking"
    sess=ort.InferenceSession(f"/root/.cache/huggingface/hub/models--pollen-robotics--microduck-policies/snapshots/{'*'}/alpha_{which}.onnx".replace('*','') if False else __import__('huggingface_hub').hf_hub_download("pollen-robotics/microduck-policies", f"alpha_{which}.onnx"))
    def teacher(obs,i):
        a=sess.run(None,{"obs":obs[None].astype(np.float32)})[0][0]
        return np.tile(a,(5,1))            # учитель даёт 1 шаг — держим до реплана
    cmd=np.zeros(13,np.float32)
    if which=="walking": cmd[0]=0.15       # vx м/с
    cmd[9]=0.115                            # body z: стоячая высота
    run(teacher,cmd,seconds=8,tag=f"teacher_{which}")