"""Refit CAD poses from observed outlines and robust depth planes; fixed CAD scale."""
import os
os.environ.setdefault('OMP_NUM_THREADS','4')
from pathlib import Path
import json,cv2,numpy as np,open3d as o3d
from scipy.spatial.transform import Rotation,Slerp
ROOT=Path(__file__).resolve().parents[1];BASE=ROOT/'output/20260915_171525_dynhamr';OUT=BASE/'object_refit';OUT.mkdir(exist_ok=True)
meta=json.loads((ROOT/'docs/20260915_171525/intrinsics.json').read_text());cam=np.load(BASE/'rgbd_camera.npz');old=np.load(BASE/'object_poses.npz');N=352
K=np.array([[meta['fx'],0,meta['cx']],[0,meta['fy'],meta['cy']],[0,0,1.]])
meshes={n:np.asarray(o3d.io.read_triangle_mesh(str(BASE/'object_assets'/f'{n}.stl')).vertices) for n in ['upper','lower']}
poses={n:[] for n in meshes};records={n:[] for n in meshes};cap=cv2.VideoCapture(str(ROOT/'docs/20260915_171525/color.mp4'))
def project(v,T):
 p=v@T[:3,:3].T+T[:3,3];return (p@K.T)[:,:2]/p[:,2,None]
def score(v,T,mask):
 uv=project(v,T); hull=cv2.convexHull(np.rint(uv).astype(np.int32)); pred=np.zeros_like(mask);cv2.fillConvexPoly(pred,hull,1)
 obs=np.zeros_like(mask);cs,_=cv2.findContours(mask,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE);cv2.fillConvexPoly(obs,cv2.convexHull(max(cs,key=cv2.contourArea)),1)
 return float(np.sum((pred>0)&(obs>0))/max(1,np.sum((pred>0)|(obs>0))))
