# 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 core module for dispel4py.
'''

import uuid
# Connection-level dict elements
NAME = 'name'
TYPE = 'type'
DESCRIPTION = 'descr'
META = 'meta'
GROUPING = 'grouping'
WRITER = 'writer'


class GenericPE(object):
    '''
Base class for Dispel4Py processing elements (PEs). Custom PEs are expected to
extend this class and override the 'process' function.

Custom PEs must override :py:func:`~dispel4py.core.GenericPE.__init__` to
declare the inputs and outputs that can be connected within the workflow
graph, by defining a NAME and possibly a TYPE.
The type of a connection is specific to the enactment system.
In the example below the target system is Storm and the type declares what
kind of tuples are produced::
    out1[TYPE] = ['timestamp', 'origin', 'streams']

In some cases, the output types are determined dynamically and depend on the
input types, for example when implementing a filter which consumes any type
of blocks but the type of the output is the same as the type of the input.
The graph framework supports this by propagating types across the workflow
before enactment and providing each PE with the input types that it can expect
in the method::
    setInputTypes(self, types)
which can be overridden to deduce output types from input types or to raise an
error if the types are not acceptable.
The PE may then override the method::
    getOutputTypes(self)
to declare the output types that it produces.
In the example of a filter PE this method would return the input types.

Custom PEs may implement the method preprocess() to initialise variables or
data before processing commences.

