#!/usr/bin/env python3
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import onnx_graphsurgeon as gs
import argparse
import onnx
import numpy as np
import torch


# Pad layer subgraph structure in ONNX (specific to opset 11):
#               Constant
#                  |
#                Shape
#                  |
#         Mul   Gather
#          \     /
#            Sub
#             |
#       ConstantOfShape
#             |
#          Concat
#             |
#          Reshape
#             |
#           Slice
#             |
#          Transpose
#             |
#          Reshape
#             |
#    Input  Cast  Constant
#       \     |    /
#            Pad
def process_pad_nodes(graph):
    """
    Fold the pad subgraph into a single layer with pad values as input
      Input
       |
      Pad
       |
      Conv
    """
    pad_nodes = [node for node in graph.nodes if node.op == "Pad"]
    for node in pad_nodes:
        fold_pad_inputs(node, graph)

    return graph


def fold_pad_inputs(node, graph):
    # Gather the amount of padding in each dimension from pytorch graph.
    if torch.__version__ < "1.5.0":
        pad_values_pyt = (
            node.i(1).i(0).i(0).i(0).i(0).i(0).i(0).i(0).attrs["value"].values
        )
    elif torch.__version__ < "2.0.0":
        pad_values_pyt = node.i(1).i(0).i(0).i(0).i(0).i(0).inputs[0].values
    else:
        pad_values_pyt = node.i(1).i(0).i(0).i(0).i(0).i(0).i(0).attrs["value"].values

    # Assumption a 4d input tensor
    onnx_pad_values = [0] * 4 * 2  # 4d tensor and 2 sides padding for each dimension
    j = 3
    for i in range(0, len(pad_values_pyt), 2):
        onnx_pad_values[j] = pad_values_pyt[i]
        onnx_pad_values[j + 4] = pad_values_pyt[i + 1]
        j -= 1

    # Change the existing pad tensor to the new onnx_pad values tensor
    pads_folded_tensor = gs.Constant(
        name=node.inputs[1].name, values=np.array(onnx_pad_values)
    )
    node.inputs[1] = pads_folded_tensor


# Pytorch-exported Upsample structure in ONNX:
#        Mul        Mul
#         |          |
#        Cast       Cast
#         |          |
#        Floor      Floor
#         |          |
#      Unsqueeze  Unsqueeze
#         \         /
#           Concat
#             |
#            Cast    Cast
#              \      /
#                Div
#                 |
#     Input     Concat
#       \         /
#         Upsample
def process_upsample_nodes(graph, opset=11):
    """
    Replace the upsample structure with structure below
      Conv   scale_factor
       |      /
      Upsample
       |
      ReLU
    """
    if opset >= 11:
        upsample_layer_name = "Resize"
    else:
        upsample_layer_name = "Upsample"

    upsample_nodes = [node for node in graph.nodes if node.op == upsample_layer_name]
    for node in upsample_nodes:
        fold_upsample_inputs(node, graph, opset)

    return graph


def fold_upsample_inputs(upsample, graph, opset=11):
    """
    Inplace transformation of the graph. The upsample subgraph is collapsed
    to single upsample node with input and scale factor (constant tensor).
    Args:
        upsample: upsample node in the original graph.
        graph: graph object.
    """

    if opset == 9:
        # Gather the scale factor from mul op in the upsample input subgraph
        scale_factor = (
            upsample.i(1).i(1).i(0).i(0).i(0).i(0).i(0).i(0).i(1).attrs["value"].values
        )

        # Create the new scales tensor
        scales = np.array([1.0, 1.0, scale_factor, scale_factor], dtype=np.float32)
        scale_tensor = gs.Constant(name=upsample.inputs[-1].name, values=scales)

        # Change the last input to the node to the new constant scales tensor.
        upsample.inputs[-1] = scale_tensor
    else:
        # In opset 11, upsample layer is exported as Resize. We will transform this Resize layer into an Upsample layer
        # and collapse the input
        sizes_tensor_name = upsample.inputs[3].name

        # Create the new scales tensor
        scale_factor = (
            upsample.i(3).i(1).i().i().i().i().i(0).i(1).attrs["value"].values
        )
        scales = np.array([1.0, 1.0, scale_factor, scale_factor], dtype=np.float32)
        scale_tensor = gs.Constant(name=sizes_tensor_name, values=scales)

        # Rename the Resize op to upsample and add the data and scales as inputs to the upsample layer.
        input_tensor = upsample.inputs[0]
        upsample.inputs = [input_tensor, scale_tensor]
        upsample.op = "Upsample"


# Pytorch-exported GroupNorm subgraph in ONNX:
# Conv
#   |
# Reshape    Scale    Bias
#     \       |       /
# InstanceNormalization
#         |
#      Reshape    Unsqueeze
#          \      /
#        Mul (scale)   Unsqueeze
#           \         /
#           Add (bias)
#              |
#            ReLU
def process_groupnorm_nodes(graph):
    """
    Gather the instance normalization nodes and the rest of the subgraph
    and convert into a single group normalization node.
    """
    instancenorms = [node for node in graph.nodes if node.op == "InstanceNormalization"]
    for node in instancenorms:
        convert_to_groupnorm(node, graph)

    return graph


def retrieve_attrs(instancenorm):
    """
    Gather the required attributes for the GroupNorm plugin from the subgraph.
    Args:
        instancenorm: Instance Normalization node in the graph.
    """
    attrs = {}
    # The 2nd dimension of the Reshape shape is the number of groups
    attrs["num_groups"] = instancenorm.i().i(1).attrs["value"].values[1]
    attrs["eps"] = instancenorm.attrs["epsilon"]

    # 1 is the default plugin version the parser will search for, and therefore can be omitted,
    # but we include it here for illustrative purposes.
    attrs["plugin_version"] = "1"

    # "" is the default plugin namespace the parser will use, included here for illustrative purposes
    attrs["plugin_namespace"] = ""

    return attrs


def convert_to_groupnorm(instancenorm, graph):
    """
    Convert the Pytorch-exported GroupNorm subgraph to the subgraph below
    Conv
      |
    GroupNorm
      |
    ReLU
    Attributes:
        instancenorm: Instance Normalization node in the graph.
        graph: Input graph object
    """
    # Retrieve the instancenorm attributes and create the replacement node
    attrs = retrieve_attrs(instancenorm)
    groupnorm = gs.Node(op="GroupNormalizationPlugin", attrs=attrs)
    graph.nodes.append(groupnorm)

    # The plugin needs to receive an input from the Conv node, and output to the ReLU node
    conv_output_tensor = instancenorm.i().inputs[0]  # Output of Conv
    relu_input_tensor = instancenorm.o().o().o().outputs[0]  # Output of Add

    # Reconnect inputs/outputs to the groupnorm plugin
    conv_output_tensor.outputs[0] = groupnorm
    relu_input_tensor.inputs[0] = groupnorm

    # Add scale and bias constant tensors to group norm plugin
    if torch.__version__ < "1.5.0":
        groupnorm.inputs.append(instancenorm.o().o().i(1).inputs[0])
        groupnorm.inputs.append(instancenorm.o().o().o().i(1).inputs[0])
    else:
        groupnorm.inputs.append(instancenorm.o().o().inputs[1])
        groupnorm.inputs.append(instancenorm.o().o().o().inputs[1])