# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import collections
import heapq
from collections import OrderedDict
from typing import Dict, List, Tuple, Union

import numpy as np

from ..core import _imperative_rt
from ..core._imperative_rt import GraphProfiler
from ..core._imperative_rt import OperatorNode as _OpNode
from ..core._imperative_rt import VarNode as _VarNode
from ..core.tensor import megbrain_graph as G
from ..core.tensor.megbrain_graph import set_priority_to_id
from ..tensor import Tensor

__all__ = [
    "get_dep_vars",
    "get_owner_opr_inputs",
    "get_owner_opr_type",
    "get_opr_type",
    "graph_traversal",
    "get_oprs_seq",
    "replace_vars",
    "replace_oprs",
    "set_priority_to_id",
    "GraphInference",
]


def get_dep_vars(
    var: Union[_VarNode, List[_VarNode]], var_type: Union[str, List[str]] = None
) -> List[_VarNode]:
    r"""Returns :class:`.tensor.core.megbrain_graph.VarNode` of type ``var_type`` that input ``var``
    depands on. If ``var_type`` is None, returns all types.
    """
    outputs = []
    memo = set()

    if isinstance(var, _VarNode):
        var = [var]

    if isinstance(var_type, str):
        var_type = [var_type]

    q = list(var)
    while q:
        v = q.pop(0)
        if v in memo:
            continue
        memo.add(v)
        q.extend(get_owner_opr_inputs(v))
        if var_type is not None:
            if get_owner_opr_type(v) in var_type:
                outputs.append(v)
        else:
            outputs.append(v)

    return outputs


def get_owner_opr_inputs(var: _VarNode) -> List[_VarNode]:
    r"""Gets the inputs of owner opr of a variable. """
    return var.owner.inputs


def get_owner_opr_type(var: _VarNode) -> str:
    r"""Gets the type of owner opr of a variable."""
    return var.owner.type


def get_opr_type(opr: _OpNode) -> str:
    r"""Gets the type of an opr."""
    assert isinstance(opr, _OpNode)
    return opr.type


class _OprStableOrderHeapq:
    r"""heap implementation for operator comparison in stable order"""

    _list = None
    _extra_priority = None
    _used_id_name_pairs = None

    def __init__(self, extra_priority):
        assert isinstance(extra_priority, collections.Callable)
        self._list = []
        self._extra_priority = extra_priority
        self._used_id_name_pairs = {}

    def pop_min(self):
        return heapq.heappop(self._list)[-1]

    def add(self, opr):
        # named as add to mimic set() interface

        id_ = opr.id
        name = opr.name

        other = self._used_id_name_pairs.setdefault((id_, name), opr)
        if other is not opr:
            raise RuntimeError(
                "duplicated (id, name) pair: opr0={} opr1={}".format(other, opr)
            )

        item = self._extra_priority(opr) + (id_, name, opr)
        heapq.heappush(self._list, item)

    def __bool__(self):
        return bool(self._list)


def graph_traversal(outputs: _VarNode):
    r"""Helper function to traverse the computing graph and return enough useful information.

    Args:
        outputs: model outputs.

    Returns:
        tuple (map_oprs, map_vars, var2oprs, opr2receivers, indegree2opr, opr2indegree)

        WHERE

        * map_oprs is dict from opr_id to actual opr
        * map_vars is dict from var_id to actual var
        * var2oprs is dict from var to dest oprs along with index
        * opr2receivers is dict from current opr to next opr
        * indegree2opr is dict from in_degree to opr in computing graph
        * opr2indegree is dict from opr in computing graph to in_degree

        (indegree2opr, opr2indegree) are only used in topological sort in get_oprs_seq function
    """
    # meta information for comp graph
    map_oprs = collections.defaultdict(set)
    map_vars = collections.defaultdict(set)

    var2oprs = collections.defaultdict(list)
    opr2receivers = collections.defaultdict(list)
    queue = []
    [queue.append(o) for o in [x.owner for x in outputs] if o not in queue]
    visited = set(map(lambda x: x.id, queue))

    # iterate through whole comp_graph, fill in meta information
    indegree2opr = collections.defaultdict(set)
    indegree2opr[0] = _OprStableOrderHeapq(lambda op: (op.priority,))
    opr2indegree = {}

    idx = 0
    while idx < len(queue):
        cur_opr = queue[idx]
        map_oprs[cur_opr.id] = cur_opr

        idx += 1

        indegree = 0
        for var_idx, var in enumerate(cur_opr.inputs):
            map_vars[var.id] = var
            var2oprs[var.id].append((cur_opr.id, var_idx))

            pre_opr = var.owner

            if pre_opr.id not in visited:
                visited.add(pre_opr.id)
                queue.append(pre_opr)

            indegree += 1
            opr2receivers[pre_opr.id].append(cur_opr.id)
        opr = cur_opr if indegree == 0 else cur_opr.id
        indegree2opr[indegree].add(opr)
        opr2indegree[cur_opr.id] = indegree

    return map_oprs, map_vars, var2oprs, opr2receivers, indegree2opr, opr2indegree


