class TestProducer(GenericPE):
    '''
    This PE produces a range of numbers
    '''

    def __init__(self, numOutputs=1):
        GenericPE.__init__(self)
        if numOutputs == 1:
            self._add_output('output', tuple_type=['number'])
        else:
            for i in range(numOutputs):
                self._add_output('output%s' % i, tuple_type=['number'])
        self.counter = 0
        self.outputnames = list(self.outputconnections.keys())

    def _process(self, inputs):
        self.counter += 1
        result = {}
        for output in self.outputnames:
            result[output] = self.counter
        # self.log("Writing out %s" % result)
        return result class TestOneInOneOut(GenericPE):
    '''
    This PE outputs the input data.
    '''

    def __init__(self):
        GenericPE.__init__(self)
        self._add_input('input')
        self._add_output('output')

    def setInputTypes(self, inputtypes):
        self.inputconnections['input']['type'] = inputtypes['input']
        self.outputconnections['output']['type'] = inputtypes['input']

    def process(self, inputs):
        # self.log('Processing inputs %s' % inputs)
        return {'output': inputs['input']} class TestOneInOneOutWriter(GenericPE):
    '''
    This PE copies the input to an output, but it uses the write method.
    Remember that the write function allows to produce more than one output
    block within one processing step.
    '''

    def __init__(self):
        GenericPE.__init__(self)
        self._add_input('input')
        self._add_output('output', tuple_type=['number'])

    def process(self, inputs):
        self.write('output', inputs['input']) class TestTwoInOneOut(GenericPE):
    '''
    This PE takes two inputs and it merges the data into one output string.
    '''

    def __init__(self):
        GenericPE.__init__(self)
        self._add_input('input0')
        self._add_input('input1')
        self._add_output('output', tuple_type=['result'])

    def process(self, inputs):
        # print '%s: inputs %s' % (self.id, inputs)
        result = ''
        for inp in self.inputconnections:
            if inp in inputs:
                result += '%s' % (inputs[inp])
        if result:
            # print '%s: result %s' % (self.id, result)
            return {'output': result} class TestIterative(IterativePE):
    '''
    This PE outputs the input data.
    '''

    def __init__(self):
        IterativePE.__init__(self)

    def _process(self, data):
        return data class IntegerProducer(ProducerPE):

    def __init__(self, start, limit):
        ProducerPE.__init__(self)
        self.start = start
        self.limit = limit

    def _process(self, inputs):
        for i in range(self.start, self.limit):
            self.write('output', i) class PrintDataConsumer(ConsumerPE):

    def __init__(self):
        ConsumerPE.__init__(self)

    def _process(self, data):
        print(data) class RandomWordProducer(GenericPE):
    '''
    This PE produces a random word as an output.
    '''
    words = ["dispel4py", "computing", "mpi", "processing",
             "simple", "analysis", "data"]

    def __init__(self):
        GenericPE.__init__(self)
        self._add_output('output', tuple_type=['word'])

    def process(self, inputs=None):
        word = random.choice(RandomWordProducer.words)
        outputs = {'output': [word]}
        return outputs class RandomFilter(GenericPE):
    '''
    This PE randomly filters the input.
    '''
    input_name = 'input'
    output_name = 'output'

    def __init__(self):
        GenericPE.__init__(self)
        self._add_input('input')
        self._add_output('output', tuple_type=['word'])

    def process(self, inputs):
        if random.choice([True, False]):
            return {'output': inputs['input']}
            # self.write('output', inputs['input'] )
        return None class WordCounter(GenericPE):
    '''
    This PE counts the number of times (counter) that it receives each word.
    And it produces as an output: the same word (the input) and its counter.
    '''
    input_name = 'input'
    output_name = 'output'

    def __init__(self):
        GenericPE.__init__(self)
        self._add_input('input', grouping=[0])
        self._add_output('output', tuple_type=['word', 'count'])
        self.mywords = defaultdict(int)

    def _process(self, inputs):
        word = inputs['input'][0]
        self.mywords[word] += 1
        return {'output': [word, self.mywords[word]]} def simpleLogger(self, msg):
def process_and_return(workflow, inputs, resultmappings=None):
def process(workflow, inputs, args=None, resultmappings=None):
    def __init__(self, pe, provided_inputs=None):
    def _read(self):
    def _write(self, name, data):
def main():                                                  # pragma: no cover
class SimpleProcessingWrapper(GenericWrapper): 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 list(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)
from dispel4py.examples.graph_testing.testing_PEs\
    import TestProducer, TestOneInOneOut, TestOneInOneOutWriter, \
    TestTwoInOneOut, TestIterative, IntegerProducer, PrintDataConsumer, \
    RandomWordProducer, RandomFilter, WordCounter
from dispel4py.new import simple_process
from dispel4py.workflow_graph import WorkflowGraph
from dispel4py.base import create_iterative_chain, CompositePE
from nose import tools
# 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.

'''
Tests for simple sequential processing engine.

Using nose (https://nose.readthedocs.org/en/latest/) run as follows::

    $ nosetests dispel4py/test/simple_process_test.py
    ....
    ----------------------------------------------------------------------
    Ran 4 tests in 0.003s

    OK
'''





def testPipeline():
    prod = TestProducer()
    cons1 = TestOneInOneOut()
    cons2 = TestOneInOneOut()
    graph = WorkflowGraph()
    graph.connect(prod, 'output', cons1, 'input')
    graph.connect(cons1, 'output', cons2, 'input')
    results = simple_process.process_and_return(graph, inputs={prod: 5})
    tools.eq_({cons2.id: {'output': list(range(1, 6))}}, results)


def testSquare():
    graph = WorkflowGraph()
    prod = TestProducer(2)
    cons1 = TestOneInOneOut()
    cons2 = TestOneInOneOut()

    last = TestTwoInOneOut()