"""Check and project all 400 Hz interpolated hand/object states, keeping objects fixed."""
import os
os.environ['OPENBLAS_NUM_THREADS']='1';os.environ['OMP_NUM_THREADS']='1'
from pathlib import Path
import json,sys,numpy as np,mujoco
from scipy.optimize import minimize
import l20_model_source as source
R=Path(__file__).resolve().parents[1];O=Path(os.environ.get('HF_FIX_OUT',R/'output/collision_fix_20260915'));T=O/'datasets/processed/current/l20/bimanual/boxes';OLD=Path(os.environ.get('HF_FIX_OLD',R/'output/foundationpose_spider_20260915'));source.REPO_ROOT=R/'third_party/l20_assets';m=mujoco.MjModel.from_xml_path(str(T/'scene_act.xml'));d=mujoco.MjData(m);video_mode='--video' in sys.argv;clearance=.0025 if '--buffer' in sys.argv else 0.;r=np.load(O/'reference_video_rate.npz' if video_mode else T/'0/trajectory_kinematic_act.npz');arrays={k:r[k] for k in r.files};q=arrays['qpos'].copy();original=q.copy();meta={}
for side in ['right','left']:
 k=source.HandKinematics(OLD/f'model_{side}/l20_{side}.xml',side);wa=np.array([m.jnt_qposadr[m.joint(side+'_hand_'+s).id] for s in ['pos_x','pos_y','pos_z','rot_x','rot_y','rot_z']]);ja=np.array([m.jnt_qposadr[m.joint(side+'_'+n).id] for n in k.joint_names]);B=np.zeros((m.nv,22));B[wa,:6]=np.eye(6);B[ja,6:]=k.expansion;meta[side]=(k,wa,ja,B)
def contacts(side=None,B=None):
 ds=[];rows=[]
 for c in d.contact:
  g0,g1=map(int,c.geom)
  if {int(m.geom_contype[g0]),int(m.geom_contype[g1])} not in ([{1,2},{1,4}] if os.environ.get('HF_TABLE','0')=='1' else [{1,2}]):continue
  hg=g0 if m.geom_contype[g0]==1 else g1
  if side is not None and not m.geom(hg).name.startswith(side+'_'):continue
  ds.append(float(c.dist))
  if B is not None:
   J=np.zeros((3,m.nv));mujoco.mj_jac(m,d,J,None,c.pos,int(m.geom_bodyid[hg]));rows.append((c.frame[:3]*(1 if hg==g1 else -1))@J@B)
 return np.asarray(rows).reshape(-1,22),np.asarray(ds)
pre=[];post=[];changed=[]
for f in range(len(q)):
 d.qpos[:]=q[f];mujoco.mj_fwdPosition(m,d);_,ds=contacts();pen=max(0,clearance-ds.min()) if len(ds) else 0;pre.append(pen)
 if pen>.00005:
  for side,(k,wa,ja,B) in meta.items():
   for it in range(15):
    A,ds=contacts(side,B)
    if not len(ds) or ds.min()>clearance-.00002:break
    x=np.r_[d.qpos[wa],d.qpos[ja][k.independent_indices]];scale=np.r_[[.01]*3,[.1]*3,[.1]*16];lo=np.maximum(-scale,np.r_[[-np.inf]*6,k.lower]-x);hi=np.minimum(scale,np.r_[[np.inf]*6,k.upper]-x);H=np.diag(1/scale**2)
    fit=minimize(lambda z:.5*z@H@z,np.zeros(22),jac=lambda z:H@z,bounds=list(zip(lo,hi)),constraints=[dict(type='ineq',fun=lambda z:A@z+ds-clearance-.0003,jac=lambda z:A)],method='SLSQP',options={'ftol':1e-9,'maxiter':80})
    old=d.qpos.copy();best=None
    for alpha in [1.,.5,.25]:
     y=x+alpha*fit.x;d.qpos[:]=old;d.qpos[wa]=y[:6];d.qpos[ja]=k.expand(y[6:]);mujoco.mj_fwdPosition(m,d);_,test=contacts(side);p=max(0,clearance-test.min()) if len(test) else 0
     if best is None or p<best[0]:best=(p,d.qpos.copy())
    d.qpos[:]=best[1];mujoco.mj_fwdPosition(m,d)
  q[f]=d.qpos.copy();changed.append(f)
 _,ds=contacts();post.append(max(0,-ds.min()) if len(ds) else 0)
 if f%800==0:print('interpolation',f,'max before/after mm',max(pre)*1000,max(post)*1000,flush=True)
arrays['qpos']=q;arrays['qvel']=np.gradient(q,1/30 if video_mode else .0025,axis=0);arrays['qvel'][0]=0;arrays['ctrl']=q[:,m.jnt_qposadr[m.actuator_trnid[:,0]]]
np.savez_compressed(O/('video_before_projection.npz' if video_mode else 'interpolation_before_projection.npz'),qpos=original,time=arrays['time']);np.savez_compressed(O/'reference_video_rate.npz' if video_mode else T/'0/trajectory_kinematic_act.npz',**arrays)
report={'requested_safety_clearance_mm':clearance*1000,'before_metric':'clearance shortfall' if clearance else 'penetration','steps':len(q),'max_penetration_before_mm':float(max(pre)*1000),'max_penetration_after_mm':float(max(post)*1000),'projected_steps':len(changed),'all_finite':bool(np.isfinite(q).all()),'object_coordinates_unchanged':bool(np.array_equal(q[:,-12:],original[:,-12:])),'max_additional_wrist_translation_mm':float(max(np.linalg.norm(q[:,:3]-original[:,:3],axis=1).max(),np.linalg.norm(q[:,27:30]-original[:,27:30],axis=1).max())*1000),'clearance_gate_0_5mm':bool(max(post)<.0005)}
(O/('video_rate_validation.json' if video_mode else 'interpolation_validation.json')).write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2),flush=True)

if not report['clearance_gate_0_5mm']:raise RuntimeError('Collision projection failed; inspect saved validation before running physics')