# Copyright (c) The University of Edinburgh 2014
#
# 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.

'''
The dispel4py workflow graph.
'''

import networkx as nx
import inspect
import sys

from dispel4py.core import GenericPE


class WorkflowNode:
    '''
    Wrapper class for workflow nodes - wraps around general subclasses
    of classes denoting PEs, that is GenericPEs.
    '''

    # Supported types of workflow nodes:
    WORKFLOW_NODE_PE = 0
    WORKFLOW_NODE_FN = 1
    WORKFLOW_NODE_CP = 2
    node_counter = 0

    def __init__(self, o):
        self.obj = o
        self.outputs = []
        self.inputs = []

        if inspect.isroutine(o):  # it's a 'function'
            self.nodeType = self.WORKFLOW_NODE_FN
        elif isinstance(o, GenericPE):
            o.id = o.name + str(WorkflowNode.node_counter)
            WorkflowNode.node_counter += 1
            self.nodeType = self.WORKFLOW_NODE_PE
            for i in o.inputconnections.values():
                # empty for the time being - only the index matters
                self.inputs.append({})
            for i in o.outputconnections.values():
                self.outputs.append({})
        elif isinstance(o, WorkflowGraph):
            self.nodeType = self.WORKFLOW_NODE_CP
            try:
                for i in o.inputmappings:
                    self.inputs.append({})
            except AttributeError:
                pass
            try:
                for i in o.outputmappings:
                    self.outputs.append({})
            except AttributeError:
                pass
        else:
            sys.stderr.write('Error: Unknown type of object passed as a \
                              Workflow Node: %s\n' % type(o))
            raise Exception("Unknown type of object passed as a \
                             Workflow Node: %s" % type(o))

    def getContainedObject(self):
        ''' Returns the wrapped PE or function. '''
        return self.obj


# Used as attribute names
FROM_CONNECTION = 'from_connection'
TO_CONNECTION = 'to_connection'
DIRECTION = 'direction'


class WorkflowGraph(object):
    """
    A graph representing the workflow and related methods
    """

    def __init__(self):
        self.graph = nx.Graph()
        self.objToNode = {}

    def add(self, n):
        '''
        Adds node n, which must be an instance of
        :py:class:`dispel4py.core.GenericPE`, and returns the created workflow
        node.

        :rtype: WorkflowNode
        '''
        nd = WorkflowNode(n)
        self.graph.add_node(nd)
        self.objToNode[n] = nd
        return nd

    def connect(self, fromNode, fromConnection, toNode, toConnection):
        '''
        Connect the two given nodes from the given output to the given input.
        If the nodes are not in the graph, they will be added.

        :param fromNode: the source PE of the connection
        :param fromConnection: the name of the output of the source node
        'fromNode'
        :type fromConnection: String
        :param toNode: the destination PE of the connection
        :param toConnection: the name of the input of the destination node
        'toNode'
        :type toConnection: String
        '''

        if fromNode not in self.objToNode:
            self.add(fromNode)
        if toNode not in self.objToNode:
            self.add(toNode)

        fromWfNode = self.objToNode[fromNode]
        toWfNode = self.objToNode[toNode]

        if self.graph.has_edge(fromWfNode, toWfNode):
            self.graph[fromWfNode][toWfNode]['ALL_CONNECTIONS']\
                .append((fromConnection, toConnection))
        else:
            self.graph.add_edge(fromWfNode, toWfNode,
                                **{'FROM_CONNECTION': fromConnection,
                                   'TO_CONNECTION': toConnection,
                                   'DIRECTION': (fromNode, toNode),
                                   'ALL_CONNECTIONS': [
                                       (fromConnection, toConnection)]})

    def getContainedObjects(self):
        nodes = [node.getContainedObject() for node in self.graph.nodes()]
        return sorted(nodes, key=lambda x: x.id)

    def propagate_types(self):
        '''
        Propagates the types throughout the graph by retrieving the output
        types from each node, starting from the root, and providing them to
        connected consumers.
        '''
        visited = set()
        for node in self.graph.nodes():
            if node not in visited:
                self.__assign_types(node, visited)

    def __assign_types(self, node, visited):
        pe = node.getContainedObject()
        inputTypes = {}
        for edge in self.graph[node].values():
            if pe == edge['DIRECTION'][1]:
                # pe is destination so look up the types produced by sources
                source = edge['DIRECTION'][0]
                sourceNode = self.objToNode[source]
                if sourceNode not in visited:
                    self.__assign_types(sourceNode, visited)
                inType = source.getOutputTypes()[edge['FROM_CONNECTION']]
                inputTypes[edge['TO_CONNECTION']] = inType
        pe.setInputTypes(inputTypes)
        visited.add(node)
        # print "%s: Assigned inputs = %s, \
        # received outputs = %s" % \
        # (pe.__class__.__name__, inputTypes, pe.getOutputTypes())

    def flatten(self):
        '''
        Subgraphs contained within composite PEs are added to the top level
        workflow.
        '''
        hasComposites = True
        toRemove = set()
        while hasComposites:
            hasComposites = False
            toRemove = set()
            for node in self.graph.nodes():
                if node.nodeType == WorkflowNode.WORKFLOW_NODE_CP:
                    hasComposites = True
                    toRemove.add(node)
                    wfGraph = node.getContainedObject()
                    subgraph = wfGraph.graph
                    self.graph.add_nodes_from(subgraph.nodes(data=True))
                    self.graph.add_edges_from(subgraph.edges(data=True))
                    self.objToNode.update(wfGraph.objToNode)
                    for inputname in wfGraph.inputmappings:
                        toPE, toConnection = wfGraph.inputmappings[inputname]
                        edge = None
                        fromPE, fromConnection = None, None
                        for e in self.graph[node].values():
                            if wfGraph == e['DIRECTION'][1] \
                                    and inputname == e['TO_CONNECTION']:
                                fromPE = e['DIRECTION'][0]
                                fromConnection = e['FROM_CONNECTION']
                                edge = self.objToNode[fromPE],\
                                    self.objToNode[wfGraph]
                                break
                        if edge is not None:
                            self.connect(fromPE, fromConnection,
                                         toPE, toConnection)
                    for outputname in wfGraph.outputmappings:
                        fromPE, fromConnection = \
                            wfGraph.outputmappings[outputname]
                        destinations = []
                        for e in self.graph[node].values():
                            if wfGraph == e['DIRECTION'][0] \
                                    and outputname == e['FROM_CONNECTION']:
                                toPE = e['DIRECTION'][1]
                                toConnection = e['TO_CONNECTION']
                                destinations.append((toPE, toConnection))
                        for (toPE, toConnection) in destinations:
                            # print 'connecting output %s.%s' %\
                            # (toPE, toConnection)
                            self.connect(fromPE, fromConnection,
                                         toPE, toConnection)
            self.graph.remove_nodes_from(toRemove)


