"""
 Copyright (c) 2018 Intel Corporation

 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 collections
import logging as log

import networkx as nx
import numpy as np

from mo.utils.error import Error
from mo.utils.utils import refer_to_faq_msg


def unique_id(graph: nx.MultiDiGraph, prefix: str = ""):
    """
    Generates a unique node id for a new node in a given graph.
    The optional string prefix can be specified.
    """
    # TODO thread safety?
    unique_id.count = max(unique_id.count, graph.number_of_nodes()) + 1
    if prefix and not graph.has_node(prefix):
        return str(prefix)
    while graph.has_node(prefix + str(unique_id.count)):
        unique_id.count += 1
    return prefix + str(unique_id.count)


unique_id.count = 0


def get_node_id_by_name(graph: nx.MultiDiGraph, name: str):
    for node in graph.nodes():
        if 'name' in graph.node[node] and graph.node[node]['name'] == name:
            return node
    raise Error('No node with name {}. ' +
                refer_to_faq_msg(51), name)


def create_graph_with_nodes(src_nodes, get_id: callable, get_attrs: callable):
    """
    Go over all nodes in src_nodes that should be enumerable and create new NX nodes
    using get_id and get_attrs functions to create node id and node attributes correspondingly.
    """
    graph = nx.MultiDiGraph()
    for node in src_nodes:
        graph.add_node(get_id(node), **get_attrs(node))
    return graph


# TODO implement merging for keys with dictionary values?
def merge_edge_props(attrs: dict, additional_attrs: dict):
    """
    Update edge attributes without changing 'in' and 'out' keys.
    It is necessary to copy edge attributes during merging of nodes when
    result of one subgraph call is passed as input to another subgraph call
    """
    result = attrs
    for (key, value) in additional_attrs.items():
        if key not in ['in', 'out']:
            if type(additional_attrs[key]) is list:
                if key not in result:
                    result[key] = []
                result[key].extend(additional_attrs[key])
                result[key] = list(set(result[key]))  # silly solution to find unique elements
            else:
                result[key] = value
    return result


def print_graph_stat(graph: nx.MultiDiGraph):
    log.debug('Number of nodes in graph: {}'.format(graph.number_of_nodes()))
    log.debug('Number of edges in graph: {}'.format(len(list(graph.edges()))))
    ops = collections.defaultdict(int)
    for _node in graph.nodes():
        node = NodeWrap(graph, _node)
        kind = node.kind if node.has('kind') else '<UNDEFINED>'
        if node.has('op'):
            ops['op/' + node.op] += 1
        else:
            ops[kind] += 1
        if node.has('shape') and np.any(node.shape == 0):
            log.error("Found bad shape: '{}' for node '{}'".format(node.shape, node.node))
    for k, v in ops.items():
        log.debug('   {} : {}'.format(k, v))


def dump_graph_for_graphviz(graph: nx.MultiDiGraph, node_attrs: list = ['kind', 'op', 'shape'],
                            edge_attrs: list = ['in', 'out'],
                            nodes_to_dump: list = None, save_to_svg = False):
    log.debug("---- GRAPHVIZ OUTPUT STARTS ----")
    if nodes_to_dump is None:
        nodes_to_dump = graph.nodes()
    string = '\ndigraph {\n'
    for src_node_name, dst_node_name, attrs in graph.edges(data=True):
        if src_node_name not in nodes_to_dump or dst_node_name not in nodes_to_dump:
            continue
        src_node = graph.node[src_node_name]
        dst_node = graph.node[dst_node_name]
        src_node_string = str(src_node_name) + '\\n' + '\\n'.join(
            [str(key) + '=' + str(src_node.get(key, 'None')) for key in node_attrs if key in src_node])
        dst_node_string = str(dst_node_name) + '\\n' + '\\n'.join(
            [str(key) + '=' + str(dst_node.get(key, 'None')) for key in node_attrs if key in dst_node])
        edge_string = ' '.join([str(key) + '=' + str(attrs.get(key, 'None')) for key in edge_attrs if key in attrs])
        string += '"{}" -> "{}" [label = "{}"];\n'.format(src_node_string, dst_node_string, edge_string)
    string += '}'
    log.debug(string)
    log.debug("---- GRAPHVIZ OUTPUT ENDS ----")

    if save_to_svg:
        try:
            import graphviz
            import os
            file_name = "{}_{}.txt".format(graph.name, 0)
            id = 1
            while os.path.exists(file_name):
                file_name = "{}_{}.txt".format(graph.name, id)
                id += 1
            with open(file_name, "w") as f:
                f.write(string)
            graphviz.render('dot','svg', file_name)
            print('Graph was saved to {}.{}'.format(file_name, 'svg'))
        except ImportError:
            raise ImportError('Can\'t import graphviz')
        except Exception as e:
            raise Error('Can\'t save graph to svg') from e

    return string


def create_sub_graph_copy(graph: nx.MultiDiGraph, nodes_to_extract: list):
    """
    Create new graph which is a sub-graph of the 'graph' that contains just nodes from 'nodes_to_extract' list. The
    returned sub-graph is a deep copy of the provided graph nodes.
    :param graph: graph to create a sub-graph from.
    :param nodes_to_extract: list of node names to extract.
    :return: new graph.
    """
    return graph.subgraph(nodes_to_extract).copy()


def get_inputs(graph: nx.MultiDiGraph, node: str, edge_attr: dict = {}):
    in_edges = graph.in_edges(node, data=True)
    return [(u, d) for u, v, d in in_edges if all([attr in d and d[attr] == edge_attr[attr] for attr in edge_attr])]


def get_outputs(graph: nx.MultiDiGraph, node: str, edge_attr: dict = {}):
    out_edges = graph.out_edges(node, data=True)
    return [(v, d) for u, v, d in out_edges if all([attr in d and d[attr] == edge_attr[attr] for attr in edge_attr])]


def get_single_input(graph: nx.MultiDiGraph, node: str, edge_attr: dict = {}):
    """
    Searches for all edges that have given attributes.
    If there no such edges or there are multiple edges, raise exception.
    If there is only one edge, returns the source node for this edge
    and the edge attributes themselves.
    """
    inputs = get_inputs(graph, node, edge_attr)
    if len(inputs) != 1:
        log.debug("Node '{}' has {} inputs with edge attributes '{}'".format(node, inputs, str(edge_attr)))
        raise AttributeError(
            "None or multiple inputs satisfy given attributes. Node: " + str(node) + ", edge_attr: " + str(edge_attr))
    return inputs[0]


def get_single_output(graph: nx.MultiDiGraph, node: str, edge_attr: dict = {}):
    outputs = get_outputs(graph, node, edge_attr)
    if len(outputs) != 1:
        log.debug("Node '{}' has {} outputs with edge attributes '{}'".format(node, outputs, str(edge_attr)))
        raise AttributeError(
            "None or multiple outputs satisfy given attributes. Node: " + str(node) + ", edge_attr: " + str(edge_attr))
    return outputs[0]


def get_graph_ops(graph: nx.MultiDiGraph):
    return [Node(graph, node) for node in graph.nodes() if Node(graph, node).soft_get('kind') == 'op']


class NodeWrap:

    def __init__(self, graph: nx.MultiDiGraph, node: str):
        super(NodeWrap, self).__setattr__('graph', graph)
        super(NodeWrap, self).__setattr__('node', node)  # obsolete
        super(NodeWrap, self).__setattr__('id', node)

    def __setattr__(self, k, v):
        # you can assign only existing attributes
        attrs = self.graph.node[self.node]
        if not k in attrs:
            raise AttributeError
        attrs[k] = v

    def __getattr__(self, k):
        # hope it raises AttributeError if k is not in the dict
        return self.graph.node[self.node][k]

    def has(self, k):
        return k in self.graph.node[self.node]

    def has_valid(self, k):
        return self.has(k) and not self.graph.node[self.node][k] is None

    def has_and_set(self, k):
        return self.has_valid(k) and self[k]

    def __getitem__(self, k):
        return self.graph.node[self.node][k]

    def __setitem__(self, k, v):
        self.graph.node[self.node][k] = v

    def __contains__(self, k):
        return self.has(k)

    def in_nodes_edges(self):
        return {x[1]['in']: (NodeWrap(self.graph, x[0]), x[1]) for x in get_inputs(self.graph, self.node)}

    def in_nodes(self):
        assert self.has('kind')
        assert self.kind in ['op', 'data']
        if self.kind == 'op':
            return {x[1]['in']: NodeWrap(self.graph, x[0]) for x in get_inputs(self.graph, self.node)}
        elif self.kind == 'data':
            return [NodeWrap(self.graph, n) for n, d in get_inputs(self.graph, self.node)]

    def in_edges(self):
        assert self.has('kind')
        assert self.kind in ['op', 'data']
        if self.kind == 'op':
            return {x[1]['in']: x[1] for x in get_inputs(self.graph, self.node)}
        elif self.kind == 'data':
            return [d for n, d in get_inputs(self.graph, self.node)]

    def out_nodes_edges(self):
        return {x[1]['out']: (NodeWrap(self.graph, x[0]), x[1]) for x in get_outputs(self.graph, self.node)}

    def out_nodes(self):
        assert self.has('kind')
        assert self.kind in ['op', 'data']
        if self.kind == 'op':
            return {x[1]['out']: NodeWrap(self.graph, x[0]) for x in get_outputs(self.graph, self.node)}
        elif self.kind == 'data':
            return [NodeWrap(self.graph, n) for n, d in get_outputs(self.graph, self.node)]

    def out_edges(self):
        assert self.has('kind')
        assert self.kind in ['op', 'data']
        if self.kind == 'op':
            return {x[1]['out']: x[1] for x in get_outputs(self.graph, self.node)}
        elif self.kind == 'data':
            return [d for n, d in get_outputs(self.graph, self.node)]

    def in_node(self, key=0):
        return self.in_nodes()[key]

    def out_node(self, key=0):
        return self.out_nodes()[key]

    def in_edge(self, key=0):
        return self.in_edges()[key]

    def out_edge(self, key=0):
        return self.out_edges()[key]

    def get_attrs(self):
        return self.graph.node[self.node]

    def soft_get(self, k):
        return self[k] if self.has_valid(k) else '<UNKNOWN>'


class Node(NodeWrap):
    pass


def get_sorted_inputs(node: Node):
    return sorted([x for x in get_inputs(node.graph, node.node) if 'in' in x[1]], key=lambda x: x[1]['in'])


def get_sorted_outputs(node: Node):
    return sorted([x for x in get_outputs(node.graph, node.node) if 'out' in x[1]], key=lambda x: x[1]['out'])


def create_edge(src_node: Node, dst_node: Node, out_port: int = 0, in_port: int = 0, edge_attrs: dict = None):
    """
    Creates edge from node 'src_node' from output with index 'out_port' to node 'dst_node' with input index 'in_port'.
    :param src_node: node to create edge from.
    :param dst_node: node to create edge to.
    :param out_port: the index of output tensor of the 'src_node'.
    :param in_port: the input index of the node 'dst_node'.
    :param edge_attrs: dictionary with edge attrs.
    :return: None
    """
    # edges must belong to the same graph
    assert src_node.graph is dst_node.graph
    graph = src_node.graph

    if edge_attrs is None:
        edge_attrs = dict()
    else:
        edge_attrs = edge_attrs.copy()
    edge_attrs.update({'in': in_port, 'out': out_port, 'in_attrs': ['in'], 'out_attrs': ['out'],
                       'data_attrs': ['fw_tensor_debug_info']})

    graph.add_edges_from([(src_node.id, dst_node.id, edge_attrs)])


def insert_node_after(node: Node, new_node: Node, node_out_port: int = 0):
    """
    Insert node 'new_node' after output with index 'node_out_port' of the node 'node'. All consumers of node 'node'
    output with index 'node_out_port' will be changed to consume node 'new_node'.
    The function should be used when graph doesn't contain data nodes yet.
    :param node: node after which new node should be inserted.
    :param new_node: node to be inserted.
    :param node_out_port: the output index for the node 'node' to insert
    :return: None
    """
    assert node.graph is new_node.graph
    assert (len([name for name in node.graph.nodes() if Node(node.graph, name).soft_get('kind') == 'data']) == 0)

    graph = node.graph
    for _, dst_name, edge_attrs in list(graph.out_edges(node.id, data=True)):
        if edge_attrs['out'] == node_out_port:
            log.debug('Create edge from "{}" to "{}"'.format(new_node.name, dst_name))
            create_edge(new_node, Node(graph, dst_name), 0, edge_attrs['in'])
            log.debug('Remove edge from "{}" to "{}"'.format(node.name, dst_name))
            graph.remove_edge(node.id, dst_name)
    create_edge(node, new_node, node_out_port, 0, {})


def erase_node(node: Node):
    """
    Erases node from the graph and reconnect edges from input node(s) to output node(s)
    Produces assertion error in case of multiple inputs and outputs node at the same time
    :param node: Node to erase
    """
    graph = node.graph
    node_id = node.id

    inputs = list(graph.in_edges(node_id, data=True))
    outputs = list(graph.out_edges(node_id, data=True))

    assert not (len(inputs) > 1 and len(outputs) > 1)

    if len(outputs) == 0 and len(inputs) != 0:
        for input, _, attrs in inputs:
            if Node(graph, node_id).has_and_set('is_output'):
                if graph.node[input]['kind'] == 'op':
                    data_nodes = [u for u, v in graph.in_edges(input)]
                    for data in data_nodes:
                        graph.node[data]['is_output'] = graph.node[node_id]['is_output']
                else:
                    graph.node[input]['is_output'] = graph.node[node_id]['is_output']

    if len(outputs) == 0 or len(inputs) == 0:
        graph.remove_node(node_id)
        return

    if len(outputs) == 1:
        output = outputs[0][1]
        for src, noop, attrs in inputs:
            graph.remove_edge(src, noop)
            graph.add_edge(src, output, **attrs)
        graph.remove_node(node_id)
        return

    if len(inputs) == 1:
        input = inputs[0][0]
        for noop, dst, attrs in outputs:
            graph.remove_edge(noop, dst)
            graph.add_edge(input, dst, **attrs)
        graph.remove_node(node_id)
        return


def replace_node(old_node: Node, new_node: Node):
    """
    Replaces node 'old_node' with a node 'new_node' preserving edge attributes.
    :param old_node: node to be replaced.
    :param new_node: node to replace with.
    :return: None
    """
    assert old_node.graph is new_node.graph
    graph = old_node.graph
    # save output edges and reconnect them to new node
    for i in range(len(old_node.out_nodes())):
        graph.add_edge(new_node.id, old_node.out_node(i).id, **old_node.out_edge(i))
    # TODO Need to check if there are other users for this node
    graph.remove_node(old_node.id)
