#!/usr/bin/env python3
"""Rewrite oemer U-Net ConvTranspose pads so ONNX Runtime can load the model.

oemer 0.1.8's 1st_model.onnx ships ConvTranspose nodes with pads
[0, 0, -1, -1]. ONNX Runtime 1.30 and newer reject those during shape
inference. pads [0, 0, 0, 0] with output_padding [1, 1] is the same
output size for these stride-2, kernel-1 nodes. Weights are left as they are.
"""

from __future__ import annotations

import sys
from pathlib import Path

import onnx
from onnx import helper

BAD_PADS = [0, 0, -1, -1]
FIXED_PADS = [0, 0, 0, 0]
FIXED_OUTPUT_PADDING = [1, 1]
CRASH_NODE = "model/conv2d_transpose_1/conv2d_transpose"


def _attr(node: onnx.NodeProto, name: str) -> onnx.AttributeProto | None:
    for attr in node.attribute:
        if attr.name == name:
            return attr
    return None


def _ints(node: onnx.NodeProto, name: str) -> list[int] | None:
    attr = _attr(node, name)
    if attr is None:
        return None
    return list(attr.ints)


def _set_ints(attr: onnx.AttributeProto, values: list[int]) -> None:
    del attr.ints[:]
    attr.ints.extend(values)


def rewrite(model_path: Path) -> int:
    model = onnx.load(model_path)
    changed = 0
    for node in model.graph.node:
        if node.op_type != "ConvTranspose":
            continue
        pads = _ints(node, "pads")
        if pads is None or pads == FIXED_PADS:
            continue
        if pads != BAD_PADS:
            if any(value < 0 for value in pads):
                raise SystemExit(
                    f"{model_path}: {node.name} has unexpected negative pads {pads}"
                )
            continue
        pads_attr = _attr(node, "pads")
        if pads_attr is None:
            raise SystemExit(f"{model_path}: {node.name} is missing pads")
        _set_ints(pads_attr, FIXED_PADS)
        output_padding = _attr(node, "output_padding")
        if output_padding is None:
            node.attribute.append(
                helper.make_attribute("output_padding", FIXED_OUTPUT_PADDING)
            )
        else:
            _set_ints(output_padding, FIXED_OUTPUT_PADDING)
        changed += 1
        print(
            f"rewrote {node.name}: pads {FIXED_PADS} output_padding {FIXED_OUTPUT_PADDING}"
        )
    if changed:
        onnx.save(model, model_path)
    verify(model_path)
    print(f"rewrote {changed} ConvTranspose node(s) in {model_path}")
    return changed


def verify(model_path: Path) -> None:
    model = onnx.load(model_path)
    found = False
    for node in model.graph.node:
        if node.op_type != "ConvTranspose":
            continue
        pads = _ints(node, "pads") or []
        if any(value < 0 for value in pads):
            raise SystemExit(f"{model_path}: {node.name} still has negative pads {pads}")
        if node.name != CRASH_NODE:
            continue
        found = True
        output_padding = _ints(node, "output_padding")
        if pads != FIXED_PADS or output_padding != FIXED_OUTPUT_PADDING:
            raise SystemExit(
                f"{model_path}: {node.name} is {pads} / {output_padding}, "
                f"expected {FIXED_PADS} / {FIXED_OUTPUT_PADDING}"
            )
    if not found:
        raise SystemExit(f"{model_path}: missing {CRASH_NODE}")


def main() -> int:
    if len(sys.argv) != 2:
        print("usage: patch_oemer_unet.py MODEL.onnx", file=sys.stderr)
        return 2
    path = Path(sys.argv[1])
    if not path.is_file():
        print(f"missing model: {path}", file=sys.stderr)
        return 1
    rewrite(path)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())