from __future__ import annotations
import argparse,json
from pathlib import Path
import numpy as np
import torch
import onnxruntime as ort
from runtime.trt_torch import TorchTensorRTEngine

def stats(a,b):
    a=a.astype(np.float32).ravel(); b=b.astype(np.float32).ravel()
    d=np.abs(a-b)
    denom=(np.linalg.norm(a)*np.linalg.norm(b)+1e-12)
    cos=float(np.dot(a,b)/denom)
    return {'max_abs':float(d.max()),'mean_abs':float(d.mean()),'cosine':cos}

def ort_run(path,feeds):
    sess=ort.InferenceSession(str(path),providers=['CUDAExecutionProvider','CPUExecutionProvider'])
    npfeeds={k:v.detach().cpu().numpy() for k,v in feeds.items()}
    return sess.run(None,npfeeds)[0]

def trt_run(path,feeds):
    e=TorchTensorRTEngine(path); out=e.run({k:v.clone() for k,v in feeds.items()}); torch.cuda.synchronize(); return next(iter(out.values())).detach().cpu().numpy()

def main():
    ap=argparse.ArgumentParser();ap.add_argument('--root',default='.');a=ap.parse_args();r=Path(a.root);res={}
    feeds={'pixel_values':torch.randn(1,3,518,518,device='cuda')}
    res['dinov2']=stats(ort_run(r/'onnx/dinov2/model.onnx',feeds),trt_run(r/'engines/l4-sm89/dinov2.plan',feeds))
    feeds={'image':torch.randn(1,3,544,544,device='cuda'),'intrins':torch.tensor([[[471.117,0,259.0],[0,471.117,259.0],[0,0,1.0]]],device='cuda')}
    res['dsine']=stats(ort_run(r/'onnx/dsine/model.onnx',feeds),trt_run(r/'engines/l4-sm89/dsine.plan',feeds))
    res['passed']=res['dinov2']['cosine']>0.999 and res['dsine']['cosine']>0.995
    print(json.dumps(res,indent=2)); (r/'validation_l4.json').write_text(json.dumps(res,indent=2))
    if not res['passed']: raise SystemExit('engine validation failed')
if __name__=='__main__':main()