#!/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.
#

"""
Replaces a subgraph in an onnx model with a plugin.
"""

from polygraphy import mod
from polygraphy.logger import G_LOGGER
from polygraphy.tools import Tool
from polygraphy.tools.args import (
    DataLoaderArgs,
    OnnxLoadArgs,
    ModelArgs,
    OnnxInferShapesArgs,
)
import os

# Your tool should lazily import any external dependencies. By doing so,
# we avoid creating hard dependencies on other packages.
# Additionally, this allows Polygraphy to automatically install required packages
# as they are needed, instead of requiring the user to do so up front.

gs = mod.lazy_import("onnx_graphsurgeon>=0.5.0")
onnx = mod.lazy_import("onnx")
yaml = mod.lazy_import("yaml", pkg_name="pyyaml")
common_backend = mod.lazy_import("polygraphy.backend.common")

def default_replace_with_plugin(graph, input_tensors: list, output_tensors: list, attrs=None, op=None):
    """
    replaces a subgraph (set of nodes) with a single plugin node
    default method to be used when the plugin does not specify a custom replacement method
    """

    def issubset_unhashable(list_a: list, list_b: list) -> bool:
        """
        Return whether list_a is a subset (or equal to) list_b
        The objects in list_a and list_b are unhashable, otherwise set(list_a) <= set(list_b) is enough
        """
        return len(list_a) <= len(list_b) and all(a in list_b for a in list_a)

    # Disconnect those output nodes of the input tensors whose inputs are a subset of the input tensors
    for in_tensor in input_tensors:
        to_remove_nodes = []
        for out_node in in_tensor.outputs:
            if issubset_unhashable(out_node.inputs, input_tensors):
                to_remove_nodes.append(out_node)
        for node in to_remove_nodes:
            in_tensor.outputs.remove(node)

    # Disconnet input nodes of all output tensors
    for out_tensor in output_tensors:
        to_remove_nodes = []
        for in_node in out_tensor.inputs:
            to_remove_nodes.append(in_node)
        for node in to_remove_nodes:
            out_tensor.inputs.remove(node)

    # Insert the new node
    new_node = graph.layer(op=op, inputs=input_tensors, outputs=output_tensors, attrs=attrs)
    graph.cleanup().toposort()
    return new_node[0].inputs[0]

class Replace(Tool):
    # Polygraphy will use the docstring of the tool child class to generate
    # the summary for the command-line help output.
    """
    Replace a subgraph in an onnx model with a plugin.
    """
    GRAPH_PATTERN_FILE_NAME="pattern.py"

    def __init__(self):
        super().__init__(name="replace")

    def get_subscriptions_impl(self):
        return [
            ModelArgs(model_opt_required=True),
            OnnxInferShapesArgs(),
            OnnxLoadArgs(),
            DataLoaderArgs(),
        ]

    def add_parser_args_impl(self, parser):
        parser.add_argument("--plugin-dir", help="Plugin directory.", required=True)
        parser.add_argument(
            "-o", "--output", help="Where to save the modified model", required=False
        )
        parser.add_argument("--config", help="location of config.yaml.")

    def run_impl(self, args):
        self.replace_plugin(
            model_file=args.model_file,
            plugin_dir=args.plugin_dir,
            output=args.output,
            config=args.config
        )
    def replace_plugin(self, model_file, plugin_dir, output=None, config=None):
        graph = gs.import_onnx(self.arg_groups[OnnxLoadArgs].load_onnx()) if self.arg_groups else gs.import_onnx(onnx.load(model_file))

        tensor_map = graph.tensors()
        config_yaml = config or os.path.join(os.path.dirname(model_file), "config.yaml")            
        
        plugin_dir = os.path.abspath(plugin_dir)
        
        with open(config_yaml, "r") as stream:
            in_yaml = yaml.safe_load_all(stream)

            for plugin in in_yaml:
                plugin_name = plugin["name"]
                plugin_op = plugin["op"]
                G_LOGGER.ultra_verbose(f"replacing {plugin_name}...")
                plugin_pattern_loc = os.path.join(plugin_dir, plugin_name, self.GRAPH_PATTERN_FILE_NAME)
                # if the plugin provides a custom replacement method, use that                    
                replace_fn = default_replace_with_plugin
                try:
                    replace_fn = mod.import_from_script(
                        plugin_pattern_loc,
                        "replace_with_plugin"
                    )
                except:
                    pass

                replace_cnt = 0
                for instance in plugin["instances"]:
                    attrs = instance.get("attributes", None)
                    if replace_fn(
                            graph=graph,
                            input_tensors=[tensor_map[ip_tensor_name] for ip_tensor_name in instance["inputs"]],
                            output_tensors=[tensor_map[op_tensor_name] for op_tensor_name in instance["outputs"]],
                            attrs=attrs,
                            op=plugin_op
                    ):
                        replace_cnt += 1
                G_LOGGER.info(f"replaced {replace_cnt} instances of {plugin_name} plugin")
                if replace_cnt != len(plugin['instances']):
                    G_LOGGER.warning(f"Warning: not all instances of {plugin_name} were replaced!")

        output_onnx = output or os.path.join(os.path.dirname(model_file), "replaced.onnx")

        onnx.save(gs.export_onnx(graph), output_onnx)