"""Offline, rotation-safe smoothing and conservative visible-surface RGB-D fit."""
import sys,json
from pathlib import Path
import numpy as np
import torch,cv2
from scipy.spatial.transform import Rotation
from scipy.ndimage import gaussian_filter1d
ROOT=Path(__file__).resolve().parents[1]
sys.path.insert(0,str(ROOT/'third_party/Dyn-HaMR/dyn-hamr'))
from body_model import MANO,run_mano
BASE=ROOT/'output/20260915_171525_dynhamr'; OUT=BASE/'corrected'; OUT.mkdir(exist_ok=True)
SRC=ROOT/'docs/20260915_171525'; meta=json.loads((SRC/'intrinsics.json').read_text())
a=dict(np.load(BASE/'optimization/prior/20260915_171525_000000_world_results.npz')); N=352
torch.set_num_threads(4)
model=MANO(model_path=str(ROOT/'third_party/Dyn-HaMR/_DATA/data/mano'),batch_size=704,pose2rot=True)
def geometry(d):
 with torch.no_grad():
  o=run_mano(model,*[torch.tensor(d[k]).float() for k in ['trans','root_orient','pose_body','is_right','betas']])
 return o['joints'].numpy(),o['vertices'].numpy()
def smooth_rotation(v,sigma=1.2):
 shape=v.shape; v=v.reshape(N,-1,3); out=np.empty_like(v)
 for j in range(v.shape[1]):
  r=Rotation.from_rotvec(v[:,j])
  for t in range(N):
   ix=np.arange(max(0,t-4),min(N,t+5)); w=np.exp(-.5*((ix-t)/sigma)**2)
   out[t,j]=r[ix].mean(weights=w).as_rotvec()
 return out.reshape(shape)
j0,v0=geometry(a); d={k:v.copy() for k,v in a.items()}
for b in range(2):
 d['root_orient'][b]=smooth_rotation(a['root_orient'][b]); d['pose_body'][b]=smooth_rotation(a['pose_body'][b])
# Correct translations in world coordinates; left MANO translation has reflected X.
R=a['cam_R'][0]; ct=a['cam_t'][0]; K=np.array([meta[k] for k in ['fx','fy','cx','cy']])
vc=np.einsum('tij,btvj->btvi',R,v0)+ct[None,:,None,:]
wc=np.einsum('tij,btj->bti',R,j0[:,:,0])+ct[None,:,:]
raw=np.zeros((2,N)); valid=np.zeros((2,N),bool); counts=np.zeros((2,N),int)
cap=cv2.VideoCapture(str(SRC/'color.mp4'))
for t in range(N):
 ok,c=cap.read(); assert ok
 dep=cv2.imread(str(SRC/'depth'/f'{t:06d}.png'),-1)*meta['depth_scale_m']
 skin=cv2.inRange(cv2.cvtColor(c,cv2.COLOR_BGR2YCrCb),np.array([0,133,77]),np.array([255,173,127]))>0
 for b in range(2):
  v=vc[b,t]; uv=np.rint(v[:,:2]/v[:,2,None]*K[:2]+K[2:]).astype(int); closest={}
  for i in np.flatnonzero((v[:,2]>.1)&(uv[:,0]>=2)&(uv[:,0]<846)&(uv[:,1]>=2)&(uv[:,1]<478)):
   x,y=uv[i]; key=(x//4,y//4)
   if key not in closest or v[i,2]<v[closest[key],2]: closest[key]=i
  errors=[]
  for i in closest.values():
   x,y=uv[i]
   if not skin[y,x]:continue
   patch=dep[y-1:y+2,x-1:x+2]; z=patch[(patch>.1)&(patch<.85)]
   if len(z)>=5 and np.ptp(z)<.02:errors.append(np.median(z)-v[i,2])
  counts[b,t]=len(errors)
  if errors:
   dz=np.median(errors); mad=np.median(np.abs(np.array(errors)-dz));raw[b,t]=dz
   valid[b,t]=len(errors)>=25 and mad<.01 and abs(dz)<.025
cap.release()
# Only directly supported neighborhoods receive a depth update. No long-gap extrapolation.
shifts=np.zeros_like(raw)
for b in range(2):
 numerator=gaussian_filter1d(raw[b]*valid[b],2); support=gaussian_filter1d(valid[b].astype(float),2)
 shifts[b]=np.where(support>.6,numerator/np.maximum(support,1e-6),0)
 shifts[b]=gaussian_filter1d(shifts[b],1)
 delta_c=wc[b]/wc[b,:,2,None]*shifts[b,:,None]
 delta_w=np.einsum('tji,tj->ti',R,delta_c)
 delta_w[:,0]*=2*a['is_right'][b]-1
 d['trans'][b]=gaussian_filter1d(a['trans'][b]+delta_w,1.5,axis=0)
j1,v1=geometry(d)
report={'frames':N,'method':'Rotation weighted means on SO(3), sigma 1.2 frames; translation Gaussian sigma 1.5 frames; conservative skin-depth residual <=25mm and MAD<10mm, support-weighted local correction','offline_uses_future_frames':True,'depth_fit_is_surface_estimate_not_ground_truth':True,'hands':{}}
for b in range(2):
 side='right' if a['is_right'][b,0]>.5 else 'left'
 def rms2(x):return float(np.sqrt(np.mean(np.sum(np.diff(x,n=2,axis=0)**2,axis=-1)))*1000)
 deviation=np.linalg.norm(j1[b,:,0]-j0[b,:,0],axis=-1)*1000
 angles=(Rotation.from_rotvec(a['root_orient'][b]).inv()*Rotation.from_rotvec(d['root_orient'][b])).magnitude()*180/np.pi
 report['hands'][side]={'depth_accepted_frames':int(valid[b].sum()),'depth_shift_max_mm':float(abs(shifts[b]).max()*1000),'wrist_second_difference_rms_mm_before':rms2(j0[b,:,0]),'wrist_second_difference_rms_mm_after':rms2(j1[b,:,0]),'wrist_change_max_mm':float(deviation.max()),'wrist_change_p95_mm':float(np.percentile(deviation,95)),'root_change_max_deg':float(angles.max())}
 assert np.isfinite(j1[b]).all() and deviation.max()<40 and angles.max()<15
phase=OUT/'prior';phase.mkdir(exist_ok=True)
np.savez_compressed(phase/'20260915_171525_000000_world_results.npz',**d)
np.savez_compressed(OUT/'depth_fit.npz',raw_shift=raw,accepted=valid,applied_shift=shifts,sample_count=counts)
np.savez_compressed(OUT/'comparison_joints.npz',before=j0,after=j1)
(OUT/'validation.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2))