"""Collision-constrained L20 retargeting: analytic contact Jacobians, exact mimic, verified steps."""
import os
os.environ.setdefault('MUJOCO_GL','osmesa')
os.environ['OPENBLAS_NUM_THREADS']='1';os.environ['OMP_NUM_THREADS']='1'
from pathlib import Path
import json,copy,argparse,xml.etree.ElementTree as E
import numpy as np,mujoco,trimesh
from scipy.optimize import minimize
from scipy.spatial.transform import Rotation as Rotation
import l20_model_source as source
R=Path(__file__).resolve().parents[1];OLD=Path(os.environ.get('HF_FIX_OLD',R/'output/foundationpose_spider_20260915'));OUT=Path(os.environ.get('HF_FIX_OUT',R/'output/collision_fix_20260915'));T=OUT/'datasets/processed/current/l20/bimanual/boxes';(T/'0').mkdir(parents=True,exist_ok=True)
p=argparse.ArgumentParser();p.add_argument('--legacy-geoms',action='store_true');p.add_argument('--pilot',action='store_true');args=p.parse_args()
root=E.parse(OLD/'datasets/processed/current/l20/bimanual/boxes/scene_act.xml').getroot();assets=root.find('asset')
if not args.legacy_geoms:
 manifest=json.loads((OUT/'collision_v2/manifest.json').read_text());handmanifest=json.loads((OUT/'hand_collision/manifest.json').read_text());assert len(manifest)==2 and len(handmanifest)==4
 for name,side in [('upper','right'),('lower','left')]:
  body=root.find(f"worldbody/body[@name='{side}_object']")
  for g in list(body):
   if g.tag=='geom' and '_collision_' in g.get('name',''):body.remove(g)
  for a in list(assets):
   if a.get('name','').startswith(name+'_collision_'):assets.remove(a)
  for i,path in enumerate(sorted((OUT/'collision_v2').glob(name+'_*.obj'))):
   part=trimesh.load(path)
   if abs(part.volume)<1e-12 or np.linalg.matrix_rank(part.vertices-part.vertices.mean(0),tol=1e-9)<3:continue
   n=f'{name}_collision_{i}';E.SubElement(assets,'mesh',name=n,file=str(path),maxhullvert='128');E.SubElement(body,'geom',name=n,type='mesh',mesh=n,contype='2',conaffinity='7',friction='.8 .005 .001',condim='3',solref='.006 1',solimp='.95 .99 .001',mass='0',group='3')
 for body in root.findall('.//body'):
  for g in list(body.findall('geom')):
   mesh=g.get('mesh')
   if mesh not in handmanifest:continue
   g.set('contype','0');g.set('conaffinity','0')
   for i,path in enumerate(sorted((OUT/'hand_collision').glob(mesh+'_*.obj'))):
    n=f'{mesh}_part_{i}';E.SubElement(assets,'mesh',name=n,file=str(path),maxhullvert='128');attrs=dict(g.attrib);attrs.update(name=n,mesh=n,contype='1',conaffinity='2',group='3',mass='0',solref='.006 1',solimp='.95 .99 .001');E.SubElement(body,'geom',**attrs)
TABLE=os.environ.get('HF_TABLE','0')=='1';PAIRS=[{1,2},{1,4}] if TABLE else [{1,2}]
if TABLE:
 for fl in root.findall(".//geom[@name='right_floor']"):fl.set('conaffinity','3')   # hands (contype 1) now collide with the table
# Keep contacts within 4 mm available to the solver, without early physical forces.
for g in root.findall('.//geom'):
 if int(g.get('contype','0')) in ([1,2,4] if TABLE else [1,2]):g.attrib.update(margin='.004',gap='.004',solref='.005 1',solimp='.95 .99 .001')
scene=T/('scene_legacy.xml' if args.legacy_geoms else 'scene_act.xml');E.ElementTree(root).write(scene)
m=mujoco.MjModel.from_xml_path(str(scene));d=mujoco.MjData(m);source.REPO_ROOT=R/'third_party/l20_assets'
oldref=np.load(OLD/'reference_video_rate.npz');original=np.load(OLD/'reference_before_contact_fit.npz')['qpos'];q=oldref['qpos'].copy();cp=oldref['contact_pos'];contact=oldref['contact'];N=len(q);meta={}
for side in ['right','left']:
 k=source.HandKinematics(OLD/('model_'+side)/('l20_'+side+'.xml'),side);wn=[side+'_hand_'+s for s in ['pos_x','pos_y','pos_z','rot_x','rot_y','rot_z']];wa=np.array([m.jnt_qposadr[m.joint(n).id] for n in wn]);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,[m.site(side+'_hand_'+f+'_track').id for f in source.FINGERS])