def get_oprs_seq(
    outputs: List[_VarNode], prune_reshape=False, prune_immtensor=True
) -> List[_OpNode]:
    r"""Gets oprs in some topological order for a dumped model.

    Args:
        outputs: model outputs.
        prune_reshape: whether to prune the useless operators used by Reshape opr during inference.
        prune_immtensor: whether to prune the ImmutableTensor opr.

    Returns:
        opr list with some correct execution order.
    """

    def topological_sort(map_oprs, opr2receivers, indegree2opr, opr2indegree):
        # generate an execution order with topological sort algorithm
        oprs_seq = []
        nr_remain = len(map_oprs)
        while indegree2opr[0]:
            opr = indegree2opr[0].pop_min()
            opr_id = opr.id
            nr_remain -= 1
            if opr.type != "ImmutableTensor" or not prune_immtensor:
                oprs_seq.append(opr)

            for post_id in opr2receivers[opr_id]:
                indegree = opr2indegree[post_id]
                indegree2opr[indegree].remove(post_id)

                indegree -= 1
                if indegree == 0:
                    indegree2opr[indegree].add(map_oprs[post_id])
                else:
                    indegree2opr[indegree].add(post_id)
                opr2indegree[post_id] = indegree

        assert nr_remain == 0, "there are {} remaining nodes; cyclic graph?".format(
            nr_remain
        )
        return oprs_seq

    # reshape op definition: reshape(input_tensor, dest_shape) -> output_tensor
    # when inferencing, shape of output_tensor is already known, so one can prune some operators related to dest_shape in the loaded graph
    def prune_reshape_oprs(outputs, oprs_seq, var2oprs):
        def iterative_pruning(cur_opr, post_opr, marked_opr_ids, visited):
            useless = True
            for oup in cur_opr.outputs:
                if "workspace" not in oup.name:
                    var_idx = post_opr.inputs.index(oup)
                    var2oprs[oup.id].remove((post_opr.id, var_idx))
                    useless = useless and (len(var2oprs[oup.id]) == 0)

            if useless:
                marked_opr_ids.append(cur_opr.id)

                for opr in set([var.owner for var in cur_opr.inputs]):
                    if (opr.id, cur_opr.id) not in visited:
                        visited.add((opr.id, cur_opr.id))
                        iterative_pruning(opr, cur_opr, marked_opr_ids, visited)

        reshape_vars = get_dep_vars(outputs, "Reshape")
        reshape_oprs = [var.owner for var in reshape_vars]

        marked_opr_ids = []
        visited = set()
        for reshape_opr in reshape_oprs:
            iterative_pruning(
                reshape_opr.inputs[1].owner, reshape_opr, marked_opr_ids, visited
            )

        # filter out all marked oprs
        return list(filter(lambda x: x.id not in marked_opr_ids, oprs_seq))

    # adjust the order of oprs, let param/data privoder oprs close to the oprs which use them as inputs.
    def reorder_oprs_seq(oprs):
        rst = []
        param_or_data_provider_oprs = []
        other_oprs = []

        for o in oprs:
            if o.type in ["ImmutableTensor", "Host2DeviceCopy"]:
                param_or_data_provider_oprs.append(o)
            else:
                other_oprs.append(o)

        for o in other_oprs:
            for inp in o.inputs:
                if inp.owner.type in ["ImmutableTensor", "Host2DeviceCopy"]:
                    if inp.owner in param_or_data_provider_oprs:
                        rst.append(inp.owner)
                        param_or_data_provider_oprs.remove(inp.owner)
            rst.append(o)
        rst = rst + param_or_data_provider_oprs
        assert len(rst) == len(oprs)
        return rst

    map_oprs, _, var2oprs, opr2receivers, indegree2opr, opr2indegree = graph_traversal(
        outputs
    )
    oprs_seq = topological_sort(map_oprs, opr2receivers, indegree2opr, opr2indegree)
    oprs_seq = reorder_oprs_seq(oprs_seq)
    if prune_reshape is True:
        oprs_seq = prune_reshape_oprs(outputs, oprs_seq, var2oprs.copy())
    return oprs_seq


