"""Build the two-hand/two-object SPIDER task from recorded, calibrated tracks."""
import os
os.environ.setdefault('MUJOCO_GL','osmesa')
from pathlib import Path
import copy,json,xml.etree.ElementTree as E
import numpy as np,cv2,mujoco,trimesh,open3d as o3d
from scipy.spatial.transform import Rotation as R,Slerp
from scipy.optimize import least_squares
import l20_model_source as source
ROOT=Path(__file__).resolve().parents[1];OUT=Path(os.environ.get('HF_SPIDER_OUT',ROOT/'output/foundationpose_spider_20260915'));BASE=Path(os.environ.get('HF_HAND_BASE',ROOT/'output/20260915_171525_dynhamr'))
TASK=OUT/'datasets/processed/current/l20/bimanual/boxes';TRIAL=TASK/'0';TRIAL.mkdir(parents=True,exist_ok=True)
source.REPO_ROOT=ROOT/'third_party/l20_assets';fingers=source.FINGERS;fp=np.load(OUT/'foundationpose_objects.npz');N=len(fp['time']);t=np.arange(N)/30
# Fit the visible table, excluding saturated foreground; z-up is table normal, not IMU gravity.
cap=cv2.VideoCapture(str(ROOT/'docs/20260915_171525/color.mp4'));ok,img=cap.read();cap.release();hsv=cv2.cvtColor(img,cv2.COLOR_BGR2HSV)
dep=cv2.imread(str(ROOT/'docs/20260915_171525/depth/000000.png'),-1)*.0001;yy,xx=np.indices(dep.shape);K=fp['K']
mask=(xx>100)&(xx<780)&(yy>200)&(yy<460)&(hsv[:,:,1]<50)&(dep>.2)&(dep<1)
pts=np.stack([(xx-K[0,2])*dep/K[0,0],(yy-K[1,2])*dep/K[1,1],dep],-1)[mask][::3]
pc=o3d.geometry.PointCloud(o3d.utility.Vector3dVector(pts));plane,ii=pc.segment_plane(.004,3,1000);normal=np.array(plane[:3]);offset=plane[3]
if normal[2]>0:normal=-normal;offset=-offset
S=R.align_vectors([[0,0,1]],[normal])[0].as_matrix();shift=np.array([0.,0.,offset]);world_from_source=np.eye(4);world_from_source[:3,:3]=S;world_from_source[:3,3]=shift
# Interpolate only upper rejected intervals. Lower occlusion gets an explicit, separate hypothesis.
poses={};provenance={}
for name in ['upper','lower']:
 raw=fp[name+'_T_world'].copy();valid=fp[name+'_valid'].astype(bool);idx=np.flatnonzero(valid)
 out=raw.copy();clip=np.clip(t,t[idx[0]],t[idx[-1]])
 out[:,:3,3]=np.stack([np.interp(t,t[idx],raw[idx,k,3]) for k in range(3)],1)
 out[:,:3,:3]=Slerp(t[idx],R.from_matrix(raw[idx,:3,:3]))(clip).as_matrix()
 if name=='lower':
  # Hold last visible world pose after frame176. No fabricated assembled pose.
  provenance[name]='Observed through frame176; held last valid world pose during occlusion, not a measurement.'
 else:provenance[name]='Rejected frames interpolated between accepted poses; endpoints held.'
 poses[name]=world_from_source[None]@out