for t in range(N):
 ok,im=cap.read();assert ok
 hsv=cv2.cvtColor(im,cv2.COLOR_BGR2HSV);dep=cv2.imread(str(ROOT/'docs/20260915_171525/depth'/f'{t:06d}.png'),-1)*meta['depth_scale_m']
 for n,v in meshes.items():
  hue=hsv[:,:,0]; mask=(((hue<12)|(hue>170)) if n=='upper' else ((hue>90)&(hue<135)))&(hsv[:,:,1]>90)&(hsv[:,:,2]>35)&(dep>.12)&(dep<.85);mask[:80]=False
  num,labels,stats,_=cv2.connectedComponentsWithStats(mask.astype('uint8'),8);label=1+np.argmax(stats[1:,4]);mask=(labels==label).astype('uint8')
  cs,_=cv2.findContours(mask,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE);hull=cv2.convexHull(max(cs,key=cv2.contourArea));poly=None
  for eps in [.015,.025,.04,.06]:
   q=cv2.approxPolyDP(hull,eps*cv2.arcLength(hull,True),True)
   if len(q)==4:poly=q[:,0].astype(float);break
  if poly is None:poly=cv2.boxPoints(cv2.minAreaRect(hull)).astype(float)
  center=poly.mean(0);poly=poly[np.argsort(np.arctan2(poly[:,1]-center[1],poly[:,0]-center[0]))];poly=np.roll(poly,-np.argmin(poly.sum(1)),axis=0)
  # Robust plane uses only colored measured pixels, never white insert/background hull fill.
  yy,xx=np.nonzero(cv2.erode(mask,np.ones((3,3),np.uint8)));z=dep[yy,xx];pts=np.c_[(xx-meta['cx'])/meta['fx']*z,(yy-meta['cy'])/meta['fy']*z,z]
  pc=o3d.geometry.PointCloud(o3d.utility.Vector3dVector(pts[::3]));plane,inliers=pc.segment_plane(.004,3,100);normal=np.array(plane[:3]);rays=np.c_[poly,np.ones(4)]@np.linalg.inv(K).T;corners=rays*(-plane[3]/(rays@normal))[:,None]
  lo=v.min(0);hi=v.max(0)
  if n=='upper':local=np.array([[lo[0],lo[1],hi[2]],[hi[0],lo[1],hi[2]],[hi[0],lo[1],lo[2]],[lo[0],lo[1],lo[2]]])
  else:local=np.array([[lo[0],-.0172,lo[2]],[hi[0],-.0172,lo[2]],[hi[0],-.0172,hi[2]],[lo[0],-.0172,hi[2]]])
  A=local-local.mean(0);B=corners-corners.mean(0);u,s,vt=np.linalg.svd(A.T@B);rr=vt.T@np.diag([1,1,np.linalg.det(vt.T@u.T)])@u.T
  T=np.eye(4);T[:3,:3]=rr;T[:3,3]=corners.mean(0)-rr@local.mean(0)
  before=np.eye(4);before[:3,:3]=Rotation.from_quat(old[n+'_quaternion_xyzw'][t]).as_matrix();before[:3,3]=old[n+'_position_world'][t];before=np.linalg.inv(cam['c2w'][t])@before
  candidates=[T]
  for shift in [0,2]:
   success,rv,tv=cv2.solvePnP(local.astype(float),np.roll(poly,shift,axis=0).astype(float),K,None,flags=cv2.SOLVEPNP_ITERATIVE)
   if success:
    candidate=np.eye(4);candidate[:3,:3]=Rotation.from_rotvec(rv.ravel()).as_matrix();candidate[:3,3]=tv.ravel();candidates.append(candidate)
  T=max(candidates,key=lambda candidate:score(v,candidate,mask)-2*abs((local@candidate[:3,:3].T+candidate[:3,3])[:,2].mean()-corners[:,2].mean()))
  rr=T[:3,:3]
  iou0=score(v,before,mask);iou1=score(v,T,mask);rmse=np.sqrt(np.mean(np.sum((local@rr.T+T[:3,3]-corners)**2,axis=1)))
  accepted=bool(iou1>.65 and rmse<.06 and len(inliers)>len(pts[::3])*.5)
  poses[n].append(cam['c2w'][t]@T);records[n].append(dict(frame=t,accepted=accepted,old_outline_iou=iou0,new_outline_iou=iou1,corner_fit_rmse_m=float(rmse)))
  if t in [0,176,351]:
   pic=im.copy()
   for transform,color in [(before,(0,200,255)),(T,(0,255,0))]:cv2.polylines(pic,[cv2.convexHull(np.rint(project(v,transform)).astype('int32'))],True,color,2)
   cv2.imwrite(str(OUT/f'{n}_{t:04d}_overlay.jpg'),pic)
 if t%80==0:print('refit',t,flush=True)
cap.release();(OUT/'all_candidates.json').write_text(json.dumps(records));arrays={'time':cam['time']};summary={}
for n in meshes:
 P=np.asarray(poses[n]);good=np.array([r['accepted'] for r in records[n]]);idx=np.flatnonzero(good);assert len(idx)>N*.5,(n,len(idx))
 ts=np.clip(np.arange(N),idx[0],idx[-1]);p=np.stack([np.interp(ts,idx,P[idx,k,3]) for k in range(3)],axis=1);rot=Slerp(idx,Rotation.from_matrix(P[idx,:3,:3]))(ts)
 arrays.update({n+'_position_world':p,n+'_quaternion_xyzw':rot.as_quat(),n+'_keyframes':np.arange(N),n+'_keyframe_valid':good})
 summary[n]=dict(accepted=int(good.sum()),old_iou_median=float(np.median([r['old_outline_iou'] for r in records[n]])),new_iou_median=float(np.median([r['new_outline_iou'] for r in records[n]])),frames=records[n])
np.savez_compressed(OUT/'object_poses.npz',**arrays);(OUT/'validation.json').write_text(json.dumps(summary,indent=2));print({n:{k:v for k,v in s.items() if k!='frames'} for n,s in summary.items()})