"""Interactive wall-clock playback of saved SPIDER states (no re-simulation)."""
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/spider_dynamics_fix_20260915')
p.add_argument('--mode',choices=['physics','reference'],default='physics')
p.add_argument('--check',action='store_true');a=p.parse_args()
os.environ['MUJOCO_GL']='osmesa' if a.check else 'glfw'
import numpy as np,mujoco
root=a.directory.resolve();task=root/'datasets/processed/current/l20/bimanual/boxes'
m=mujoco.MjModel.from_xml_path(str(task/'scene_act.xml'));d=mujoco.MjData(m)
source=root/('reference_video_rate.npz' if a.mode=='reference' else 'physics_motion.npz')
mode='collision-corrected kinematic reference (middle video panel)' if a.mode=='reference' else 'saved physical state playback, no re-simulation'
with np.load(source) as z:
 q=z['qpos'];t=z['time']
 if a.mode=='physics':u=z['ctrl']
if a.mode=='reference':
 v=np.zeros((len(q),m.nv));u=q[:,m.jnt_qposadr[m.actuator_trnid[:,0]]]
else:
 with np.load(task/'0/trajectory_mjwp_act.npz') as z:v=z['qvel'].reshape(len(q),-1)
assert q.shape==(len(t),m.nq) and np.isfinite(q).all() and np.all(np.diff(t)>0)
if a.check:
 for i in [0,len(q)//2,len(q)-1]:
  d.qpos[:]=q[i];d.qvel[:]=v[i];d.ctrl[:]=u[i];mujoco.mj_forward(m,d)
 print(json.dumps({'steps':len(q),'duration_s':float(t[-1]),'nq':m.nq,'finite':True,'mode':mode}));raise SystemExit
import mujoco.viewer
state={'paused':False,'time':0.,'speed':1.,'reset_camera':False}
def key(k):
 if k==32:state['paused']=not state['paused']
 elif k in [82,114]:state['time']=0.
 elif k in [262,263]:state['time']=float(np.clip(state['time']+(1 if k==262 else -1)/30,0,t[-1]));state['paused']=True
 elif k==265:state['speed']=min(4.,state['speed']*2)
 elif k==264:state['speed']=max(.125,state['speed']/2)
 elif k in [67,99]:state['reset_camera']=True
m.vis.headlight.ambient[:]=.7
d.qpos[:]=q[0];mujoco.mj_forward(m,d)
center=np.mean([d.site_xpos[m.site(s+'_object_track').id] for s in ['right','left']],axis=0)
cid=m.camera('source_camera').id;offset=d.cam_xpos[cid]-center
camera={'distance':float(np.linalg.norm(offset)),'azimuth':float(np.degrees(np.arctan2(offset[1],offset[0]))),'elevation':float(-np.degrees(np.arctan2(offset[2],np.linalg.norm(offset[:2]))))}
status=root/'interactive_playback_status.json'
print(mode+' | SPACE pause | R restart | LEFT/RIGHT video-frame step | UP/DOWN speed | C camera reset | mouse orbit/zoom',flush=True)
with mujoco.viewer.launch_passive(m,d,key_callback=key) as viewer:
 def reset_camera():
  viewer.cam.type=mujoco.mjtCamera.mjCAMERA_FREE;viewer.cam.lookat[:]=center
  for k,x in camera.items():setattr(viewer.cam,k,x)
 with viewer.lock():
  reset_camera();viewer.opt.sitegroup[:]=0;viewer.opt.geomgroup[3:]=0
 last=time.perf_counter();last_status=0
 while viewer.is_running():
  now=time.perf_counter();elapsed=now-last;last=now
  if not state['paused']:state['time']=(state['time']+elapsed*state['speed'])%float(t[-1])
  i=int(np.clip(np.searchsorted(t,state['time'],side='right')-1,0,len(t)-1))
  with viewer.lock():
   if state['reset_camera']:reset_camera();state['reset_camera']=False
   d.qpos[:]=q[i];d.qvel[:]=v[i];d.ctrl[:]=u[i];d.time=float(t[i]);mujoco.mj_forward(m,d)
  viewer.sync()
  if now-last_status>1:
   status.write_text(json.dumps({'running':True,'pid':os.getpid(),'mode':mode,'source':str(source),'step':i,'time_s':float(t[i]),'speed':state['speed'],'paused':state['paused'],'updated_unix':time.time()},indent=2));last_status=now
  time.sleep(.008)
status.write_text(json.dumps({'running':False,'pid':os.getpid(),'updated_unix':time.time()}))