"""Isolate early CPU/GPU contact instability with identical saved controls."""
import os,json
from pathlib import Path
ROOT=Path(__file__).resolve().parents[1];O=ROOT/'output/spider_dynamics_fix_20260915';T=O/'datasets/processed/current/l20/bimanual/boxes'
os.environ['WARP_CACHE_PATH']=str(ROOT/'.spider_cache/warp');os.environ['OMP_NUM_THREADS']='2'
import numpy as np,mujoco,mujoco_warp as mw,warp as wp
r=np.load(T/'0/trajectory_kinematic_act.npz');rq,rv,rc=(r[k] for k in ['qpos','qvel','ctrl']);u=np.load(O/'hard_contact_attempt/control_replay.npz')['ctrl'][:80];wp.init();reports=[]
variants=[('hard',.999,.9999,.0001,.005,2),('medium',.95,.99,.001,.005,2),('default',.9,.95,.001,.005,2),('softtime',.9,.95,.001,.01,2),('cg',.95,.99,.001,.005,1),('tau20',.9,.95,.001,.02,2),('tau40',.9,.95,.001,.04,2),('imp90',.8,.9,.001,.01,2),('imp80',.7,.8,.001,.01,2),('buffer_tau10',.9,.95,.001,.01,2),('buffer_tau20',.9,.95,.001,.02,2)]
for name,dmin,dmax,width,tau,solver in variants:
 m=mujoco.MjModel.from_xml_path(str(O/'hard_contact_attempt/scene_act.xml'));m.opt.iterations=80;m.opt.ls_iterations=50;m.opt.solver=solver;m.opt.integrator=mujoco.mjtIntegrator.mjINT_IMPLICITFAST;m.geom_solimp[:]=[dmin,dmax,width,.5,2];m.geom_solref[:]=[tau,1]
 if name.startswith('buffer'):m.geom_gap[m.geom_contype>0]=.002
 d=mujoco.MjData(m);d.qpos[:]=rq[0];d.qvel[:]=rv[0];d.ctrl[:]=rc[0];mujoco.mj_step(m,d);cpu=[];gpu=[]
 with wp.ScopedDevice('cuda:0'):
  wm=mw.put_model(m);wd=mw.put_data(m,d,nworld=2,nconmax=1024,njmax=3072)
  with wp.ScopedCapture() as cap:mw.step(wm,wd)
  for ctrl in u:
   wd.ctrl.assign(np.tile(ctrl.astype(np.float32),(2,1)));wp.capture_launch(cap.graph);wp.synchronize();gpu.append(wd.qpos.numpy()[0].copy());d.ctrl[:]=ctrl;mujoco.mj_step(m,d);cpu.append(d.qpos.copy())
 gpu=np.array(gpu);cpu=np.array(cpu);pen=[];shift=[]
 for q in gpu:
  d.qpos[:]=q;mujoco.mj_fwdPosition(m,d);pen.append(max([max(0,-c.dist) for c in d.contact if {int(m.geom_contype[c.geom[0]]),int(m.geom_contype[c.geom[1]])}=={1,2}],default=0)*1000);shift.append([d.xipos[m.body(n).id].copy() for n in ['right_object','left_object']])
 shift=np.array(shift);report={'name':name,'max_cpu_gpu_qpos_difference':float(np.max(abs(gpu-cpu))),'max_penetration_mm':float(max(pen)),'max_object_com_displacement_mm':float(np.max(np.linalg.norm(shift-shift[0],axis=2))*1000),'finite':bool(np.isfinite(gpu).all())};reports.append(report);np.savez_compressed(O/('contact_probe_'+name+'.npz'),gpu=gpu,cpu=cpu);(O/'gpu_contact_probe.json').write_text(json.dumps(reports,indent=2));print(report,flush=True)