import onnx
import onnx.utils

try:
    from onnx import optimizer
except ImportError:
    import onnxoptimizer as optimizer
import argparse

import tools.modhelper as helper
import tools.other as other
import tools.replacing as replacing

# Main process
# Argument parser
parser = argparse.ArgumentParser(
    description="Edit an ONNX model.\nThe processing sequense is 'delete "
                "nodes/values' -> 'add nodes' -> 'change shapes'.\nCutting "
                "cannot be done with other operations together"
)
parser.add_argument("in_file", type=str, help="input ONNX FILE")
parser.add_argument("out_file", type=str, help="ouput ONNX FILE")
parser.add_argument(
    "-c",
    "--cut",
    dest="cut_node",
    type=str,
    nargs="+",
    help="remove nodes from the given nodes(inclusive)",
)
parser.add_argument(
    "--cut-type",
    dest="cut_type",
    type=str,
    nargs="+",
    help="remove nodes by type from the given nodes(inclusive)",
)
parser.add_argument(
    "-d",
    "--delete",
    dest="delete_node",
    type=str,
    nargs="+",
    help="delete nodes by names and only those nodes",
)
parser.add_argument(
    "--delete-input",
    dest="delete_input",
    type=str,
    nargs="+",
    help="delete inputs by names",
)
parser.add_argument(
    "--delete-output",
    dest="delete_output",
    type=str,
    nargs="+",
    help="delete outputs by names",
)
parser.add_argument(
    "-i",
    "--input",
    dest="input_change",
    type=str,
    nargs="+",
    help="change input shape (e.g. -i 'input_0 1 3 224 224')",
)
parser.add_argument(
    "-o",
    "--output",
    dest="output_change",
    type=str,
    nargs="+",
    help="change output shape (e.g. -o 'input_0 1 3 224 224')",
)
parser.add_argument(
    "--add-conv",
    dest="add_conv",
    type=str,
    nargs="+",
    help="add nop conv using specific input",
)
parser.add_argument(
    "--add-bn",
    dest="add_bn",
    type=str,
    nargs="+",
    help="add nop bn using specific input",
)
parser.add_argument(
    "--rename-output",
    dest="rename_output",
    type=str,
    nargs="+",
    help="Rename the specific output(e.g. --rename-output old_name new_name)",
)
parser.add_argument(
    "--pixel-bias-value",
    dest="pixel_bias_value",
    type=str,
    nargs="+",
    help='(per channel) set pixel value bias bn layer at model front for '
         'normalization( e.g. --pixel_bias_value "[104.0, 117.0, 123.0]" )',
)
parser.add_argument(
    "--pixel-scale-value",
    dest="pixel_scale_value",
    type=str,
    nargs="+",
    help='(per channel) set pixel value scale bn layer at model front for '
         'normalization( e.g. --pixel_scale_value '
         '"[0.0078125, 0.0078125, 0.0078125]" )',
)

args = parser.parse_args()

# Load model and polish
m = onnx.load(args.in_file)
m = other.polish_model(m)
g = m.graph
replacing.replace_initializer_with_Constant(g)
other.topological_sort(g)

# Remove nodes according to the given arguments.
if args.delete_node is not None:
    helper.delete_nodes(g, args.delete_node)

if args.delete_input is not None:
    helper.delete_input(g, args.delete_input)

if args.delete_output is not None:
    helper.delete_output(g, args.delete_output)

# Add do-nothing Conv node
if args.add_conv is not None:
    other.add_nop_conv_after(g, args.add_conv)
    other.topological_sort(g)

# Add do-nothing BN node
if args.add_bn is not None:
    other.add_nop_bn_after(g, args.add_bn)
    other.topological_sort(g)

# Add bias scale BN node
if args.pixel_bias_value is not None or args.pixel_scale_value is not None:

    if len(g.input) > 1:
        raise ValueError(
            " '--pixel-bias-value' and '--pixel-scale-value' "
            "only support one input node model currently"
        )

    i_n = g.input[0]

    pixel_bias_value = [0] * i_n.type.tensor_type.shape.dim[1].dim_value
    pixel_scale_value = [1] * i_n.type.tensor_type.shape.dim[1].dim_value

    if args.pixel_bias_value is not None and len(args.pixel_bias_value) == 1:
        pixel_bias_value = [
            float(n)
            for n in args.pixel_bias_value[0]
            .replace("[", "")
            .replace("]", "")
            .split(",")
        ]

    if args.pixel_scale_value is not None and len(args.pixel_scale_value) == 1:
        pixel_scale_value = [
            float(n)
            for n in args.pixel_scale_value[0]
            .replace("[", "")
            .replace("]", "")
            .split(",")
        ]

    if i_n.type.tensor_type.shape.dim[1].dim_value != len(
        pixel_bias_value
    ) or i_n.type.tensor_type.shape.dim[1].dim_value != len(pixel_scale_value):
        raise ValueError(
            "--pixel-bias-value ("
            + str(pixel_bias_value)
            + ") and --pixel-scale-value ("
            + str(pixel_scale_value)
            + ") should be same as input dimension:"
            + str(i_n.type.tensor_type.shape.dim[1].dim_value)
        )
    other.add_bias_scale_bn_after(
        g, i_n.name, pixel_bias_value, pixel_scale_value
    )

# Change input and output shapes as requested
if args.input_change is not None:
    other.change_input_shape(g, args.input_change)
if args.output_change is not None:
    other.change_output_shape(g, args.output_change)

# Cutting nodes according to the given arguments.
if args.cut_node is not None or args.cut_type is not None:
    if args.cut_node is None:
        other.remove_nodes(g, cut_types=args.cut_type)
    elif args.cut_type is None:
        other.remove_nodes(g, cut_nodes=args.cut_node)
    else:
        other.remove_nodes(g, cut_nodes=args.cut_node, cut_types=args.cut_type)
    other.topological_sort(g)

# Rename nodes
if args.rename_output:
    if len(args.rename_output) % 2 != 0:
        print("Rename output should be paires of names.")
    else:
        for i in range(0, len(args.rename_output), 2):
            other.rename_output_name(
                g, args.rename_output[i], args.rename_output[i + 1]
            )

# Remove useless nodes
if (
    args.delete_node
    or args.delete_input
    or args.input_change
    or args.output_change
):
    # If shape changed during the modification, redo shape inference.
    while len(g.value_info) > 0:
        g.value_info.pop()
passes = ["extract_constant_to_initializer"]
m = optimizer.optimize(m, passes)
g = m.graph
replacing.replace_initializer_with_Constant(g)
other.topological_sort(g)
# Polish and output
m = other.polish_model(m)
other.add_output_to_value_info(m.graph)
onnx.save(m, args.out_file)
