#!/usr/bin/env python3
"""
det04 — Find the first op in det_500m.onnx (via onnx2torch) where B=1 and
        frame[0] of B=2 outputs diverge.

Registers forward hooks on all submodules, runs B=1 and B=2 on the same image,
then reports the first module whose output differs — the culprit causing cross-frame
contamination.
"""
import argparse
import cv2
import torch
import onnx2torch

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 image ────────────────────────────────────────────────────────────────
import cv2, numpy as np
img = cv2.imread(args.image)
assert img is not None, f'Could not read image: {args.image}'
blob = cv2.dnn.blobFromImage(
    cv2.resize(img, DET_SIZE), 1.0 / 128.0, DET_SIZE,
    (127.5, 127.5, 127.5), swapRB=True,
)
t = torch.from_numpy(blob)

# ── Load model ────────────────────────────────────────────────────────────────
print('Loading onnx2torch model ...')
pt = onnx2torch.convert(args.model)
pt.eval()

# ── Hook: capture one output tensor per named module ─────────────────────────
acts_b1 = {}
acts_b2 = {}

def make_hook(store, name):
    def hook(module, input, output):
        out = output[0] if isinstance(output, (tuple, list)) else output
        store[name] = out.detach()
    return hook

handles = []
for name, module in pt.named_modules():
    if name == '':
        continue
    handles.append(module.register_forward_hook(make_hook(acts_b1, name)))

# ── B=1 forward ───────────────────────────────────────────────────────────────
with torch.no_grad():
    pt(t)

for h in handles: h.remove()
handles.clear()

for name, module in pt.named_modules():
    if name == '':
        continue
    handles.append(module.register_forward_hook(make_hook(acts_b2, name)))

# ── B=2 forward ───────────────────────────────────────────────────────────────
t2 = torch.cat([t, t], dim=0)
with torch.no_grad():
    pt(t2)

for h in handles: h.remove()

# ── Compare: find first divergence ───────────────────────────────────────────
print(f'\nComparing {len(acts_b1)} intermediate activations ...')
print(f'{"Module":<55} {"B1 shape":<22} {"B2 shape":<22} {"max_diff":>10}  status')
print('-' * 120)

first_culprit = None
for name in acts_b1:
    if name not in acts_b2:
        continue
    a1 = acts_b1[name]
    a2 = acts_b2[name]

    try:
        if a2.shape[0] == 2:
            a2_f0  = a2[0]
            a1_cmp = a1[0] if a1.shape[0] == 1 else a1
        elif a1.numel() > 0 and a2.numel() == 2 * a1.numel():
            a2_f0  = a2.flatten()[:a1.numel()].reshape(a1.shape)
            a1_cmp = a1
        else:
            continue
    except Exception:
        continue

    try:
        diff = (a1_cmp.float() - a2_f0.float()).abs().max().item()
    except Exception:
        continue

    ok = diff < 1e-4
    if not ok and first_culprit is None:
        first_culprit = name
        status = '<<< FIRST MISMATCH'
    elif not ok:
        status = 'MISMATCH'
    else:
        status = 'OK'

    if not ok:
        print(f'{name:<55} {str(tuple(a1.shape)):<22} {str(tuple(a2.shape)):<22} {diff:>10.5f}  {status}')

if first_culprit:
    print(f'\n>>> Culprit: {first_culprit}')
else:
    print('\nNo divergence found — model is batch-independent.')

print('\nDone.')