"""Interpolate rigid rotations on SO(3), choosing a continuous Euler chart for SPIDER."""
from pathlib import Path
import numpy as np,json
from scipy.spatial.transform import Rotation as R,Slerp
O=Path(__file__).resolve().parents[1]/'output/collision_fix_20260915';T=O/'datasets/processed/current/l20/bimanual/boxes/0';a=np.load(O/'reference_video_rate.npz');old=np.load(T/'trajectory_kinematic_act.npz');arrays={k:old[k] for k in old.files};times=arrays['time'];t=a['time'];q=a['qpos'];qi=np.stack([np.interp(times,t,q[:,j]) for j in range(q.shape[1])],1)
def continuous_euler(rot,initial):
 angles=rot.as_euler('XYZ');prev=initial.copy();rows=[]
 for v in angles:
  candidates=np.array([v,[v[0]+np.pi,np.pi-v[1],v[2]+np.pi]])
  candidates+=2*np.pi*np.round((prev-candidates)/(2*np.pi));v=candidates[np.argmin(np.linalg.norm(candidates-prev,axis=1))];rows.append(v);prev=v
 return np.array(rows)
for start in [0,27,54,60]:
 rotation=Slerp(t,R.from_euler('XYZ',q[:,start+3:start+6]))(np.clip(times,t[0],t[-1]));qi[:,start+3:start+6]=continuous_euler(rotation,q[0,start+3:start+6])
arrays['qpos']=qi;arrays['qvel']=np.gradient(qi,.0025,axis=0);arrays['qvel'][0]=0
import mujoco
m=mujoco.MjModel.from_xml_path(str(T.parent/'scene_act.xml'));arrays['ctrl']=qi[:,m.jnt_qposadr[m.actuator_trnid[:,0]]]
np.savez_compressed(T/'trajectory_kinematic_act.npz',**arrays)
print('SO3 interpolation done',qi.shape)