#!/usr/bin/env python3
"""
det02 — ONNX graph surgery to patch det_500m.onnx for dynamic batch.

Patches 3 things in the graph metadata:
  1. Input dim[0]: 1 → symbolic 'batch'
  2. Reshape shape constants starting with 1: that 1 → 0 (copy batch dim from input)
  3. Output dim[0]: 1 → symbolic 'batch'

Saves the result, then prints declared output shapes and actual runtime shapes for N=1 and N=2.
"""
import argparse
import numpy as np
import onnx
from onnx import numpy_helper
import onnxruntime as ort

parser = argparse.ArgumentParser()
parser.add_argument('--model', required=True, help='Path to det_500m.onnx')
parser.add_argument('--out',   required=True, help='Output path for patched model')
args = parser.parse_args()

model = onnx.load(args.model)

# ── 1. Patch input batch dim ──────────────────────────────────────────────────
for inp in model.graph.input:
    d0 = inp.type.tensor_type.shape.dim[0]
    d0.ClearField('dim_value')
    d0.dim_param = 'batch'
    print(f"  Input '{inp.name}': batch dim → dynamic")

# ── 2. Patch Reshape shape initializers that start with 1 ────────────────────
reshape_shape_names = set()
for node in model.graph.node:
    if node.op_type == 'Reshape':
        reshape_shape_names.add(node.input[1])

patched = 0
for init in list(model.graph.initializer):
    if init.name not in reshape_shape_names:
        continue
    arr = numpy_helper.to_array(init).copy()
    if arr.ndim == 1 and arr[0] == 1:
        print(f"  Reshape initializer '{init.name}': {arr} → ", end='')
        arr[0] = 0
        print(arr)
        new_t = numpy_helper.from_array(arr, init.name)
        model.graph.initializer.remove(init)
        model.graph.initializer.append(new_t)
        patched += 1

print(f'  Patched {patched} Reshape shape constant(s)')

# ── 3. Patch output batch dims ────────────────────────────────────────────────
for out in model.graph.output:
    d0 = out.type.tensor_type.shape.dim[0]
    d0.ClearField('dim_value')
    d0.dim_param = 'batch'
    print(f"  Output '{out.name}': batch dim → dynamic")

# ── Save ──────────────────────────────────────────────────────────────────────
onnx.checker.check_model(model)
onnx.save(model, args.out)
print(f'\nSaved: {args.out}')

# ── Declared output shapes (from patched graph) ───────────────────────────────
print('\nDeclared output shapes after surgery:')
for out in model.graph.output:
    shape = [d.dim_param if d.HasField('dim_param') else d.dim_value
             for d in out.type.tensor_type.shape.dim]
    print(f'  {out.name}: {shape}')

# ── Actual output shapes (from runtime with N=1 and N=2) ─────────────────────
print('\nActual output shapes at runtime:')
session = ort.InferenceSession(args.out, providers=['CPUExecutionProvider'])
inp0 = session.get_inputs()[0]
out_names = [o.name for o in session.get_outputs()]

for N in (1, 2):
    dummy = np.random.randint(0, 255, (N, 3, 640, 640)).astype(np.float32)
    try:
        outs = session.run(out_names, {inp0.name: dummy})
        print(f'  N={N}:')
        for name, o in zip(out_names, outs):
            print(f'    {name}: {list(o.shape)}')
    except Exception as e:
        print(f'  N={N}: FAILED — {e}')