import abc
from collections import OrderedDict

import networkx as nx
import networkx.algorithms.dag as dag

from modelflow.elementarynodes import InputChannel
from .node import Callable, StaticInputs, NameMixin, Node


class Graph(NameMixin):
    """
    Graph base class.
    """
    def __init__(self):
        self._graph = nx.MultiDiGraph()
        self._output = None

    def __setitem__(self, key, value):
        try:
            self._inputs[key]['value'] = value
        except KeyError as keyerryr:
            raise KeyError(key + " in {}".format(self))

    def __getitem__(self, key):
        return self._inputs[key]()

    @property
    def nodes(self):
        """Return the computational nodes in the graph"""
        return self._graph.nodes()

    def get_nodes_by_name(self, name):
        """Get a node by its name"""
        return (node for node in self.nodes if node.name == name)

    def _null_input(self, name):
        return InputChannel(name=name)

    def connect_input_to_node(self, from_input: str, to_node: Node, to_channel: str):
        """
        Connect an input to a node

        Parameters
        ----------
        from_input: str
            Input of the graph which to connect.

        to_node: Node
            Connect the input to this node here.

        to_channel: str
            Connect the input to the channel of the node specified here.
        """
        self._graph.add_edge(self._inputs[from_input], to_node, to_channel=to_channel)

    def set_output(self, node: Node):
        """
        Set a node as output

        Parameters
        ----------
        node: Node
            The node which to set as output.
        """
        self._output = node

    def connect_node_to_node(self, from_node: Node, to_node: Node, to_channel: str):
        """
        Connect a node within the graph to another one.

        Parameters
        ----------
        from_node: Node
            Source node.

        to_node: Node
            Target node.

        to_channel: str
            Target channel
        """
        self._graph.add_edge(from_node, to_node, to_channel=to_channel)

    def add_path(self, *path_elements, last_is_output=True):
        """
        Add a computational path to the graph.

        This is a convenience method.

        Parameters
        ----------
        path_elements
            an arbitrary number of arguments, where each is of the form
            ``({"GraphInput1": "NodeInput1", "GraphInput2": ("NodeInput2", "NodeInput2.2")}, node)``

            "GraphInput" is here an input of the Graph which is mapped to "NodeInput" of the  node
            "GraphInput" can also be replaced by the special "%" sign, which represents
            the output of the predecessor node

        last_is_output: Boolean
            whether to set the last element automatically as graph output
        """
        last_node = None
        for input_mapping, node in path_elements:
            self._add_node_to_path(input_mapping, last_node, node)
            last_node = node
        if last_is_output:
            self.set_output(last_node)

    def _add_node_to_path(self, input_mapping, last_node, node):
        for source, targets in input_mapping.items():
            if isinstance(targets, str):
                targets = targets,
            for target in targets:
                if source != "%":
                    self.connect_input_to_node(source, node, target)
                if source == "%":
                    if last_node is None:
                        raise Exception("No predecessor not to take input from for node " + str(node))
                    self.connect_node_to_node(last_node, node, target)

    def run(self):
        sorted_nodes = dag.topological_sort(self._graph)
        for node in sorted_nodes:
            result = node()
            if node == self._output:
                return result
            self._propagate_result(node, result)

    def _propagate_result(self, node, result):
        for target, edge_collection in self._graph[node].items():
            for edge in edge_collection.values():
                target[edge['to_channel']] = result

    @property
    def input_names(self):
        """Return list of the input names."""
        return list(self._inputs.keys())


class DynamicGraph(Graph, Callable):
    """
    A computational graph which can be assembled at run time.
    """
    def __init__(self):
        super().__init__()
        self._inputs = OrderedDict()
        self.inputs_set_at_initialization = []

    def add_inputs(self, input_id_list):
        for input_id in input_id_list:
            if input_id in self._inputs:
                raise Exception("Input exists already: " + str(input_id))
            self._inputs.update({input_id: InputChannel(input_id)})

    def update_inputs(self, input_dict):
        for key, value in input_dict.items():
            self[key] = value


class StaticGraph(Graph, Callable, StaticInputs):
    """
    A computational graph which should be assembled at import time.
    """
    def __init__(self, **kwargs):
        Graph.__init__(self)
        StaticInputs.__init__(self, **kwargs)
        self.construct_graph()

    @abc.abstractmethod
    def construct_graph(self):
        "Construct the computational graph here."