def _create_dot(graph, instanceNames={}, counter=0):
    dot = ''
    # assign unique names
    for node in graph.graph.nodes():
        try:
            name = node.getContainedObject().id, counter
        except:
            name = node.getContainedObject().__class__.__name__, counter
        instanceNames[node] = name
        counter += 1

    # now add all the nodes and their input and output connections
    cluster_index = 0
    for node in graph.graph.nodes():
        pe = node.getContainedObject()
        if isinstance(pe, WorkflowGraph):
            dot += _create_cluster(pe, cluster_index, instanceNames, counter)
            cluster_index += 1
            continue

        name, index = instanceNames[node]
        dot += name + str(index) + "[label=\"{ "
        # add inputs
        inputNames = []
        outputNames = []
        for edge in graph.graph[node].values():
            if pe == edge['DIRECTION'][1]:
                inputName = edge['TO_CONNECTION']
                dotName = "<in_" + inputName + ">" + inputName
                if dotName not in inputNames:
                    inputNames.append(dotName)
            else:
                outputName = edge['FROM_CONNECTION']
                dotName = "<out_" + outputName + ">" + outputName
                if dotName not in outputNames:
                    outputNames.append(dotName)

        if inputNames:
            dot += '{' + ' | '.join(inputNames) + '} | '
        dot += name
        if outputNames:
            dot += ' | {' + ' | '.join(outputNames) + '}'
        dot += " }\"];\n"

    # connect the inputs and outputs
    for node in graph.graph.nodes():
        pe = node.getContainedObject()
        for edge in graph.graph[node].values():
            if pe == edge['DIRECTION'][0]:
                if isinstance(pe, WorkflowGraph):
                    inner_source, source_output = \
                        pe.outputmappings[edge['FROM_CONNECTION']]
                    node = pe.objToNode[inner_source]
                else:
                    source_output = edge['FROM_CONNECTION']
                # pe is the source so look up the connected destination
                dest = edge['DIRECTION'][1]
                if isinstance(dest, WorkflowGraph):
                    inner_dest, dest_input = \
                        dest.inputmappings[edge['TO_CONNECTION']]
                    destNode = dest.objToNode[inner_dest]
                else:
                    destNode = graph.objToNode[dest]
                    dest_input = edge['TO_CONNECTION']
                dot += '%s%s' % instanceNames[node] + ':out_' + source_output
                dot += ' -> '
                dot += '%s%s' % instanceNames[destNode] + ':in_' + dest_input
                dot += ';\n'
    return dot


def _create_cluster(graph, index, instanceNames, counter):
    dot = 'subgraph cluster_%s {\n' % index
    try:
        # names for composite PEs are optional
        dot += 'label = "%s";' % graph.name
    except:
        pass
    dot += 'style=filled;\n'
    dot += 'color=lightgrey;\n'
    if index % 2:
        dot += 'fillcolor=lightgrey;\n'
    dot += _create_dot(graph, instanceNames, counter)
    dot += '}\n'
    return dot


def draw(graph):
    '''
    Creates a representation of the workflow graph in the dot language.
    '''
    dot = 'digraph request\n{\nnode [shape=Mrecord, \
           style=filled, fillcolor=white];\n'
    dot += _create_dot(graph)
    dot += '}\n'
    return dot


def drawDot(graph):
    '''
    Draws the workflow as a graph and creates a PNG image using graphviz dot.
    '''
    from subprocess import Popen, PIPE
    dot = draw(graph)
    p = Popen(['dot', '-T', 'png'], stdout=PIPE, stdin=PIPE, stderr=PIPE)
    stdout, stderr = p.communicate(dot.encode('utf-8'))
    return stdout


def getConnectedInputs(node, graph):
    names = []
    for edge in graph.edges(node, data=True):
        direction = edge[2]['DIRECTION']
        dest = direction[1]
        dest_input = edge[2]['TO_CONNECTION']
        if dest == node.getContainedObject():
            names.append(dest_input)
    return names
