#!/usr/bin/env python3
"""Downgrade the current RobotLab policy ONNX from opset18 to opset11.

This is intentionally narrow: it only rewrites ops whose axes changed from an
attribute in opset11 to an input tensor in newer opsets. The policy graph uses
otherwise opset11-compatible operators.
"""

import argparse
from pathlib import Path

import numpy as np
import onnx
from onnx import helper, numpy_helper


AXES_INPUT_OPS = {"Unsqueeze", "Squeeze", "ReduceL2", "ReduceSum"}
REMOVE_ATTRS = {
    "Reshape": {"allowzero"},
}


def constant_values(model):
    values = {}
    for initializer in model.graph.initializer:
        values[initializer.name] = numpy_helper.to_array(initializer)

    for node in model.graph.node:
        if node.op_type != "Constant" or len(node.output) != 1:
            continue
        for attr in node.attribute:
            if attr.name == "value":
                values[node.output[0]] = numpy_helper.to_array(attr.t)
                break
    return values


def axes_from_value(values, name):
    if name not in values:
        raise KeyError(f"Cannot find constant axes tensor {name!r}")
    axes = np.asarray(values[name]).astype(np.int64).reshape(-1).tolist()
    return [int(axis) for axis in axes]


def set_opset(model, target):
    for opset in model.opset_import:
        if opset.domain in ("", "ai.onnx"):
            opset.version = target
            return
    model.opset_import.append(helper.make_operatorsetid("", target))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--target-opset", type=int, default=11)
    parser.add_argument("--ir-version", type=int, default=7)
    args = parser.parse_args()

    model = onnx.load(str(args.input))
    values = constant_values(model)
    converted = 0

    for node in model.graph.node:
        if node.op_type in REMOVE_ATTRS:
            remove_names = REMOVE_ATTRS[node.op_type]
            kept_attrs = [attr for attr in node.attribute if attr.name not in remove_names]
            del node.attribute[:]
            node.attribute.extend(kept_attrs)

        if node.op_type not in AXES_INPUT_OPS or len(node.input) < 2:
            continue
        axes = axes_from_value(values, node.input[1])
        del node.input[1:]
        node.attribute.extend([helper.make_attribute("axes", axes)])
        converted += 1

    set_opset(model, args.target_opset)
    model.ir_version = min(model.ir_version, args.ir_version)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    onnx.save(model, str(args.output))
    print(f"Wrote {args.output}")
    print(f"converted_axes_nodes={converted}")
    print(f"ir={model.ir_version}")
    print(f"opsets={[(op.domain, op.version) for op in model.opset_import]}")


if __name__ == "__main__":
    main()