# All hand/object geometry is used, rather than only tip points.
def collision_rows(side,B,derivatives=True):
 rows=[];dist=[]
 for c in d.contact:
  g0,g1=map(int,c.geom);types={int(m.geom_contype[g0]),int(m.geom_contype[g1])}
  if types not in PAIRS:continue
  hg=g0 if m.geom_contype[g0]==1 else g1
  if not m.geom(hg).name.startswith(side+'_'):continue
  dist.append(float(c.dist))
  if derivatives:
   jac=np.zeros((3,m.nv));mujoco.mj_jac(m,d,jac,None,c.pos,int(m.geom_bodyid[hg]));normal=c.frame[:3]*(1 if hg==g1 else -1);rows.append(normal@jac@B)
 return np.asarray(rows).reshape(-1,22),np.array(dist)
def put(x,base,k,wa,ja):
 d.qpos[:]=base;d.qpos[wa]=x[:6];d.qpos[ja]=k.expand(x[6:]);mujoco.mj_fwdPosition(m,d)
frames=[0,36,46,80,160,240,351] if args.pilot else range(N);report=[];previous={};margin=float(os.environ.get('HF_MARGIN','.0005'))
for f in frames:
 for si,side in enumerate(['right','left']):
  k,wa,ja,B,sites=meta[side];base=q[f].copy();xorig=np.r_[original[f,wa],original[f,ja][k.independent_indices]];x=np.r_[q[f,wa],q[f,ja][k.independent_indices]]
  low=np.r_[xorig[:3]-.15,xorig[3:6]-.65,k.lower];high=np.r_[xorig[:3]+.15,xorig[3:6]+.65,k.upper];x=np.clip(x,low,high)
  target=cp[f,si*5:si*5+5];on=contact[f,si*5:si*5+5].astype(bool);prior=xorig.copy()
  if side in previous and not args.pilot:
   prior=np.clip(xorig+previous[side]*(1. if on.any() else .95),low,high);x=prior.copy()
   prev=np.r_[q[f-1,wa],q[f-1,ja][k.independent_indices]]
   temporal=np.r_[[.012]*3,[.18]*3,[.25]*16];low=np.maximum(low,prev-temporal);high=np.minimum(high,prev+temporal)
   # Transport the previous hand with its currently contacted object to stay on the same side.
   if on.any():
    d.qpos[:]=base;mujoco.mj_kinematics(m,d);center=target[on].mean(0);choices=[54] if f>=177 else [54,60]
    start=min(choices,key=lambda j:np.linalg.norm(center-d.site_xpos[m.site(('right' if j==54 else 'left')+'_object_track').id]))
    old_obj=q[f-1,start:start+6];new_obj=q[f,start:start+6];delta=Rotation.from_euler('XYZ',new_obj[3:])*Rotation.from_euler('XYZ',old_obj[3:]).inv()
    x[:3]=delta.apply(prev[:3]-old_obj[:3])+new_obj[:3]
    e=(delta*Rotation.from_euler('XYZ',prev[3:6])).as_euler('XYZ');candidates=np.array([e,[e[0]+np.pi,np.pi-e[1],e[2]+np.pi]]);candidates+=2*np.pi*np.round((prev[3:6]-candidates)/(2*np.pi));x[3:6]=candidates[np.argmin(np.linalg.norm(candidates-prev[3:6],axis=1))];x[6:]=prev[6:]
   x=np.clip(x,low,high)
  put(x,base,k,wa,ja);_,dist=collision_rows(side,B,False);initial=max(0,-dist.min()) if len(dist) else 0
  # Exit deep or contradictory overlap first; no object pose changes or mesh shrinking.
  if initial>.004 and (f==0 or args.pilot):
   seeds=[x.copy()];directions=np.array([[i,j,z] for i in [-1,0,1] for j in [-1,0,1] for z in [-1,0,1] if (i,j,z)!=(0,0,0)],float);directions/=np.linalg.norm(directions,axis=1)[:,None]
   for radius in [.025,.05,.08,.12]:
    for direction in directions:
     seed=x.copy();seed[:3]=np.clip(x[:3]+direction*radius,low[:3],high[:3]);seeds.append(seed)
   best=None
   for seed in seeds:
    put(seed,base,k,wa,ja);_,ds=collision_rows(side,B,False);pen=max(0,-ds.min()) if len(ds) else 0;gap=np.linalg.norm(d.site_xpos[sites][on]-target[on],axis=1).mean() if on.any() else 0
    score=500*pen+np.linalg.norm(seed[:3]-xorig[:3])+gap*.3
    if best is None or score<best[0]:best=(score,seed.copy())
   x=best[1]
  weights=np.r_[[float(os.environ.get('HF_WRIST_WEIGHT','25'))]*3,[1.5]*3,[.22]*16];bounds_step=np.r_[[.015]*3,[.10]*3,[.18]*16]
  for it in range(45):
   put(x,base,k,wa,ja);A,ds=collision_rows(side,B);J=[];res=[]
   for tip in np.flatnonzero(on):
    jac=np.zeros((3,m.nv));mujoco.mj_jacSite(m,d,jac,None,int(sites[tip]));J.extend((jac@B)/.012);res.extend((d.site_xpos[sites[tip]]-target[tip])/.012)
   J=np.asarray(J).reshape(-1,22);res=np.array(res);H=np.diag(weights**2)+J.T@J;g=weights**2*(x-prior)+J.T@res
   lo=np.maximum(low-x,-bounds_step);hi=np.minimum(high-x,bounds_step)
   constraints=[] if not len(ds) else [dict(type='ineq',fun=lambda z:A@z+ds-margin,jac=lambda z:A)]
   fit=minimize(lambda z:.5*z@H@z+g@z,np.zeros(22),jac=lambda z:H@z+g,bounds=list(zip(lo,hi)),constraints=constraints,method='SLSQP',options={'maxiter':80,'ftol':1e-9})
   dx=fit.x;pen=max(0.,-ds.min()) if len(ds) else 0.;oldcost=.5*np.sum((weights*(x-prior))**2)+.5*(res@res)+1e5*pen
   chosen=None
   for alpha in [1.,.5,.25,.1]:
    trial=np.clip(x+alpha*dx,low,high);put(trial,base,k,wa,ja);_,td=collision_rows(side,B,False);tp=max(0.,-td.min()) if len(td) else 0.;rr=(d.site_xpos[sites][on]-target[on]).ravel()/.012;cost=.5*np.sum((weights*(trial-prior))**2)+.5*(rr@rr)+1e5*tp
    if cost<oldcost-1e-7 or (pen>.0001 and tp<pen*.8):chosen=trial;break
   if chosen is None:break
   x=chosen
   if np.linalg.norm(alpha*dx)<1e-5:break
  put(x,base,k,wa,ja);_,ds=collision_rows(side,B,False);final=max(0.,-ds.min()) if len(ds) else 0
  q[f]=d.qpos.copy();previous[side]=x-xorig
  report.append(dict(frame=int(f),side=side,initial_penetration_mm=float(initial*1000),final_penetration_mm=float(final*1000),wrist_shift_mm=float(np.linalg.norm(x[:3]-xorig[:3])*1000),contact_gap_mm=float(np.linalg.norm(d.site_xpos[sites][on]-target[on],axis=1).mean()*1000) if on.any() else None,iterations=it+1))
 if f%20==0 or args.pilot:print('frame',f,report[-2:],flush=True)
 (OUT/('pilot_fit.json' if args.pilot else 'fit_progress.json')).write_text(json.dumps(report,indent=2))