np.savez_compressed(OUT/'object_reference.npz',**{k:v for k,v in poses.items()},source_world_to_sim=world_from_source,time=t,upper_observed=fp['upper_observed_raw'] if 'upper_observed_raw' in fp else fp['upper_valid'],lower_observed=fp['lower_observed_raw'] if 'lower_observed_raw' in fp else fp['lower_valid'])
root=None;handmeta={}
for side in ['right','left']:
 build=OUT/('model_'+side);path=source.build_model(side,build);rr=E.parse(path).getroot()
 for mesh in rr.findall('asset/mesh'):mesh.set('file',str((build/mesh.get('file')).resolve()));mesh.set('maxhullvert','32')
 # Prefix all named entities and their references before combining.
 for e in rr.iter():
  for a in ['name','mesh','material','joint1','joint2']:
   if e.get(a):e.set(a,side+'_'+e.get(a))
 hand=rr.find('worldbody/body');wn=[]
 for i in range(6):
  n=f'{side}_hand_'+('pos_'+'xyz'[i] if i<3 else 'rot_'+'xyz'[i-3]);wn.append(n)
  hand.insert(i,E.Element('joint',name=n,type='slide' if i<3 else 'hinge',axis=['1 0 0','0 1 0','0 0 1'][i%3],limited='false',damping='0',armature='.02'))
 for g in hand.iter('geom'):g.attrib.update(contype='1',conaffinity='2',friction='.8 .005 .001',condim='3',solref='.015 1')
 for j,f in enumerate(fingers):hand.find(f".//site[@name='{side}_landmark_{4+4*j:02d}']").set('name',f'{side}_hand_{f}_track')
 kin=source.HandKinematics(path,side);handmeta[side]=(kin,wn)
 if root is None:
  root=rr;world=root.find('worldbody');assets=root.find('asset');act=E.SubElement(root,'actuator')
  floor=world.find("geom[@name='right_floor']");floor.attrib.update(pos='0 0 0',size='2 2 .01',contype='4',conaffinity='2',rgba='.6 .62 .65 1')
 else:
  assets.extend(list(rr.find('asset')));world.append(hand);root.find('equality').extend(list(rr.find('equality')))
 for i,n in enumerate(wn):E.SubElement(act,'position',name=n+'_act',joint=n,kp='2000' if i<3 else '40',kv='60' if i<3 else '2',forcerange='-200 200' if i<3 else '-20 20')
 for n,lo,hi in zip(kin.independent_joint_names,kin.lower,kin.upper):E.SubElement(act,'position',name=side+'_'+n+'_act',joint=side+'_'+n,kp='8',kv='.15',forcerange='-1 1',ctrlrange=source._fmt([lo,hi]))
root.find('option').attrib.update(timestep='.0025',gravity='0 0 -9.81',integrator='implicitfast',cone='elliptic')
root.find('default/joint').attrib.update(damping='.05',armature='.002')
for eq in root.findall('equality/joint'):eq.set('solref','.01 1')
meshes={};scenes={}
for side,name,stl,color in [('right','upper','上半.stl','.85 .08 .06 1'),('left','lower','下半.stl','.06 .18 .8 1')]:
 mesh=trimesh.load(ROOT/'docs'/stl);meshes[name]=mesh;assets.append(E.Element('mesh',name=name+'_visual',file=str(ROOT/'docs'/stl)))
 obj=E.SubElement(world,'body',name=side+'_object')
 # Mass is an explicit simulation assumption, not a sensor measurement.
 E.SubElement(obj,'inertial',pos=source._fmt(mesh.center_mass),mass='.12',diaginertia='.0004 .0006 .0004')
 for i in range(6):
  n=side+'_object_'+('pos_'+'xyz'[i] if i<3 else 'rot_'+'xyz'[i-3]);E.SubElement(obj,'joint',name=n,type='slide' if i<3 else 'hinge',axis=['1 0 0','0 1 0','0 0 1'][i%3],limited='false',damping='0.001',armature='.001');E.SubElement(act,'position',name=n,joint=n,kp='0',kv='0')
 E.SubElement(obj,'geom',name=name+'_visual',type='mesh',mesh=name+'_visual',rgba=color,contype='0',conaffinity='0',mass='0',group='1')
 for i,p in enumerate(sorted((OUT/'collision').glob(name+'_*.obj'))):
  mn=f'{name}_collision_{i}';E.SubElement(assets,'mesh',name=mn,file=str(p),maxhullvert='32');E.SubElement(obj,'geom',name=mn,type='mesh',mesh=mn,contype='2',conaffinity='7',friction='.8 .005 .001',condim='3',solref='.015 1',mass='0',group='3')
 E.SubElement(obj,'site',name=side+'_object_track',pos=source._fmt(mesh.centroid))
 sc=o3d.t.geometry.RaycastingScene();sc.add_triangles(o3d.t.geometry.TriangleMesh.from_legacy(o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(mesh.vertices),o3d.utility.Vector3iVector(mesh.faces))));scenes[name]=sc
