from __future__ import annotations
import argparse,json,time,torch

def chamfer(a,b):
    d=torch.cdist(a.float()[None],b.float()[None]).squeeze(0);return float(d.min(1).values.mean()+d.min(0).values.mean())
def symmetry_x(v):
    a=v.float().clone();a[:,0]*=-1;d=torch.cdist(v.float()[None],a[None]).squeeze(0);return float(d.min(1).values.mean())
def joint_error(a,b):return float(torch.linalg.vector_norm(a.float()-b.float(),dim=-1).mean()) if a.shape==b.shape else None

def main():
    ap=argparse.ArgumentParser();ap.add_argument('--teacher',required=True);ap.add_argument('--student',required=True);ap.add_argument('--out',required=True);a=ap.parse_args();t=torch.load(a.teacher,map_location='cpu',weights_only=False);s=torch.load(a.student,map_location='cpu',weights_only=False)
    r={}
    if 'vertices' in t and 'vertices' in s:r['chamfer']=chamfer(s['vertices'],t['vertices']);r['symmetry_x']=symmetry_x(s['vertices'])
    if 'joints' in t and 'joints' in s:r['joint_error']=joint_error(s['joints'],t['joints'])
    if 'runtime_s' in s:r['runtime_s']=s['runtime_s'];r['speedup']=t.get('runtime_s',0)/s['runtime_s'] if t.get('runtime_s') else None
    open(a.out,'w').write(json.dumps(r,indent=2));print(json.dumps(r))
if __name__=='__main__':main()