"""Interactive MuJoCo full replay, or export the exact same poses as video."""
import argparse,os,time,json
from pathlib import Path
p=argparse.ArgumentParser(description=__doc__)
p.add_argument('--directory',type=Path,default=Path(__file__).resolve().parents[1]/'output/l20_full_replay')
p.add_argument('--start',type=int,default=0);p.add_argument('--render',action='store_true');p.add_argument('--frames',type=int,default=0);p.add_argument('--check',action='store_true')
a=p.parse_args()
if a.render or a.check:os.environ.setdefault('MUJOCO_GL','osmesa')
import numpy as np
import mujoco
motion=np.load(a.directory/'motion.npz');q=motion['qpos'];fps=float(motion['fps']);N=len(q)
m=mujoco.MjModel.from_xml_path(str(a.directory/'scene.xml'));d=mujoco.MjData(m)
option=mujoco.MjvOption();option.sitegroup[:]=0;option.geomgroup[3:]=0
cam=mujoco.MjvCamera();cam.lookat[:]=motion['camera_center'];cam.distance=1.05;cam.azimuth=135;cam.elevation=-23
if a.check:
 for row in q:d.qpos[:]=row;mujoco.mj_forward(m,d)
 print(json.dumps(dict(frames=N,nq=m.nq,actuators=m.nu,finite=bool(np.isfinite(q).all()),mode='kinematic')))
elif a.render:
 import imageio.v2 as imageio
 from PIL import Image,ImageDraw,ImageFont
 total=min(a.frames or N,N-a.start)
 renderer=mujoco.Renderer(m,height=720,width=640)
 close=mujoco.MjvCamera();close.distance=.62;close.azimuth=35;close.elevation=-12
 font=ImageFont.truetype('/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf',16)
 path=a.directory/('preview.mp4' if a.frames else 'full_replay.mp4')
 writer=imageio.get_writer(path,fps=fps,codec='libx264',quality=8,macro_block_size=1)
 try:
  for t in range(a.start,a.start+total):
   d.qpos[:]=q[t];mujoco.mj_forward(m,d)
   close.lookat[:]=(d.xpos[m.body('hand_base_link').id]+d.xpos[m.body('right_object').id])/2;close.lookat[2]+=.06
   panels=[]
   for c in [cam,close]:renderer.update_scene(d,camera=c,scene_option=option);panels.append(renderer.render().copy())
   img=Image.fromarray(np.concatenate(panels,1));draw=ImageDraw.Draw(img)
   draw.rectangle((0,0,1280,64),fill=(28,38,49))
   draw.text((18,9),f'L20 + bottle | full video motion | {t/fps:05.2f} / {N/fps:.2f} s',font=font,fill='white')
   draw.text((18,36),'World view - moving wrist',font=font,fill=(195,220,235));draw.text((660,36),'Grasp detail - linear URDF coupling',font=font,fill=(195,220,235))
   draw.rectangle((0,690,1280,720),fill=(28,38,49));draw.text((18,696),'Kinematic contact replay | estimated bottle trajectory | floor + ambient lighting',font=font,fill=(210,220,230))
   writer.append_data(np.asarray(img))
   if t in [a.start,0,160,600,N//2,N-1]:img.save(a.directory/f'preview_{t:04d}.png')
   if t%100==0:print(f'Render {t}/{total}',flush=True)
 finally:writer.close();renderer.close()
 print('RENDER_COMPLETE',path,flush=True)
else:
 import mujoco.viewer
 state={'paused':False,'frame':0.,'last':time.perf_counter()}
 def key(k):
  if k==32:state['paused']=not state['paused']
  elif k in [82,114]:state['frame']=0.
  elif k==262:state['frame']=(int(state['frame'])+1)%N;state['paused']=True
  elif k==263:state['frame']=(int(state['frame'])-1)%N;state['paused']=True
 with mujoco.viewer.launch_passive(m,d,key_callback=key) as viewer:
  viewer.cam.lookat[:]=cam.lookat;viewer.cam.distance=cam.distance;viewer.cam.azimuth=cam.azimuth;viewer.cam.elevation=cam.elevation
  viewer.opt.sitegroup[:]=0;viewer.opt.geomgroup[3:]=0
  while viewer.is_running():
   now=time.perf_counter();elapsed=now-state['last'];state['last']=now
   if not state['paused']:state['frame']=(state['frame']+elapsed*fps)%N
   with viewer.lock():d.qpos[:]=q[int(state['frame'])];d.time=int(state['frame'])/fps;mujoco.mj_forward(m,d)
   viewer.sync();time.sleep(.005)