def replace_vars(
    dst: List[_VarNode], varmap: Dict[_VarNode, _VarNode]
) -> List[_VarNode]:
    r"""Replaces vars in the graph.

    Args:
        dst: target vars representing the graph.
        varmap: the map that specifies how to replace the vars.

    Returns:
        new vars that correspond to ``dst`` with all the dependencies replaced.
    """
    dst_vec = []
    repl_src_vec = []
    repl_dst_vec = []
    for i in dst:
        assert isinstance(i, _VarNode)
        dst_vec.append(i)

    for i, j in getattr(varmap, "items", lambda: varmap)():
        assert isinstance(i, _VarNode)
        assert isinstance(j, _VarNode)
        repl_src_vec.append(i)
        repl_dst_vec.append(j)

    return _imperative_rt.graph._replace_vars(repl_src_vec, repl_dst_vec, dst_vec)


def replace_oprs(dst: List[_VarNode], oprmap: Dict[_OpNode, _OpNode]) -> List[_VarNode]:
    """Replaces operators in the graph.

    Args:
        dst: target vars representing the graph.
        oprmap: the map that specifies how to replace the operators.

    Returns:
        new vars that correspond to ``dst`` with all the dependencies replaced.
    """
    dst_vec = []
    repl_src_vec = []
    repl_dst_vec = []
    for i in dst:
        assert isinstance(i, _VarNode)
        dst_vec.append(i)

    for i, j in getattr(oprmap, "items", lambda: oprmap)():
        assert isinstance(i, _OpNode)
        assert isinstance(j, _OpNode)
        repl_src_vec.append(i)
        repl_dst_vec.append(j)

    return _imperative_rt.graph._replace_oprs(repl_src_vec, repl_dst_vec, dst_vec)


def find_vars_by_name(dst: List[_VarNode], names: List[str]) -> List[_VarNode]:
    r"""Gets VarNode list by names in the graph.

    Args:
        dst: target vars representing the graph.
        names: name list for target VarNode.

    Returns:
        results found by names.
    """
    output_names = names.copy()
    all_vars = get_dep_vars(dst) + dst
    # use dict to keep outputs order the same as names.
    output_dict = {}
    for i in all_vars:
        if i.name in output_names:
            output_dict[i.name] = i
            output_names.remove(i.name)
    assert len(output_names) == 0, "Can not find varnode {} in this model".format(
        output_names
    )
    return [output_dict[i] for i in names]


def convert_inputs(
    dst: List[_VarNode], inputs: List[_VarNode] = None
) -> Tuple[List[_VarNode], Dict[str, _VarNode]]:
    r"""Replaces ``Host2DeviceCopy`` with :class:`~.InputNode` in the graph
    to :meth:`~.InputNode.set_value` and run.

    Args:
        dst: target vars representing the graph.
        inputs: indicates which inputs to be replaced. All
            inputs(``Host2DeiceCopy``) will be replaced if not specified.

    Returns:
        new vars that correspond to ``dst`` with all inputs replaced, and new inputs dict.
    """
    if inputs is None:
        inputs = get_dep_vars(dst, "Host2DeviceCopy")
    input_dict = OrderedDict()
    replace_dict = {}
    for inp in inputs:
        inp_node = G.InputNode(
            device=inp.comp_node, dtype=inp.dtype, shape=inp.shape, graph=inp.graph,
        )
        inp_node.name = inp.name
        input_dict[inp.name] = inp_node
        replace_dict[inp] = inp_node.outputs[0]
    new_output_nodes = replace_vars(dst, replace_dict)
    for old, new in zip(dst, new_output_nodes):
        new.name = old.name
    return new_output_nodes, input_dict


