"""Compare rigid rotations and separately labelled reflection diagnostics.

No scaling, nonrigid warping, or deletion of unmatched geometry is permitted.
Use Euclidean residuals over the full spatially sampled scan to rank poses.
"""
from ac28_audit import *
from scipy.spatial import cKDTree
import itertools

def frame(p):
    c=p.mean(0); _,v=np.linalg.eigh(np.cov(p.T)); v=v[:,::-1]
    if np.linalg.det(v)<0:v[:,-1]*=-1
    return c,v

def transform(p,t): return p@t[:3,:3].T+t[:3,3]

def fit(a,b):
    ac=a.mean(0);bc=b.mean(0);u,_,vt=np.linalg.svd((a-ac).T@(b-bc))
    r=vt.T@u.T
    if np.linalg.det(r)<0:vt[-1]*=-1;r=vt.T@u.T
    t=np.eye(4);t[:3,:3]=r;t[:3,3]=bc-r@ac
    return t

def icp(p, cp, initial):
    tree=cKDTree(cp); t=initial.copy()
    for threshold in (150,60,25,12,6):
        for _ in range(45):
            q=transform(p,t); d,idx=tree.query(q,workers=2)
            mask=d<min(threshold,np.quantile(d,.98))
            if mask.sum()<30:break
            delta=fit(q[mask],cp[idx[mask]]);t=delta@t
            if np.linalg.norm(delta[:3,3])<.002 and np.linalg.norm(delta[:3,:3]-np.eye(3))<1e-6:break
    d=tree.query(transform(p,t),workers=2)[0]
    return t,float(np.mean(np.minimum(d,60)**2)**.5),np.quantile(d,[.5,.9,.95]).tolist()

def main():
    p0=np.load(OUT/'16h_all_merge_pointcloud_full.npz')['points']
    cp0=np.load(OUT/'cad_sample.npy')
    p=np.asarray(cloud(p0).voxel_down_sample(6).points)
    cp=np.asarray(cloud(cp0).voxel_down_sample(4).points)
    cs,vs=frame(p);cc,vc=frame(cp)
    candidates=[];start=time.time()
    for perm in itertools.permutations(range(3)):
        for signs in itertools.product((-1,1),repeat=3):
            a=np.eye(3)[:,perm]@np.diag(signs)
            r=vc@a@vs.T;t=np.eye(4);t[:3,:3]=r;t[:3,3]=cc-r@cs
            t,score,quant=icp(p,cp,t)
            row={'det':int(round(np.linalg.det(t[:3,:3]))),'score':score,'quantiles':quant,'transform':t.tolist(),'seed':str((perm,signs))}
            candidates.append(row)
        best=sorted(candidates,key=lambda x:x['score'])[:2]
        print('elapsed',round(time.time()-start),'perm',perm,'best',[(x['det'],round(x['score'],2)) for x in best],flush=True)
    candidates.sort(key=lambda x:x['score'])
    (OUT/'pose_candidates.json').write_text(json.dumps(candidates,indent=2))
    mesh=o3d.io.read_triangle_mesh(str(DATA/'viewer/part.stl'));mesh.remove_duplicated_vertices()
    scene=o3d.t.geometry.RaycastingScene();scene.add_triangles(o3d.t.geometry.TriangleMesh.from_legacy(mesh))
    results={}
    for kind,det in [('rigid',1),('reflection_diagnostic',-1)]:
        best=next(x for x in candidates if x['det']==det);t=np.array(best['transform'])
        for _ in range(50):
            q=transform(p0,t);closest=scene.compute_closest_points(o3d.core.Tensor(q.astype('float32')))['points'].numpy()
            d=np.linalg.norm(q-closest,axis=1);mask=d<12
            delta=fit(q[mask],closest[mask]);t=delta@t
            if np.linalg.norm(delta[:3,3])<.0005:break
        q=transform(p0,t);closest=scene.compute_closest_points(o3d.core.Tensor(q.astype('float32')))['points'].numpy()
        d=np.linalg.norm(q-closest,axis=1)
        results[kind]={'transform':t.tolist(),'det':det,'median_mm':float(np.median(d)),'p90_mm':float(np.quantile(d,.9)),'p95_mm':float(np.quantile(d,.95)),'within5':float(np.mean(d<5)),'within10':float(np.mean(d<10)),'capped_rmse60':float(np.mean(np.minimum(d,60)**2)**.5)}
        np.savez_compressed(OUT/f'16h_{kind}.npz',points=q,distance=d,transform=t)
    (OUT/'registration_comparison.json').write_text(json.dumps(results,indent=2))
    fig,axes=plt.subplots(2,3,figsize=(17,9))
    for row,kind in enumerate(results):
        q=np.load(OUT/f'16h_{kind}.npz')['points']
        for ax,(i,j) in zip(axes[row],[(0,2),(0,1),(1,2)]):
            ax.scatter(cp0[::3,i],cp0[::3,j],s=.3,c='#546578',label='Design')
            ax.scatter(q[:,i],q[:,j],s=.3,c='#e17437',label='Scan')
            ax.set_aspect('equal');ax.set_title(f'{kind}: {"XYZ"[i]} / {"XYZ"[j]}');ax.legend(markerscale=8)
    fig.tight_layout();fig.savefig(OUT/'rotation_vs_reflection.png',dpi=140)
    print(json.dumps(results,indent=2),flush=True)

if __name__=='__main__':main()