if args.pilot:
 print('PILOT_FINISHED',max(x['final_penetration_mm'] for x in report),flush=True);raise SystemExit
np.savez_compressed(OUT/'reference_video_rate.npz',qpos=q,contact=contact,contact_pos=cp,time=oldref['time'],camera_world_from_cv=oldref['camera_world_from_cv'])
# Interpolated states are separately collision-checked before physical optimization.
times=np.load(OLD/'datasets/processed/current/l20/bimanual/boxes/0/trajectory_kinematic_act.npz')['time'];qi=np.stack([np.interp(times,oldref['time'],q[:,i]) for i in range(m.nq)],1);vel=np.gradient(qi,.0025,axis=0);vel[0]=0;ctrl=qi[:,m.jnt_qposadr[m.actuator_trnid[:,0]]];idx=np.minimum(np.searchsorted(oldref['time'],times),N-1);ci=np.stack([np.interp(times,oldref['time'],cp.reshape(N,-1)[:,j]) for j in range(30)],1).reshape(-1,10,3)
np.savez_compressed(T/'0/trajectory_kinematic_act.npz',qpos=qi,qvel=vel,ctrl=ctrl,contact=contact[idx],contact_pos=ci,time=times)
(T/'task_info.json').write_text((OLD/'datasets/processed/current/l20/bimanual/boxes/task_info.json').read_text());config=json.loads((OLD/'config.json').read_text());config.update(dataset_dir=str(OUT/'datasets'),num_samples=64,nconmax_per_env=1024,njmax_per_env=3072);(OUT/'config.json').write_text(json.dumps(config,indent=2));print('COLLISION_FIT_DONE',flush=True)