E.SubElement(world,'camera',name='source_camera',fovy=str(np.degrees(2*np.arctan(480/(2*K[1,1])))))
scene=TASK/'scene_act.xml';E.ElementTree(root).write(scene)
m=mujoco.MjModel.from_xml_path(str(scene));d=mujoco.MjData(m);q=np.zeros((N,m.nq));contact=np.zeros((N,10));cp=np.zeros((N,10,3));fit_report={};initialq=None
for si,(side,name) in enumerate([('right','upper'),('left','lower')]):
 kin,wn=handmeta[side];motion=np.load(BASE/f'l20_{side}_stable/motion.npz');human=np.load(BASE/f'human_joints_{side}.npz')
 wp=(motion['wrist_pos']-motion['scene_translation'])@motion['scene_rotation'];wr=np.einsum('ij,njk->nik',motion['scene_rotation'].T,R.from_quat(motion['wrist_quat_wxyz'][:,[1,2,3,0]]).as_matrix())
 wp=wp@S.T+shift;wr=np.einsum('ij,njk->nik',S,wr);angles=np.unwrap(R.from_matrix(wr).as_euler('XYZ'),axis=0)
 wa=[m.jnt_qposadr[m.joint(n).id] for n in wn];ja=[m.jnt_qposadr[m.joint(side+'_'+n).id] for n in kin.joint_names]
 full=motion['qpos'][:,[list(motion['joint_names']).index(n) for n in kin.joint_names]];a=full[:,kin.independent_indices];a=np.clip(a,kin.lower,kin.upper)
 q[:,wa]=np.c_[wp,angles];q[:,ja]=kin.expand(a)
 oa=[m.jnt_qposadr[m.joint(side+'_object_'+s).id] for s in ['pos_x','pos_y','pos_z','rot_x','rot_y','rot_z']];q[:,oa]=np.c_[poses[name][:,:3,3],np.unwrap(R.from_matrix(poses[name][:,:3,:3]).as_euler('XYZ'),axis=0)]
 # Human fingertips yield geometric contact hypotheses, independently of retargeted hand error.
 hp=(human['joints']+human['wrist_world'][:,None,:])@S.T+shift
 candidates=[];surface=[]
 for candidate in ['upper','lower']:
  local=np.einsum('nji,nkj->nki',poses[candidate][:,:3,:3],hp[:,[4,8,12,16,20]]-poses[candidate][:,None,:3,3])
  nearest=scenes[candidate].compute_closest_points(o3d.core.Tensor(local.astype(np.float32).reshape(-1,3)))['points'].numpy().reshape(N,5,3)
  distance=np.linalg.norm(nearest-local,axis=-1)
  if candidate=='lower':distance[177:]=100.
  candidates.append(distance);surface.append(np.einsum('nij,nkj->nki',poses[candidate][:,:3,:3],nearest)+poses[candidate][:,None,:3,3])
 choose=np.argmin(candidates,axis=0);dist=np.min(candidates,axis=0);on=np.zeros(5,dtype=bool)
 for f in range(N):
  on=np.where(on,dist[f]<.035,dist[f]<.025);contact[f,si*5:si*5+5]=on
  for tip in range(5):cp[f,si*5+tip]=surface[choose[f,tip]][f,tip]
 fit_report[side]={'contact_frames':int(np.any(contact[:,si*5:si*5+5],axis=1).sum()),'human_tip_distance_median_mm':float(np.median(dist)*1000)}