def convert_outputs(dst: List[_VarNode]) -> Tuple[List[_VarNode], Dict[str, _VarNode]]:
    r"""Wraps ``dst`` with :class:`~.OutputNode` in the graph to get outputs
    with :meth:`~.OutputNode.get_value`.

    Args:
        dst: target vars representing the graph.

    Returns:
        new vars that correspond to ``dst`` with all inputs replaced, and outputs dict.
    """
    output_dict = OrderedDict([(i.name, G.OutputNode(i)) for i in dst])
    new_output_nodes = [i.outputs[0] for i in output_dict.values()]
    return new_output_nodes, output_dict


def embed_inputs(
    dst: List[_VarNode], data: List[np.ndarray], inputs: List[_VarNode] = None
) -> Tuple[List[_VarNode], Dict[str, _VarNode]]:
    r"""Embeds ``data`` to the graph's inputs of ``dst``.

    Args:
        dst: target vars representing the graph.
        data: data to be embeded.
        inputs: indicates which inputs to be replaced. All
            inputs(``Host2DeiceCopy``) will be replaced if not specified.

    Returns:
      new vars that correspond to ``dst`` with all inputs replaced, and new inputs dict.
    """
    if inputs is None:
        inputs = get_dep_vars(dst, "Host2DeviceCopy")
    assert len(data) == len(inputs)
    input_dict = OrderedDict()
    replace_dict = {}
    for inp, d in zip(inputs, data):
        new_inp = _imperative_rt.make_shared(inp.graph, Tensor(d)._dev_tensor())
        new_inp.name = inp.name
        input_dict[inp.name] = new_inp
        replace_dict[inp] = new_inp
    new_output_nodes = replace_vars(dst, replace_dict)
    for old, new in zip(dst, new_output_nodes):
        new.name = old.name
    return new_output_nodes, input_dict


class GraphInference:
    r"""Loads a serialized computing graph as a GraphInference object which can be used
    to execute the computing graph.

    Args:
        file: could be file object or filename.
        outputs: only compile the subgraph with outputs as its endpoints.
    """

    def __init__(
        self,
        file,
        outputs: List[str] = None,
        profiling: bool = False,
        optimize_for_inference: bool = False,
        **kwargs
    ):
        ret = G.load_graph(file)
        self._graph, output_nodes = ret.graph, ret.output_vars_list
        if outputs is not None:
            output_nodes = find_vars_by_name(output_nodes, outputs)
        self._origin_outputs = output_nodes

        # replace inputs with `InputNode`
        output_nodes, self._inp_dict = convert_inputs(output_nodes)

        # replace outputs with `OutputNode`
        output_nodes, self._oup_dict = convert_outputs(output_nodes)

        self._func = self._graph.compile(output_nodes)

    def run(
        self, *inp_args: np.ndarray, inp_dict: Dict[str, np.ndarray] = None
    ) -> Dict[str, np.ndarray]:
        r"""

        Args:
            inp_args: list of input datas.
            inp_dict: dict of named input datas.

        Returns:
            a dict {output_name: output_value}.

        Note:
            Note that the order of the Graph's input nodes may be different from the order of the origin traced function's arguments.
            It is recommended to use ``inp_dict`` to provide input data by name.
        """
        assert len(inp_args) <= len(
            self._inp_dict
        ), "This model expects {} inputs".format(len(self._inp_dict))
        inputs = {}
        inp_keys = list(self._inp_dict.keys())
        for ind, data in enumerate(inp_args):
            inputs[inp_keys[ind]] = data
        if inp_dict is not None:
            inputs.update(inp_dict)
        assert (
            inputs.keys() == self._inp_dict.keys()
        ), "This model expects inputs {}, but gets inputs {}".format(
            list(self._inp_dict.keys()), list(inputs.keys())
        )
        for key in self._inp_dict:
            self._inp_dict[key].set_value(
                Tensor(inputs[key], device=self._inp_dict[key].device)._dev_tensor()
            )
        self._func.execute()
        self._func.wait()

        result = OrderedDict()
        for key in self._oup_dict:
            result[key] = self._oup_dict[key].get_value().numpy()
        return result