
import theano
import theano.tensor as T
import pickle
from collections import OrderedDict
import time

from .function_module import Function, WorkerFunction
from .collectives import shareds_registry
from . import exct
from .util import struct, PKL_FILE
# import ipdb
sync = None

synk_functions = list()


###############################################################################
#                                                                             #
#                       API for building Functions                            #
#                                                                             #
###############################################################################


def function(inputs, outputs=None, bcast_inputs=None, updates=None,
             givens=None, sliceable_shareds=None, **kwargs):
    """Replacement for ``theano.function()``, with a similar interface.  Builds
    underlying Theano functions, including support for function slicing.

    Args:
        inputs: as in Theano, to be scattered among workers
        outputs: as in Theano, with option to specify reduce operation (see
            notes below)
        bcast_inputs:  as inputs in Theano, to be broadcast to all workers
        updates: as in Theano, with option to specify reduct operation (see
            notes below)
        givens: as in Theano
        sliceable_shareds: any implicit inputs (Theano shared variables) acting
            as data-parallel data (i.e. to be subjected to the kwarg ``batch_s``
            and /or to function slicing) must be listed here
        **kwargs: passed on to all internal calls to ``theano.function()``

    Reduce Operations:
      Outputs: May be specified simply as Theano tensor variables, as in normal
      Theano, or as two-tuples, as in (var, reduce-op), where reduce-op can be:
      "avg", "sum", "max", "min", "prod", or None.  Default is "avg".

      Updates: May be specified as a list of two-tuples, as in normal Theano, or
      may include triples, as in (var, update, reduce-op).  Unlike for outputs,
      the reduce-op here applies only when using function slicing.  Every slice
      is computed using the original values, and the update is accumulated over
      the slices.  (This may impose some limits on the form of the update
      expression.) At the end of the function call, all updates are applied only
      locally, within each worker.  This provides clear control to user over
      when to communicate.

    Returns:
        sykhronos.function_module.Function: callable object, replacing a
        theano.Function

    Raises:
        RuntimeError: If Sykhronos not yet forked, or if already distributed
        TypeError: If incorrect format for arguments.
        ValueError: If entry in ``sliceable_shareds`` is not used in function,
            or for invalid reduce operation requested.
    """
    if not exct.state.forked:
        raise RuntimeError("Must fork before making functions for GPU.")
    if exct.state.distributed:
        raise RuntimeError("Cannot make new functions after distributing (for now).")

    if not isinstance(inputs, list):
        raise TypeError("Input 'inputs' must be list.")
    bcast_inputs = [] if bcast_inputs is None else bcast_inputs
    if not isinstance(bcast_inputs, list):
        raise TypeError("Input 'bcast_inputs' must be list if not None.")

    reg_outputs, gpu_outputs, to_cpu, output_modes = process_outputs(outputs)
    updates, update_vars, sliced_update_outs, update_modes = process_updates(updates)

    theano_function = theano.function(
        inputs=inputs + bcast_inputs,
        outputs=reg_outputs,
        updates=updates,