np.savez_compressed(OUT/'reference_before_contact_fit.npz',qpos=q,time=t)
# Bounded kinematic warm start; subsequent optimization is actual upstream SPIDER.
for si,(side,name) in enumerate([('right','upper'),('left','lower')]):
 kin,wn=handmeta[side];wa=[m.jnt_qposadr[m.joint(n).id] for n in wn];ja=[m.jnt_qposadr[m.joint(side+'_'+n).id] for n in kin.joint_names];sid=[m.site(f'{side}_hand_{f}_track').id for f in fingers];errors=[]
 handgeom=[g for g in range(m.ngeom) if m.geom_contype[g]==1 and m.geom(g).name.startswith(side+'_')]
 for f in range(N):
  active=contact[f,si*5:si*5+5].astype(bool)
  old=q[f].copy();base=old[ja][kin.independent_indices];target=cp[f,si*5:si*5+5]
  def residual(x):
   d.qpos[:]=old;d.qpos[wa[:3]]=old[wa[:3]]+x[16:];d.qpos[ja]=kin.expand(x[:16]);mujoco.mj_fwdPosition(m,d)
   penetrations=np.zeros(len(handgeom));lookup={g:k for k,g in enumerate(handgeom)}
   for cc in d.contact:
    g0,g1=map(int,cc.geom)
    if g0 in lookup and m.geom_contype[g1]==2:penetrations[lookup[g0]]=max(penetrations[lookup[g0]],-cc.dist-.001)
    if g1 in lookup and m.geom_contype[g0]==2:penetrations[lookup[g1]]=max(penetrations[lookup[g1]],-cc.dist-.001)
   return np.r_[penetrations/.002,(d.site_xpos[sid][active]-target[active]).ravel()/.008,.15*(x[:16]-base),.3*x[16:]/.03]
  fit=least_squares(residual,np.r_[base,np.zeros(3)],bounds=(np.r_[kin.lower,[-.05]*3],np.r_[kin.upper,[.05]*3]),max_nfev=30,diff_step=1e-4)
  residual(fit.x);q[f]=d.qpos;
  if active.any():errors.append(np.linalg.norm(d.site_xpos[sid][active]-target[active],axis=-1).mean()*1000)
  if f%80==0:print('warmstart',side,f,flush=True)
 fit_report[side]['warmstart_contact_gap_mean_mm']=float(np.mean(errors)) if errors else None
 print('fit',side,fit_report[side],flush=True)
# Interpolate unwrapped scalar coordinates, with exact mimic preserved by linear interpolation.
dt=.0025;times=np.arange(int(np.ceil(t[-1]/.1)*.1/dt)+1)*dt
qi=np.stack([np.interp(times,t,q[:,j]) for j in range(m.nq)],1);ci=np.stack([np.interp(times,t,cp.reshape(N,-1)[:,j]) for j in range(30)],1).reshape(-1,10,3)
nearestidx=np.minimum(np.searchsorted(t,times),N-1);ct=contact[nearestidx];vel=np.gradient(qi,dt,axis=0);vel[0]=0
ctrl=qi[:,m.jnt_qposadr[m.actuator_trnid[:,0]]];siteids=[m.site(f'{s}_hand_{f}_track').id for s in ['right','left'] for f in fingers]
np.savez_compressed(TRIAL/'trajectory_kinematic_act.npz',qpos=qi,qvel=vel,ctrl=ctrl,contact=ct,contact_pos=ci,time=times)
np.savez_compressed(OUT/'reference_video_rate.npz',qpos=q,contact=contact,contact_pos=cp,time=t,camera_world_from_cv=world_from_source[None]@fp['c2w'])
(TASK/'task_info.json').write_text(json.dumps({'ref_dt':dt,'contact_site_ids':siteids},indent=2))
c=json.loads((ROOT/'output/spider_l20_contact/config.json').read_text());c.update(dataset_dir=str(OUT/'datasets'),task='boxes',embodiment_type='bimanual',object_action_dims=12,num_samples=96,max_num_iterations=3,horizon=.4,max_sim_steps=len(times)-1,nconmax_per_env=512,njmax_per_env=1536,contact_rew_scale=2.,pos_rew_scale=10.)
(OUT/'config.json').write_text(json.dumps(c,indent=2));report=dict(nq=m.nq,nv=m.nv,nu=m.nu,frames=N,duration_s=float(times[-1]),table_plane=np.asarray(plane).tolist(),table_inliers=len(ii),source_world_to_sim=world_from_source.tolist(),object_provenance=provenance,contact=fit_report,contact_rule='Human fingertip to CAD distance: engage25mm release35mm; nearest observed object per fingertip; lower excluded after176. Geometric hypothesis, no tactile labels.',mass_assumption_kg=.12,friction_assumption=.8,collision='20 CoACD convex parts per object; approximate concavity',warmstart='Bounded finger+translation least-squares with collision penalty; not physics and not SPIDER',object_joint_armature=.001,object_real_actuator_gains_zero=bool(not m.actuator_gainprm[-12:].any()))
(OUT/'preparation.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2),flush=True)