#!/usr/bin/env python
"""VLA closed-loop на MicroDuck: 61D obs из сима + head_camera рендер -> чанк -> исполнение."""
import sys, json
from pathlib import Path
import numpy as np
import torch, mujoco
sys.path.insert(0,'/root/tinyvla_h100'); sys.path.insert(0,'/root/tinyvla_h100/src')
from closedloop import model, data, key, HOME, SPC, build_obs
from scripts.eval_microduck import BEHAVIORS, DUCK_MORPH, DUCK_TEXT, load_split, behavior_of, decode_img
from tinyvla.modeling_tinyvla import TinyVLAPolicy
from tinyvla.modules.embodiment import MORPH_FIELDS
from transformers import AutoTokenizer
from collections import defaultdict

torch.backends.cuda.enable_cudnn_sdp(False)
CKPT='/root/outputs/microduck_ft/final'
beh=sys.argv[1] if len(sys.argv)>1 else 'walking'
SECONDS=float(sys.argv[2]) if len(sys.argv)>2 else 8

man=json.loads(Path('/root/microduck/train/manifest.json').read_text())
s_mean=np.asarray(man['norm_stats']['state']['mean'],np.float32)
s_std=np.maximum(np.asarray(man['norm_stats']['state']['std'],np.float32),1e-6)
a_mean=np.asarray(man['norm_stats']['action']['mean'],np.float32)
a_std=np.maximum(np.asarray(man['norm_stats']['action']['std'],np.float32),1e-6)

pol=TinyVLAPolicy.from_pretrained(CKPT).cuda().eval(); cfg=pol.config
tok=AutoTokenizer.from_pretrained(cfg.lm_model_name)
sc={"arm_dof":0.1,"reach_m":2,"gripper_width_m":10,"num_cameras":1/3,"control_hz":1/30,
    "joint_lo_mean":1/3.1416,"joint_hi_mean":1/3.1416,"workspace_x":2,"workspace_y":2,"workspace_z":2,"payload_kg":0.2}
morph=torch.tensor([DUCK_MORPH.get(f,0)*sc.get(f,1) for f in MORPH_FIELDS],dtype=torch.float32)[None].cuda()
mt=tok([DUCK_TEXT],padding="max_length",truncation=True,max_length=cfg.morph_text_max_len,return_tensors="pt")
tt=tok([f"MicroDuck: perform {beh.replace('_',' ')}"],padding="max_length",truncation=True,
       max_length=cfg.tokenizer_max_length,return_tensors="pt")
# демо из train-эпизодов этого поведения + КОМАНДА из реального состояния датасета
eps=load_split(Path('/root/microduck'),'train')
by=defaultdict(list)
for e in sorted(eps): by[behavior_of(e)].append(e)
sup_ep=by[beh][0]; sup=eps[sup_ep]
sidx=np.linspace(0,len(sup)-1,3).astype(int)
sup_img=torch.stack([decode_img(sup[i][1]) for i in sidx])[None].cuda()
sup_act=torch.stack([torch.nn.functional.pad(torch.from_numpy(sup[i][3]),(0,cfg.max_action_dim-14)) for i in sidx])[None].cuda()
raw_state0=sup[len(sup)//2][2]*s_std+s_mean   # денорм состояния из середины эпизода
command=raw_state0[48:61].astype(np.float32)   # родная команда поведения из данных
print(f"команда из датасета ({beh}): vx,vy,vyaw={command[:3].round(3)} bodyz={command[9]:.3f}")

# камера головы: детерминированная установка (yaw180, pitch25 от мира на STAND),
# найдена сверкой с кадрами датасета; RL-сценная quat смотрела внутрь головы
import mujoco as _mj
_cid=model.camera("head_camera").id
_bid=model.cam_bodyid[_cid]
_mj.mj_resetDataKeyframe(model,data,key.id); _mj.mj_forward(model,data)
_bpos=data.xpos[_bid].copy(); _bR=data.xmat[_bid].reshape(3,3).copy()
import numpy as _np
_p=_np.deg2rad(25); _y=_np.deg2rad(180)
_fwd=_np.array([_np.cos(_p)*_np.cos(_y),_np.cos(_p)*_np.sin(_y),-_np.sin(_p)])
_r=_np.cross(_fwd,[0,0,1.0]); _r/=_np.linalg.norm(_r); _u=_np.cross(_r,_fwd)
_Rw=_np.stack([_r,_u,-_fwd],axis=1); _Rl=_bR.T@_Rw
_q=_np.empty(4); _mj.mju_mat2Quat(_q,_Rl.flatten())
model.cam_quat[_cid]=_q
model.cam_pos[_cid]=_bR.T@((_bpos+_np.array([-0.06,0,0.0]))-_bpos)
model.vis.map.znear=0.02
DS_MEAN=0.116  # средняя яркость кадров датасета
head_renderer=mujoco.Renderer(model,height=256,width=256)
_gain=[None]
def head_img():
    head_renderer.update_scene(data,camera="head_camera")
    a=head_renderer.render().astype("float32")/255.
    if _gain[0] is None:
        _gain[0]=DS_MEAN/max(a.mean(),1e-3)
    a=(a*_gain[0]).clip(0,1)
    return torch.from_numpy(a).permute(2,0,1).float()

def vla_step(obs61,i):
    st=torch.from_numpy((obs61-s_mean)/s_std)
    st=torch.nn.functional.pad(st,(0,cfg.max_state_dim-61))
    img=head_img()
    b={"observation.images.cam0":img[None].cuda(),
       "observation.images.cam1":torch.zeros_like(img)[None].cuda(),
       "observation.state":st[None].cuda(),
       "observation.language.tokens":tt["input_ids"].cuda(),
       "observation.language.attention_mask":tt["attention_mask"].bool().cuda(),
       "morph_text_ids":mt["input_ids"].cuda(),"morph_text_mask":mt["attention_mask"].bool().cuda(),
       "morphology":morph,"support_images":sup_img,"support_actions":sup_act,
       "embodiment_id":torch.tensor([0]).cuda()}
    with torch.no_grad(), torch.autocast("cuda",torch.bfloat16):
        pr=pol.predict_action_chunk(b)[0].float().cpu().numpy()[:,:14]
    return pr*a_std+a_mean                     # денорм: единицы учителя (оффсеты от home)

from closedloop import run
run(vla_step, command, seconds=SECONDS, tag=f"vla_{beh}")