import onnx, sys
from onnx import helper
src, dst, mode = sys.argv[1], sys.argv[2], sys.argv[3]  # mode: pads | static:H,W
m = onnx.load(src)
# 1) auto_pad SAME_UPPER -> explicit pads (kernel k, stride 1 => total pad k-1, upper side)
n_fixed = 0
for n in m.graph.node:
    if n.op_type in ("MaxPool", "AveragePool", "Conv"):
        ap = [a for a in n.attribute if a.name == "auto_pad"]
        if ap and ap[0].s in (b"SAME_UPPER", b"SAME_LOWER"):
            ks = [a for a in n.attribute if a.name == "kernel_shape"][0].ints
            st = [a for a in n.attribute if a.name == "strides"]
            st = st[0].ints if st else [1]*len(ks)
            assert all(s == 1 for s in st), (n.name, list(st))
            tot = [k-1 for k in ks]
            if ap[0].s == b"SAME_UPPER":
                pads = [t//2 for t in tot] + [t - t//2 for t in tot]
            else:
                pads = [t - t//2 for t in tot] + [t//2 for t in tot]
            n.attribute.remove(ap[0])
            n.attribute.append(helper.make_attribute("pads", pads))
            n_fixed += 1
            print("  fixed", n.op_type, n.name, "pads", pads)
# 2) optionally fix input dims
if mode.startswith("static"):
    hw = [int(x) for x in mode.split(":")[1].split(",")]
    x = m.graph.input[0]
    d = x.type.tensor_type.shape.dim
    d[0].dim_value = 1; d[0].ClearField("dim_param")
    if not d[2].dim_value: d[2].dim_value = hw[0]; d[2].ClearField("dim_param")
    d[3].dim_value = hw[1]; d[3].ClearField("dim_param")
    # clear output shapes so inference recomputes
    for o in m.graph.output:
        o.type.tensor_type.ClearField("shape")
    del m.graph.value_info[:]
    m = onnx.shape_inference.infer_shapes(m)
    o = m.graph.output[0]
    if not o.type.tensor_type.HasField("shape") or len(o.type.tensor_type.shape.dim)==0:
        # data-dependent graph (Shape/Reshape): set the known shape explicitly
        C = int(sys.argv[4]) if len(sys.argv) > 4 else None
        shp = [1, 1, hw[0], hw[1]] if C is None else [1, hw[1]//8, C]
        o.type.CopyFrom(helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, shp))
        print("  output shape forced to", shp)
    print("  input", [dd.dim_value or dd.dim_param for dd in m.graph.input[0].type.tensor_type.shape.dim],
          "output", [dd.dim_value or dd.dim_param for dd in m.graph.output[0].type.tensor_type.shape.dim])
onnx.checker.check_model(m)
onnx.save(m, dst)
print(dst, "auto_pad fixed:", n_fixed)