"""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_tracked';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))))
polygons={'upper':np.array([[491,268],[601,275],[631,399],[488,379]],dtype=np.float32),'lower':np.array([[309,254],[436,264],[442,382],[291,375]],dtype=np.float32)}
previous_gray=None
face_choice={}
for t in range(N):
 ok,im=cap.read();assert ok
 gray=cv2.cvtColor(im,cv2.COLOR_BGR2GRAY)
 if previous_gray is not None:
  for n,poly in polygons.items():
   featuremask=np.zeros_like(gray);cv2.fillConvexPoly(featuremask,poly.astype('int32'),255)
   points=cv2.goodFeaturesToTrack(previous_gray,300,.01,5,mask=featuremask)
   if points is not None and len(points)>=8:
    nxt,st,_=cv2.calcOpticalFlowPyrLK(previous_gray,gray,points,None,winSize=(21,21),maxLevel=3)
    back,st2,_=cv2.calcOpticalFlowPyrLK(gray,previous_gray,nxt,None,winSize=(21,21),maxLevel=3)
    good=(st.ravel()>0)&(st2.ravel()>0)&(np.linalg.norm(back-points,axis=(1,2))<1.)
    if good.sum()>=8:
     aff,inl=cv2.estimateAffinePartial2D(points[good],nxt[good],method=cv2.RANSAC,ransacReprojThreshold=2.)
     H=None if aff is None else np.vstack([aff,[0,0,1]])
     if H is not None and inl.sum()>=8:
      proposed=cv2.perspectiveTransform(poly[None],H)[0]
      if np.max(np.linalg.norm(proposed-poly,axis=1))<30:polygons[n]=proposed
 previous_gray=gray
 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
  mask &= ~(cv2.inRange(cv2.cvtColor(im,cv2.COLOR_BGR2YCrCb),np.array([0,133,77]),np.array([255,173,127]))>0)
  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))
  if cv2.contourArea(hull)>6500 and mask.sum()/cv2.contourArea(hull)>.55:
   for eps in [.015,.025,.04,.06]:
    quad=cv2.approxPolyDP(hull,eps*cv2.arcLength(hull,True),True)
    if len(quad)==4:
     poly=quad[:,0].astype('float32');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)
     poly=min([np.roll(poly,k,axis=0) for k in range(4)],key=lambda q:np.sum((q-polygons[n])**2))
     polygons[n]=poly;break
  poly=polygons[n].astype(float)
  observed=np.zeros_like(mask);cv2.fillConvexPoly(observed,poly.astype('int32'),1)
  mask=mask*observed
  # 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]
  if len(pts)<30:pts=np.array([[0,0,.5],[.1,0,.5],[0,.1,.5]])
  pc=o3d.geometry.PointCloud(o3d.utility.Vector3dVector(pts[::max(1,len(pts)//1500)]));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=[]
  for face in [lo[1],hi[1]]:
   for flip in [False,True]:
    loc=local.copy();loc[:,1]=face
    if flip:loc=loc[[3,2,1,0]]
    success,rv,tv=cv2.solvePnP(loc.astype(float),poly.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)
  previous=None if not poses[n] else np.linalg.inv(cam['c2w'][t])@poses[n][-1]
  if n not in face_choice:face_choice[n]=int(np.argmax([score(v,candidate,observed) for candidate in candidates]))
  T=candidates[face_choice[n]]
  rr=T[:3,:3]
  iou0=score(v,before,observed);iou1=score(v,T,observed);rmse=np.sqrt(np.mean(np.sum((local@rr.T+T[:3,3]-corners)**2,axis=1)))
  center_step=0 if previous is None else np.linalg.norm((T[:3,:3]@v.mean(0)+T[:3,3])-(previous[:3,:3]@v.mean(0)+previous[:3,3]))
  angle_step=0 if previous is None else Rotation.from_matrix(previous[:3,:3].T@T[:3,:3]).magnitude()
  accepted=bool(iou1>.60 and np.isfinite(T).all() and center_step<.04 and angle_step<.3)
  # Keep candidate for continuity diagnostics; rejected poses are excluded at interpolation.
  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()})