from collections import OrderedDict

from chainer.utils import type_check

from chainer_wing import compiler
from chainer_wing import runner
from chainer_wing.node import Node, MetaNode
from chainer_wing.node import NODECLASSES
from chainer_wing import util
from chainer_wing.subwindows.train_config import TrainParamServer


def dummy(node_class):
    return node_class


class Graph(object):
    """
    Managing all nodes. This class provides interfaces for spawning/removing
    nodes and connections.
    """

    def __init__(self, painter=None):
        self.nodes = {}
        self.runner = None
        # from self to other
        self.connections = {}
        # from other to self
        self.reverseConnections = {}

        if painter:
            self.painter = painter
            painter.registerGraph(self)
        else:
            self.painter = dummy

    def __getattr__(self, item):
        return super(Graph, self).__getattr__(item)

    def spawnNode(self, node_class, connections=None, position=(0, 0),
                  silent=False, id=None, name=None):
        """
        Spawns a new node of a given class at a given position with optional
        connections to other nodes.
        :param node_class: subclass object of 'Node'.
        :param connections: Dictionary
        :param position: Tuple of two integer representing the nodes position
        on the screen relative the the graph's
        origin.
        :param silent: Boolean. Suppresses all notifications that a node was
        spawned if True.
        :return: newly created Node instance.
        """
        newNode = node_class(self, position, id)
        newNode.name = name
        self.reverseConnections[newNode] = set()
        self.connections[newNode] = set()
        if connections:
            self._spawnConnections(connections, newNode)
        try:
            self.painter.registerNode(newNode, silent)
        except AttributeError:
            pass
        self.nodes[newNode.ID] = newNode

        return newNode

    def pasteNode(self, node, pos):
        new_node = self.spawnNode(node.__class__, position=(pos.x(), pos.y()))
        node.clone_param(new_node)

    def createCustomNodeClass(self, name, inputs, outputs, parents=(Node,)):
        node_class = MetaNode(name, parents, {})
        node_class.__inputs__ = OrderedDict()
        node_class.__outputs__ = OrderedDict()
        for inp in inputs:
            node_class._addInput(data=inp, cls=node_class)
        for out in outputs:
            node_class._addOutput(data=out, cls=node_class)
        NODECLASSES[name] = node_class
        return node_class

    def _spawnConnections(self, connections, newNode):
        try:
            outs = connections['outputs']
        except KeyError:
            pass
        else:
            for out in outs:
                self.connect(newNode, out[0], out[1], out[2])
        try:
            ins = connections['inputs']
        except KeyError:
            pass
        else:
            for inp in ins:
                self.connect(inp[1], inp[2], newNode, inp[0])

    def connect(self, outNode, out, inpNode, inp):
        """
        Creates a logical connection between two nodes.
        Before the actual connection is established checks will be performed
         to make sure that the input is actually legal.
        If necessary, previously created connections that are in conflict with
         the new one will be deleted.
        :param outNode: Node instance that has the output involved in the connection.
        :param out: string representing the Output's name.
        :param inpNode: Node instance that has the input involved in the connection.
        :param inp: string representing the Input's name.
        :return:
        """
        if isinstance(outNode, str):
            outNode = self.nodes[outNode]
        if isinstance(inpNode, str):
            inpNode = self.nodes[inpNode]
        outInfo = outNode.getOutputInfo(out)
        inpInfo = inpNode.getInputInfo(inp)
        if not set(outInfo.var_type) & set(inpInfo.var_type):
            raise TypeError(
                'Output \'{}\' of node {} and input \'{}\' of not {} don\'t match.'.format(
                    out,
                    str(outNode),
                    inp,
                    str(inpNode)))
        conn = Connection(outNode, out, inpNode, inp)
        for oldCon in self.reverseConnections[inpNode]:
            if oldCon.input_name == inp:
                self.reverseConnections[inpNode].remove(oldCon)
                self.connections[oldCon.output_node].remove(oldCon)
                break
        inpInfo.setConnected(True)
        self.connections[outNode].add(conn)
        self.reverseConnections[inpNode].add(conn)
        # self.update()

    def getConnectionsFrom(self, node):
        """
        Returns a list of all connections that involve 'node's' outputs.
        :param node:
        :return: List of connection dictionaries.
        :rtype: list
        """
        return self.connections[node]

    def getConnectionsTo(self, node):
        """
        Returns a list of all connections that involve 'node's' inputs.
        :param node:
        :return:
        """
        return self.reverseConnections[node]

    def getConnectionOfInput(self, inp):
        """
        Returns the connection involving an input
        :param inp: InputInfo instance.
        :return: Connection instance. If there is no connection, return None.
        """
        for con in self.getConnectionsTo(
                self.nodes[inp.ID.partition(':')[0]]):
            if con.input_name == inp.name:
                return con
        return None

    def getConnectionsOfOutput(self, output):
        """
        Returns a list of connections involving an output.
        :param output: OutputInfo instance.
        :return: list of Connection instances.
        """
        node = self.nodes[output.ID.partition(':')[0]]
        return [con for con in self.getConnectionsFrom(node) if
                con.output_name == output.name]

    def update(self):
        """
        Updates and repaints the painter instance.
        :return:
        """
        self.painter.repaint()
        self.painter.update()

    def compile(self):
        """
        Compile the Graph as chainer code.
        :return: If compilation was succeeded, return True.
        """
        try:
            result = compiler.Compiler()(self.nodes)
        except util.ExistsInvalidParameter as error:
            util.disp_error('{0} is not set @{1}'.format(error.args[1][1:],
                                                         error.args[0]))
            self.nodes[error.args[0]].runtime_error_happened = True
            return False
        except compiler.NoLossError:
            util.disp_error('Please place loss function.')
            return False
        return result

    def run(self):