Example implementation::

    import traceback
    import cStringIO
    import base64
    from obspy.core import read,UTCDateTime,Stream,Trace
    from dispel4py.core import GenericPE, NAME, TYPE

    INPUT_NAME = 'input'
    OUTPUT_NAME = 'output'

    class AppendAndSynchronize(GenericPE):

        def __init__(self):
            GenericPE.__init__(self)
            self._add_input(INPUT_NAME)
            self._add_output(OUTPUT_NAME, ['timestamp', 'origin', 'streams'])

        def process(self, inputs):
            values = inputs[INPUT_NAME]
            parameters = values[0]
            origin = values[1]
            data = values[2:]
            if not data:
                self.error+= "No Data";
                raise Exception("No Data!")

            streams=list();
            while data:
                streamItem=data.pop(0);
                streams.append(eval(streamItem["data"]))

            # Reads the first file
            st = read(self.rootpath+"%s" % (streams[0].pop(0),))

            #Reads the following files
            while streams[0]:
                ff= "%s" % (streams[0].pop(0),)
                st += read(self.rootpath+ff)

            starttime="%s" % (parameters["starttime"])
            endtime="%s" % (parameters["endtime"])

            st=st.slice(UTCDateTime(starttime),UTCDateTime(endtime));
            streamtransfer={}
            if type(st) == Stream:
                memory_file = cStringIO.StringIO()
                mseed = st.write(memory_file, format="MSEED")
                streamtransfer={"data":base64.b64encode(memory_file.getvalue())}
            output = [ parameters, origin, streamtransfer ]

            return { OUTPUT_NAME : output }
    '''

    def __init__(self, numprocesses=1):
        self.inputconnections = {}
        self.outputconnections = {}
        self.wrapper = 'simple'
        self.pickleIgnore = []
        self.pickleIgnore = vars(self).keys()
        self.numprocesses = numprocesses

        def log(self, message):
            ''' To be implemented '''
            pass
        self.log = log
        self.name = self.__class__.__name__
        self.id = self.name + str(uuid.uuid4())

    def _add_input(self, name, grouping=None, tuple_type=None):
        '''
        Declares an input for this PE.
        This method may be used when initialising a PE instead of modifying
        :py:attr:`~dispel4py.core.GenericPE.inputconnections` directly.

        :param name: name of the input
        :param grouping: the grouping type that this input expects (optional)
        :param tuple_type: type of tuples accepted by this input (optional)
        '''
        self.inputconnections[name] = {NAME: name}
        if grouping:
            self.inputconnections[name][GROUPING] = grouping
        if tuple_type:
            self.inputconnections[name][TYPE] = tuple_type

    def _add_output(self, name, tuple_type=None):
        '''
        Declares an output for this PE.
        This method may be used when initialising a PE instead of modifying
        :py:attr:`~dispel4py.core.GenericPE.outputconnections` directly.

        :param name: name of the output
        :param tuple_type: type of tuples produced by this output (optional)
        '''
        self.outputconnections[name] = {NAME: name}
        if tuple_type:
            self.outputconnections[name][TYPE] = tuple_type

    def setInputTypes(self, types):
        '''
        Sets the input types of this PE, in the form of a dictionary.
        It is meant to be overridden, e.g. if output types depend on input.

        .. note::

            This method is always called before
            :py:func:`~dispel4py.core.GenericPE.getOutputTypes`.

        :param types: object types for each input stream
        :type types: dictionary mapping input name to input type

        Usage example::

            pe.setInputTypes({'input1':['t1', 't2', 't3'], \
                              'input2':['t4', 't5']})
        '''
        pass

    def getOutputTypes(self):
        '''
        Returns the output types of this PE, in the form of a dictionary.
        This method may be overridden if output types are not static and
        depend on input types.

        .. note::

            This method is only called after the input types have been
            initialised in :py:func:`~dispel4py.core.GenericPE.setInputTypes`.

        :rtype: a dictionary mapping each output name to its type

        By default it returns a dictionary of the types defined in the
        'outputconnections' instance variable.

        Usage example::

            def getOutputTypes(self):
                output = { 'output1' : myInputs['input1'],
                           'output2' : [ 'comment' ] }

        '''
        ret = {}
        # print '%s: %s' % (self.id, self.outputconnections)
        for name, output in self.outputconnections.iteritems():
            try:
                ret[name] = output[TYPE]
            except KeyError:
                raise Exception("%s: No output type defined for '%s'"
                                % (self.id, name))
        return ret

    def _preprocess(self):
        '''
        Subclasses may override this method for variable and data
        initialisation before data processing commences.
        '''
        None

    def preprocess(self):
        '''
        This method called once before processing commences.
        '''
        self._preprocess()

    def _process(self, data):
        '''
        To be overridden by a PE implementation subclass. The data input
        parameter may contain data in any format that is prepared by a PE
        base pattern class implementing
        :py:func:`~dispel4py.core.GenericPE.process`.
        By default, the data parameter contains the dictionary of all inputs
        and the return value must be the output data dictionary.
        PE base classes may override
        :py:func:`~dispel4py.core.GenericPE.process` as they wish to pre and
        postprocess data to and from any format.
        To produce more than one output block data can be written at any point
        during processing using the :py:func:`~dispel4py.core.GenericPE.write`
        method.
        '''
        None

    def process(self, inputs):
        '''
        The 'inputs' dictionary contains data from any or all of the streams
        that are connected to this PE, in any order. The return value of this
        function is a single output dictionary, with the names of the output
        streams as keys.
        PE base classes may override this method to allow preprocessing of data
        before calling :py:func:`~dispel4py.core.GenericPE._process` with the
        prepared data on the PE implementation.
        The data returned by the implementation may then be postprocessed
        before it is returned to the implementation framework.

        :param inputs: the input data for this iteration
        :type inputs: dictionary
        :rtype: a dictionary with the output data
        '''
        return self._process(inputs)

    def _postprocess(self):
        None

    def postprocess(self):
        '''
        This method is called once after the last block has been processed and
        a terminate message was sent to this PE.
        '''
        self._postprocess()

    def write(self, name, data, **kwargs):
        '''
        Allows for preprocessing of data to be written to the output pipe.
        This method should be overridden by PE base classes.
        '''
        self._write(name, data)

    def _write(self, name, data, **kwargs):
        '''
        This writes the 'data' to the output pipe with name 'name' of this PE.
        '''
        try:
            output = self.outputconnections[name]
        except KeyError:
            raise Exception("Can't write data: Unknown output connection\
                            '%s' for PE '%s'" % (name, type(self).__name__))
        output[WRITER].write(data)


class LockstepPE(GenericPE):
    '''
    Representation of a PE which consumes its input in lockstep. The inputs
    dictionary that is passed to the
    :py:func:`~dispel4py.GenericPE.LockStepPE.process` function is guaranteed
    to contain one data item from each of the connected input streams.
    '''

    def __init__(self):
        GenericPE.__init__(self)
        self.wrapper = 'lockstep'


class SourcePE(GenericPE):
    '''
    Representation of data-producing PE, i.e. a PE with no input connections.
    '''

    def __init__(self):
        GenericPE.__init__(self)
        self.wrapper = 'source'
