# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
"""Pipeline executor that executes a series of modules in a pipeline fashion."""
import json
import os
import tvm._ffi
from tvm import relay
from tvm.relay.transform import InferType
from tvm.contrib.pipeline_executor import PipelineExecutorFactoryModule


def pipeline_executor_build_enabled():
    """Check if the pipeline executor build is enabled.

    Return
    -------
    enable: bool
        Return whether the pipeline executor is enabled.
    """
    return tvm.contrib.pipeline_executor.pipeline_executor_enabled()


def build(pipe_configs):
    """Build modules used in the pipeline executor, then use these modules and configuration
    to create a pipeline executor.

    Parameters
    ----------
    pipe_configs: PipelineConfig
        Build Configuration information.

    Returns
    -------
    ret: PipelineExecutorFactoryModule
        Common interface for pipeline executor factory modules.
    """
    libs = {}
    config = pipe_configs.get_config()
    if "module_connection" not in config:
        raise RuntimeError('"module_connection" is missing')
    if "input_connection" not in config:
        raise RuntimeError('"input_connection" is missing')
    if "param_connection" not in config:
        raise RuntimeError('"param_connection" is missing')

    mod_n_configs = config["module_connection"]
    config_len = len(mod_n_configs)
    module_string_config = [{} for _ in range(config_len)]
    # Use hardware configurations to build backend modules for each subgraph.
    for ir_mod, mod_config in mod_n_configs.items():
        pipe_config = mod_config["pipeline"].copy()
        mod_idx = pipe_config["mod_idx"]
        dev = mod_config["dev"]
        target = mod_config["target"]
        build_func = relay.build
        # Callers may need to use a customized building function to wrap the pre-building logic
        # and the backend building logic. For example, in order to support a backend which only
        # can do "int8" computation, the caller may need to merge the "quantization" logic
        # into the building logic to creat a customized building function.
        if "build" in mod_config and mod_config["build"]:
            build_func = mod_config["build"]

        lib = build_func(
            ir_mod,
            target,
            params=mod_config["params"],
            target_host=mod_config["target_host"],
            mod_name=mod_config["mod_name"],
        )

        pipe_config["dev"] = "{},{}".format(dev.device_type, dev.device_id)
        # Use "mod_idx" as the key to create a "module_connection" map which is not only
        # for the module index but also for the module connection used to build the pipeline.
        module_string_config[mod_idx] = pipe_config
        libs[mod_idx] = {
            "lib": lib,
            "dev": dev,
            "fcompile": mod_config["fcompile"],
            "export_cc": mod_config["export_cc"],
        }

    # Creating a text form configuration to record the "input_connection" and the
    # "module_connection" information. The "input_connection" is used to record the
    # map of global input and subgraph input, and the "module_connection" is used to
    # record module dependency.
    string_config = {}
    string_config["param_connection"] = config["param_connection"]
    string_config["input_connection"] = config["input_connection"]
    string_config["module_connection"] = module_string_config

    return PipelineExecutorFactoryModule(libs, string_config)


def export_library(factory, directory_path):
    """Export the pipeline executor into disk files.

    Parameters
    ----------
    factory : PipelineExecutorFactoryModule
        The pipeline executor factory
    directory_path : str
        Export the files to this directory.
    """
    if not factory.pipeline_mods:
        raise RuntimeError("The pipeline executor has not been initialized.")

    # Check if the directory_path exists.
    if not directory_path or not os.path.exists(directory_path):
        raise RuntimeError("The directory {directory_path} does not exist.")
    # Create an load configuration.
    load_config_file_name = "{}/load_config".format(directory_path)
    pipeline_config_file_name = "{}/pipeline_config".format(directory_path)
    config = {}
    config["load_config"] = load_config_file_name
    config["pipeline_config"] = pipeline_config_file_name
    load_config = []
    # Export the library, JSON, and parameter into files, then export these files path
    # into a configuration file.
    for lib_index in factory.pipeline_mods:
        mconfig = {}
        mconfig["mod_idx"] = lib_index
        mconfig["lib_name"] = "{}/lib{}.so".format(directory_path, lib_index)
        mconfig["json_name"] = "{}/json{}".format(directory_path, lib_index)
        mconfig["params_name"] = "{}/params{}".format(directory_path, lib_index)
        lib_config = factory.pipeline_mods[lib_index]
        mconfig["dev"] = "{},{}".format(lib_config["dev"].device_type, lib_config["dev"].device_id)
        fcompile = lib_config["fcompile"]
        if not fcompile:
            fcompile = False

        # Get the graph, lib, and parameters from GraphExecutorFactoryModule.
        lib = factory.pipeline_mods[lib_index]["lib"]
        # Export the lib, graph, and parameters to disk.
        lib.export_library(mconfig["lib_name"], fcompile)
        with open(mconfig["json_name"], "w") as file_handle:
            file_handle.write(lib.graph_json)
        with open(mconfig["params_name"], "wb") as file_handle:
            file_handle.write(relay.save_param_dict(lib.params))

        load_config.append(mconfig)

    with open(load_config_file_name, "w") as file_handle:
        json.dump(load_config, file_handle)

    with open(pipeline_config_file_name, "w") as file_handle:
        json.dump(factory.mods_config, file_handle)

    config_file_name = "{}/config".format(directory_path)
    with open(config_file_name, "w") as file_handle:
        json.dump(config, file_handle)

    return config_file_name


