#!/usr/bin/env python3
"""Convert policy ONNX to an older opset accepted by X5 hb_mapper."""

import argparse
from pathlib import Path

import onnx
from onnx import version_converter


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

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


if __name__ == "__main__":
    main()
