"""FoundationPose RGB-D tracking for both CAD parts, preserving camera/world transforms."""
import os
os.environ.setdefault('OMP_NUM_THREADS','4')
import sys,json,time,subprocess
from pathlib import Path
ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT/'third_party/FoundationPose'))
import cv2,numpy as np,torch,trimesh,nvdiffrast.torch as dr
from estimater import FoundationPose,ScorePredictor,PoseRefinePredictor
from Utils import nvdiffrast_render,make_mesh_tensors
OUT=ROOT/'output/foundationpose_spider_20260915';SRC=ROOT/'docs/20260915_171525';m=json.loads((SRC/'intrinsics.json').read_text());N=m['frames'];K=np.array([[m['fx'],0,m['cx']],[0,m['fy'],m['cy']],[0,0,1.]])
cam=np.load(ROOT/'output/20260915_171525_dynhamr/rgbd_camera.npz');c=cv2.VideoCapture(str(SRC/'color.mp4'));frames=[]
while True:
 ok,b=c.read()
 if not ok:break
 frames.append(b)
assert len(frames)==N
ctx=dr.RasterizeCudaContext();score=ScorePredictor();refine=PoseRefinePredictor();arrays={'K':K,'time':np.array(m['timestamps_ms'])/1000-m['timestamps_ms'][0]/1000,'c2w':cam['c2w'],'camera_valid':cam['valid']};meshes={};allmetrics={}
for name,file,init,color in [('upper','上半.stl',0,[210,35,35,255]),('lower','下半.stl',80,[35,65,210,255])]:
 mesh=trimesh.load(ROOT/'docs'/file);mesh.visual.vertex_colors=np.tile(color,(len(mesh.vertices),1));meshes[name]=mesh
 est=FoundationPose(model_pts=mesh.vertices,model_normals=mesh.vertex_normals,mesh=mesh,scorer=score,refiner=refine,glctx=ctx,debug=0,debug_dir=str(OUT/name/'debug'));mt=make_mesh_tensors(mesh);P=np.full((N,4,4),np.nan);valid=np.zeros(N,bool);metrics=[None]*N;start_state=None
 for i in list(range(init,N))+list(range(init-1,-1,-1)):
  b=frames[i];z=cv2.imread(str(SRC/'depth'/f'{i:06d}.png'),-1).astype('float32')*m['depth_scale_m'];z[(z<.1)|(z>1.)]=0
  h=cv2.cvtColor(b,cv2.COLOR_BGR2HSV);hue=h[:,:,0];mask=((((hue<12)|(hue>170)) if name=='upper' else ((hue>95)&(hue<135)))&(h[:,:,1]>90)&(h[:,:,2]>40)&(z>.1)&(z<1.)).astype('uint8');mask[:200]=0
  num,lab,st,_=cv2.connectedComponentsWithStats(mask,8)
  if num>1:mask=(lab==1+np.argmax(st[1:,4])).astype('uint8')
  if i==init-1:est.pose_last=start_state.clone()
  rgb=cv2.cvtColor(b,cv2.COLOR_BGR2RGB)
  if i==init:
   cv2.imwrite(str(OUT/name/'initial_mask.png'),mask*255);p=est.register(K=K,rgb=rgb,depth=z,ob_mask=mask,iteration=5);start_state=est.pose_last.clone()
  else:p=est.track_one(rgb=rgb,depth=z,K=K,iteration=2)
  P[i]=p
  with torch.inference_mode():_,depth,_=nvdiffrast_render(K=K,H=480,W=848,ob_in_cams=torch.as_tensor(p[None],device='cuda'),glctx=ctx,mesh_tensors=mt)
  depth=depth[0].cpu().numpy();over=(depth>0)&(mask>0)&(z>0);coverage=float(over.sum()/max(1,mask.sum()));err=float(np.median(np.abs(depth[over]-z[over]))) if over.any() else 1.
  valid[i]=mask.sum()>500 and coverage>.65 and err<.02
  # The lower part becomes occluded during assembly; independent lower tracking is not reliable there.
  if name=='lower' and i>=177:valid[i]=False
  metrics[i]={'frame':i,'valid':bool(valid[i]),'visible_color_pixels':int(mask.sum()),'coverage':coverage,'depth_residual_m':err}
  if i%60==0:print(name,metrics[i],flush=True)
 arrays[name+'_T_camera']=P;arrays[name+'_T_world']=cam['c2w']@P;arrays[name+'_valid']=valid;allmetrics[name]=metrics
np.savez_compressed(OUT/'foundationpose_objects.npz',**arrays);(OUT/'object_metrics.json').write_text(json.dumps(allmetrics,indent=2))
writer=subprocess.Popen(['ffmpeg','-y','-v','error','-f','rawvideo','-pix_fmt','bgr24','-s','848x480','-r','30','-i','-','-an','-c:v','libx264','-crf','20','-pix_fmt','yuv420p',str(OUT/'foundationpose_overlay.mp4')],stdin=subprocess.PIPE)
for i,b in enumerate(frames):
 pic=b.copy()
 for name,col in [('upper',[0,220,0]),('lower',[220,180,0])]:
  with torch.inference_mode():_,d,_=nvdiffrast_render(K=K,H=480,W=848,ob_in_cams=torch.tensor(arrays[name+'_T_camera'][i:i+1],device='cuda',dtype=torch.float32),glctx=ctx,mesh_tensors=make_mesh_tensors(meshes[name]))
  mask=d[0].cpu().numpy()>0
  if arrays[name+'_valid'][i]:pic[mask]=(pic[mask]*.65+np.array(col)*.35).astype('uint8')
 cv2.putText(pic,f'{i:03d} red:{arrays["upper_valid"][i]} blue:{arrays["lower_valid"][i]}',(8,25),0,.55,(0,255,255),2)
 writer.stdin.write(pic.tobytes())
 if i in [0,40,80,120,160,200,240,280,351]:cv2.imwrite(str(OUT/f'objects_{i:04d}.jpg'),pic)
writer.stdin.close();assert writer.wait()==0
print('FOUNDATIONPOSE_DONE', {n:int(arrays[n+'_valid'].sum()) for n in meshes},flush=True)