#!/usr/bin/env python3
"""Wrap a flat policy ONNX input as a fixed 4D feature-map input for X5 hb_mapper."""

import argparse
from pathlib import Path

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


def tensor_shape(value_info):
    return [dim.dim_value for dim in value_info.type.tensor_type.shape.dim]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--flat-dim", type=int, default=450)
    parser.add_argument("--input-name", default="obs_4d")
    args = parser.parse_args()

    model = onnx.load(str(args.input))
    if len(model.graph.input) != 1:
        raise ValueError(f"Expected one input, got {len(model.graph.input)}")

    old_input = model.graph.input[0]
    old_name = old_input.name
    old_shape = tensor_shape(old_input)
    if old_shape != [1, args.flat_dim]:
        raise ValueError(f"Expected input shape [1, {args.flat_dim}], got {old_shape}")

    reshape_out_name = "__bpu_flat_input"
    flatten_node = helper.make_node(
        "Flatten",
        inputs=[args.input_name],
        outputs=[reshape_out_name],
        name="Bpu4dInputFlatten",
        axis=1,
    )

    for node in model.graph.node:
        for index, name in enumerate(node.input):
            if name == old_name:
                node.input[index] = reshape_out_name

    new_input = helper.make_tensor_value_info(
        args.input_name, TensorProto.FLOAT, [1, 1, 1, args.flat_dim]
    )
    model.graph.input.remove(old_input)
    model.graph.input.insert(0, new_input)
    model.graph.node.insert(0, flatten_node)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    onnx.save(model, str(args.output))
    print(f"Wrote {args.output}")


if __name__ == "__main__":
    main()
