#!/usr/bin/env python3
"""
det03 — Convert det_500m.onnx to PyTorch (onnx2torch) and test batch independence.

Runs three forward passes with the same real image:
  A: ONNXRuntime,  B=1
  B: onnx2torch,   B=1
  C: onnx2torch,   B=5  (same image duplicated)

Prints C output shapes, then compares:
  A vs B   — confirms onnx2torch faithfully reproduces the ONNX computation
  B vs C0  — checks if frame[0] of B=5 matches the B=1 output (batch independence test)
  C0 vs C1 — checks if both frames in B=5 match each other (same input → same output)
"""
import argparse
import cv2
import numpy as np
import torch
import onnx2torch
import onnxruntime as ort

DET_SIZE = (640, 640)

parser = argparse.ArgumentParser()
parser.add_argument('--model', required=True, help='Path to det_500m.onnx')
parser.add_argument('--image', required=True, help='Path to any face image (jpg/png)')
args = parser.parse_args()

# ── Load models ───────────────────────────────────────────────────────────────
print('Loading ONNXRuntime session ...')
sess = ort.InferenceSession(args.model, providers=['CPUExecutionProvider'])
inp_name = sess.get_inputs()[0].name

print('Loading onnx2torch model ...')
pt_model = onnx2torch.convert(args.model)
pt_model.eval()
print('Done.\n')

# ── Preprocess image ──────────────────────────────────────────────────────────
img = cv2.imread(args.image)
assert img is not None, f'Could not read image: {args.image}'
print(f'Image: {args.image}  original shape={img.shape}')

blob = cv2.dnn.blobFromImage(
    cv2.resize(img, DET_SIZE),
    1.0 / 128.0, DET_SIZE,
    (127.5, 127.5, 127.5), swapRB=True,
)  # shape: (1, 3, 640, 640)

t = torch.from_numpy(blob)

# ── A) ONNXRuntime B=1 ───────────────────────────────────────────────────────
ort_out = sess.run(None, {inp_name: blob})

# ── B) onnx2torch B=1 ────────────────────────────────────────────────────────
with torch.no_grad():
    pt_out_b1 = pt_model(t)

# ── C) onnx2torch B=5 [img, img, img, img, img] ──────────────────────────────
t2 = torch.cat([t] * 5, dim=0)
with torch.no_grad():
    pt_out_b2 = pt_model(t2)

# ── Print B=5 output shapes ───────────────────────────────────────────────────
print('\n[C] Output shapes for B=5:')
for i, c in enumerate(pt_out_b2):
    print(f'  out[{i}]: {list(c.shape)}')

# ── Compare ───────────────────────────────────────────────────────────────────
print('\n[D] Comparison across all 9 output heads\n')
print(f'  {"head":<6} {"AvsB":>10}  {"BvsC0":>10}  {"C0vsC1":>10}  note')
print('  ' + '-' * 60)

for i, (a, b, c) in enumerate(zip(ort_out, pt_out_b1, pt_out_b2)):
    a = a.flatten()
    b = b.detach().numpy().flatten()
    c = c.detach().flatten()

    K = b.shape[0]
    c0 = c[:K]
    c1 = c[K:]

    diff_ab    = np.abs(a - b).max()
    diff_b_c0  = (torch.from_numpy(b) - c0).abs().max().item()
    diff_c0_c1 = (c0 - c1).abs().max().item()

    note = ''
    if diff_b_c0 > 1e-4:
        note = '<-- CONTAMINATION'

    print(f'  out[{i}]  {diff_ab:>10.5f}  {diff_b_c0:>10.5f}  {diff_c0_c1:>10.5f}  {note}')

print('\nPass criteria: BvsC0 < 1e-4 and C0vsC1 < 1e-4 for all heads.')
print('Done.')