class PipelineConfig(object):
    """Pipeline configuration information, this class contains the DAG that expresses
    the dependency of each module involved in a pipeline and the parameters for building
    each module.
    """

    class Binding:
        """This class defines the module connections information.
        The type can only be "input" or "output".

        Parameters
        ----------
        owner : ModuleWrapper
            The class who owns this interface.

        io_type : str
            The I/O type of this interface. It can only be "input" or "output".

        name : str/integer
            Name, for input it is string such as "data0", for output it is an integer such as 0.

        data_type: TensorType
            The data type of this interface.
        """

        def __init__(self, owner, io_type, name, data_type=None):
            self.io_owner = owner
            self.io_type = io_type
            self.name = str(name)
            # Child interfaces that depend on this interface.
            self.bindings = []
            # Parents interfaces that this interface depend on.
            self.parents = []

            self.data_type = data_type

        def get_name(self):
            # Return name of this interface and the name of owner who owns this interface.
            owner_name = ""
            if isinstance(self.io_owner, PipelineConfig.ModuleWrapper):
                owner_name = self.io_owner.name

            return owner_name, self.name

        def get_owner_idx(self):
            # If the owner is ModuleWrapper return the owner index, if not return 0.
            if isinstance(self.io_owner, PipelineConfig.ModuleWrapper):
                return self.io_owner.idx

            return -1

        def is_pipeline_executor_interface(self):
            """The pipeline interface is used to interact with the caller. There are two types
            of interfaces, one is 'input' another is 'output'. The pipeline input interface
            is responsible for passing parameters to the internal module interface, and the
            pipeline output interface is responsible for outputting the results computed by
            the pipeline executor to the caller.
            """
            return not isinstance(self.io_owner, PipelineConfig.ModuleWrapper)

        def __repr__(self):
            # Geting the binding information in the form of text.
            str_format = "  |{}: ".format(self.name)
            for binding in self.bindings:
                mname, dname = binding.get_name()
                str_format += "{0}:{1} ".format(mname, dname)

            return str_format

        def check_binding_dict(self, connection_dict):
            """Checking the binding dictionary.
            Parameter
            ---------
            connection_dict : Dict[str, Any]
                It is a dictionary of module connections.
            """
            if "interface_name" not in connection_dict:
                raise RuntimeError('"inteface_name" is missing in global config!"')
            if "connection" not in connection_dict:
                raise RuntimeError(f'"connection" is missing!"')
            # The global interface mapping should be one-to-one.
            if not connection_dict["connection"]:
                raise RuntimeError("The global interface map is empty!")
            if len(connection_dict["connection"]) > 1:
                raise RuntimeError("A global interface maps multiple module interfaces!")
            if "mod_idx" not in connection_dict["connection"][0]:
                raise RuntimeError('"mod_idx" is missing!')

        def get_binding_dict(self):
            """Returning the binding information in the form of dictionary.
            Returns
            -------
            data : Dict[str, Any]
                The binding information is in the form of dictionary.
            """
            dict_format = {"interface_name": self.name, "connection": []}
            for binding in self.bindings:
                _, dname = binding.get_name()
                midx = binding.get_owner_idx()
                dict_format["connection"].append({"mod_idx": midx, "interface_name": dname})

            self.check_binding_dict(dict_format)
            return dict_format

        def check_dag_acyclic(self, start, inputs):
            """This is to check whether the DAG containing these input interfaces is acyclic.
            Parameters
            ----------
            start: ModuleWrapper
                The starting node of the cycle check algorithm.

            inputs: Binding
                These interfaces are used to connect to each other to build DAG.

            Return
            ------
                Return true if there is no cycle in the DAG.
            """
            for binding in inputs.values():
                if start == binding.io_owner:
                    return False
                for p in binding.parents:
                    if not self.check_dag_acyclic(start, p.io_owner.input_bindings.bindings):
                        return False

            return True

        def connect(self, binding):
            """Connect the current interface to the destination interface.
            Correct connections are as follows: 1. the pipeline input connected to a module input,
            2. the module output connected to a pipeline output, 3. the module output connected to
            a module input.

            Parameters
            ----------
            binding: Binding
                The destination of this connection.
            """

            # Check whether the binding setting is correct or not.
            if self.io_owner == binding.io_owner:
                raise RuntimeError("Can not bind itself.")

            if self.io_type == "param" and not self.is_pipeline_executor_interface():
                raise RuntimeError(
                    'The "param" binding can only be used by a pipeline executor interface!'
                )

            if not self.is_pipeline_executor_interface() and self.io_type == "input":
                raise RuntimeError("Module can only bind from output interface!")

            if self.io_type == "param" and binding.io_type != "param":
                raise RuntimeError(
                    'A global "param" interface can only be bind with a module "param" interface!'
                )

            if (
                not self.is_pipeline_executor_interface()
                and not binding.is_pipeline_executor_interface()
                and binding.io_type == "output"
            ):
                raise RuntimeError("Can not bind module output with another module output!")

            if (
                not self.is_pipeline_executor_interface()
                and binding.is_pipeline_executor_interface()
                and binding.io_type == "input"
            ):
                raise RuntimeError("Can not bind module output with pipeline input!")

            if self.is_pipeline_executor_interface() and self.io_type == "output":
                raise RuntimeError("Global output can not be used as binding start point.")

            if (
                self.is_pipeline_executor_interface()
                and self.io_type == "input"
                and binding.io_type != "input"
            ):
                raise RuntimeError("Global input can only bind with module input.")

            self.bindings.append(binding)
            if not self.is_pipeline_executor_interface():
                # Check whether the data types of the source and destination are the same.
                if (
                    isinstance(binding.io_owner, PipelineConfig.ModuleWrapper)
                    and self.data_type != binding.data_type
                ):
                    raise RuntimeError(
                        f"Illegal type (%s vs. %s): binding type is not same!"
                        % (self.data_type, binding.data_type)
                    )

                binding.parents.append(self)

                # Do acyclic check after increasing the in-degree of child node by setting
                # current interface as a parent of the child node.

                if not self.check_dag_acyclic(
                    binding.io_owner, self.io_owner.input_bindings.bindings
                ):
                    raise RuntimeError("Illegal connection: Cause a cycle!")

    class BindingList:
        """Container for bindings(input or output interface).

        Parameters
        ----------
        owner : ModuleWrapper/PipelineConfig
            The owner of this class can be ModuleWrapper or PipelineConfig.

        io_type : str
            The type of this class can be "input" or "output".
        """

        def __init__(self, owner, io_type):
            self.bindings = {}
            self.io_owner = owner
            self.binding_type = io_type

        def get_binding_data_type(self, key):
            if isinstance(self.io_owner, PipelineConfig.ModuleWrapper):
                return self.io_owner.get_data_type(key, self.binding_type)
            return None

        def __getitem__(self, key):
            if key not in self.bindings:
                data_type = self.get_binding_data_type(key)
                if not data_type and isinstance(self.io_owner, PipelineConfig.ModuleWrapper):
                    raise RuntimeError(f"Can not find {key} in binding list {self.binding_type}.")

                self.bindings[key] = PipelineConfig.Binding(
                    self.io_owner, self.binding_type, key, data_type
                )

            return self.bindings[key]

    class ModuleWrapper:
        """This class is a wrapper representing the module and contains information such as
        module information, binding information and building information.
        """

        def __init__(self, mod=None):
            self.target_host = None
            self.build_func = None
            self.params = None
            self.target = None
            self.fcompile = None
            self.name = None
            self.dev = None
            self.export_cc = None
            self.cpu_affinity = ""
            self.idx = None
            self.mod = mod
            self.input_params = InferType()(mod)["main"].params
            self.output_type = InferType()(mod)["main"].checked_type.ret_type
            self.input_bindings = PipelineConfig.BindingList(self, "input")
            self.output_bindings = PipelineConfig.BindingList(self, "output")
            self.param_binding = PipelineConfig.Binding(self, "param", "param")

        def __eq__(self, other):
            if isinstance(other, PipelineConfig.ModuleWrapper):
                return self.mod == other.mod

            return False

        def __getitem__(self, key):
            if isinstance(key, str):
                if key == "input":
                    return self.input_bindings

                if key == "output":
                    return self.output_bindings

                if key == "param":
                    return self.param_binding

                raise RuntimeError(f"{key} not found!")

            raise RuntimeError('The data type of "key" is not supported!')

        def get_data_type(self, key, interface_type):
            """Get the module interface data type according to the key value and interface type.
            Parameters
            ----------
            key: str
                The interface name.

            interface_type:
                The interface type.

            Return
            -------
                Return data type.
            """
            if interface_type == "input":
                for param in self.input_params:
                    if param.name_hint == key:
                        return param._checked_type_

            if interface_type == "output":
                if isinstance(self.output_type, tvm.ir.type.TupleType):
                    if int(key) < len(self.output_type.fields):
                        return self.output_type.fields[int(key)]
                elif int(key) == 0:
                    return self.output_type

            return None

        def set_idx_name(self, idx):
            # Set the index value and generate the module name.
            self.idx = idx
            self.name = "mod{}".format(str(idx))

        def is_root_mod(self):
            """Check whether this node is the root node in DAG, this function is used
            in topological sort.
            """
            return all([not b.parents for b in self.input_bindings.bindings.values()])

        def remove_self_from_bindings(self):
            """Remove the current node from child dependencies to reduce the in-degree
            of child node, this function is used in topological sort.
            """
            for binding in self.output_bindings.bindings.values():
                for child in binding.bindings:
                    if binding in child.parents:
                        child.parents.remove(binding)

    def __init__(self):
        self.mod_wrapper = {}
        self.input_bindings = self.BindingList(self, "input")
        self.output_bindings = self.BindingList(self, "output")
        # There is a map of global parameters group and module index.
        self.param_group_bindings = self.BindingList(self, "param")

    def __str__(self):
        # Get configuration information as a string.

        # Use topological sort to get correct module order.
        self.dag_topology_sort()
        # Getting the parameters dependencies.
        param_dump = "Params\n"
        for param_name in self.param_group_bindings.bindings:
            inf = self.param_group_bindings.bindings[param_name]
            param_dump += str(inf) + "\n"
        # Get the input dependencies.
        input_dump = "\nInputs\n"
        for input_name in self.input_bindings.bindings:
            inf = self.input_bindings.bindings[input_name]
            input_dump += str(inf) + "\n"

        # Get the connections information of each module.
        output = {}
        connections_dump = "\nconnections\n"
        for mod in self.mod_wrapper:
            for interface in self.mod_wrapper[mod].output_bindings.bindings.values():
                if interface.bindings:
                    mname, dname = interface.get_name()
                    iname = mname + ".output(" + dname + ")->"
                    for dep in interface.bindings:
                        dep_mname, dep_dname = dep.get_name()
                        if isinstance(dep.io_owner, PipelineConfig.ModuleWrapper):
                            iname += f" {dep_mname}.{dep_dname}"
                            connections_dump += f"  |{iname}\n"
                        else:
                            output[dep_dname] = f"{mname}.output({dname})"

        # Get the output dependencies.
        output_dump = "\noutput\n"
        for name in sorted(output.keys()):
            output_dump += f"  |output({name}) : {output[name]}\n"

        return param_dump + input_dump + output_dump + connections_dump

    def __getitem__(self, key):
        if isinstance(key, tvm.ir.module.IRModule):
            if key not in self.mod_wrapper:
                self.mod_wrapper[key] = self.ModuleWrapper(key)
            return self.mod_wrapper[key]

        if isinstance(key, str):
            if key == "input":
                return self.input_bindings
            if key == "output":
                return self.output_bindings
            if key == "param_group":
                return self.param_group_bindings

            raise RuntimeError(f"{key} not found!")

        raise RuntimeError(f'The key type "{type(key)}" is not supported!')

    def get_config(self):
        """Get the configuration information in dictionary form, this configuration
        will be used to create pipeline executor.
        """

        # Use topological sort to get the correct order of modules.
        self.dag_topology_sort()
        mconfig = {}
        module_connection = {}
        for mod in self.mod_wrapper:
            # Generate pipeline configuration.
            mconf = {}
            output_conf = []
            module = self.mod_wrapper[mod]
            for _, binding in module.output_bindings.bindings.items():
                dep_conf = []
                output = {}
                if binding.bindings:
                    for dep in binding.bindings:
                        dep_item = {}
                        _, dname = dep.get_name()
                        if dep.is_pipeline_executor_interface():
                            dep_item["global_output_index"] = int(dname)
                        else:
                            dep_item["mod_idx"] = dep.get_owner_idx()
                            dep_item["input_name"] = dname
                        dep_conf.append(dep_item)

                # The value of output_idx start from 0.
                output["output_idx"] = int(binding.name)
                output["dependencies"] = dep_conf
                output_conf.append(output)

            mconf["mod_idx"] = module.idx
            mconf["cpu_affinity"] = module.cpu_affinity
            mconf["output"] = output_conf

            module_connection[mod] = {
                "pipeline": mconf,
                "target_host": module.target_host,
                "mod_name": "default",
                "build": module.build_func,
                "params": module.params,
                "target": module.target,
                "fcompile": module.fcompile,
                "dev": module.dev,
                "export_cc": module.export_cc,
            }

        # Creating a map including pipeline inputs and subgraph inputs.
        input_connection = []
        for input_name in self.input_bindings.bindings:
            input_dict = self.input_bindings.bindings[input_name].get_binding_dict()
            if "interface_name" not in input_dict["connection"][0]:
                raise RuntimeError("interface_name is missing in connection config!")
            # Creating the map including global interfaces and subgraph interfaces.
            input_map = {
                "global_interface_name": input_dict["interface_name"],
                "mod_idx": input_dict["connection"][0]["mod_idx"],
                "module_interface_name": input_dict["connection"][0]["interface_name"],
            }
            input_connection.append(input_map)

        # Create a map including global parameters groups and modules.
        param_connection = []
        for param_name in self.param_group_bindings.bindings:
            param_dict = self.param_group_bindings.bindings[param_name].get_binding_dict()
            param_map = {
                "global_param_name": param_dict["interface_name"],
                "mod_idx": param_dict["connection"][0]["mod_idx"],
            }
            param_connection.append(param_map)

        mconfig["module_connection"] = module_connection
        mconfig["input_connection"] = input_connection
        mconfig["param_connection"] = param_connection
        return mconfig

    def dag_topology_sort(self):
        """Use topological sort to get order of pipeline modules."""
        mlist = []
        mod_wrapper = self.mod_wrapper.copy()
        while mod_wrapper:
            temp_list = []
            for mod, wrapper in mod_wrapper.items():
                if wrapper.is_root_mod():
                    temp_list.append(mod)
                    wrapper.remove_self_from_bindings()

            for mod in temp_list:
                mod_wrapper.pop(mod, None)

            mlist += temp_list

        mod_wrapper_sort = {}
        for mod, i in zip(mlist, range(len(mlist))):
            self.mod_wrapper[mod].set_idx_name(i)
            mod_wrapper_sort[mod] = self.mod_wrapper[mod]

        self.mod_wrapper = mod_wrapper_sort

    def get_mod_idx(self, mod):
        # Return the module index.
        idx = self.mod_wrapper[mod].idx
        return idx

    def pipe_input(self, name):
        # Return the input interface according to the name.
        return self.input_bindings[name]

    def pipe_output(self, idx):
        # Return the output interface according to the name.
        return self.output_bindings[idx]