"""
ODIN: ODin Isn't Numpy
"""

from distarray.client import DistArrayProxy
from operator import add


def flatten(lst):
    """ Given a list of lists, return a flattened list.

    Only flattens one level.  For example,

    >>> flatten(zip(['a', 'b', 'c'], [1, 2, 3]))
    ['a', 1, 'b', 2, 'c', 3]

    >>> flatten([[1, 2], [3, 4, 5], [[5], [6], [7]]])
    [1, 2, 3, 4, 5, [5], [6], [7]]
    """
    if len(lst) == 0:
        return []
    else:
        return list(reduce(add, lst))


def key_and_push_args(context, args):
    arg_keys = []
    for arg in args:
        if isinstance(arg, DistArrayProxy):
            # if a DistArrayProxy, use its existing key
            arg_keys.append(arg.key)
            is_self = (context == arg.context)
            err_msg_fmt = "distarray context mismatch: {} {}"
            assert is_self, err_msg_fmt.format(context, arg.context)
        else:
            # if not a DistArrayProxy, key it and push it to engines
            arg_keys.extend(context._key_and_push(arg))
    return arg_keys


def key_and_push_kwargs(context, kwargs):
    kwarg_names = []
    kwarg_keys = []
    for kwarg_name, kwarg_value in kwargs.items():
        kwarg_names.append(kwarg_name)
        if isinstance(kwarg_value, DistArrayProxy):
            # if a DistArrayProxy, use its existing key
            kwarg_keys.append(kwarg_value.key)
            is_self = (context == kwarg_value.context)
            err_msg_fmt = "distarray context mismatch: {} {}"
            assert is_self, err_msg_fmt.format(context, kwarg_value.context)
        else:
            # if not a DistArrayProxy, key it and push it to engines
            kwarg_keys.extend(context._key_and_push(kwarg_value))
    return kwarg_names, kwarg_keys


def local(context):
    """ Decorator indicating a function is run locally on engines.

    Parameters
    ----------
    context : DistArrayContext

    Returns
    -------
    fn : function wrapped to run locally on engines
    """

    def wrap(fn):

        func_key = context._key_and_push(fn)[0]
        result_key = context._generate_key()

        def inner(*args, **kwargs):
            # generate keys for each parameter
            # push to engines if not a DistArrayProxy
            arg_keys = key_and_push_args(context, args)
            kwarg_names, kwarg_keys = key_and_push_kwargs(context, kwargs)

            # build up a python statement as a string
            args_fmt = ','.join(['{}'] * len(arg_keys))
            kwargs_fmt = ','.join(['{}={}'] * len(kwarg_keys))
            statement_fmt = '{} = {}(' + args_fmt + ',' + kwargs_fmt + ')'
            replacement_values = ([result_key, func_key] + arg_keys +
                                  flatten(zip(kwarg_names, kwarg_keys)))
            statement = statement_fmt.format(*replacement_values)

            # execute it locally and return the result as a DistArrayProxy
            context._execute(statement)
            return DistArrayProxy(result_key, context)

        return inner

    return wrap

@@ -515,20 +515,393 @@
 ext, arg
-s
+list
 ):%0A
+    %22%22%22 For each arg in arglist, get or generate a key (UUID).%0A%0A    For DistArrayProxy objects, just get the existing key.  For%0A    everything else, generate a key and push the value to the engines%0A%0A    Parameters%0A    ----------%0A    context : DistArrayContext%0A    arglist : List of objects to key and/or push%0A%0A    Returns%0A    -------%0A    arg_keys : list of keys%0A    %22%22%22%0A
     arg_
@@ -928,17 +928,20 @@
 g in arg
-s
+list
 :%0A      
@@ -1414,716 +1414,8 @@
 s%0A%0A%0A
-def key_and_push_kwargs(context, kwargs):%0A    kwarg_names = %5B%5D%0A    kwarg_keys = %5B%5D%0A    for kwarg_name, kwarg_value in kwargs.items():%0A        kwarg_names.append(kwarg_name)%0A        if isinstance(kwarg_value, DistArrayProxy):%0A            # if a DistArrayProxy, use its existing key%0A            kwarg_keys.append(kwarg_value.key)%0A            is_self = (context == kwarg_value.context)%0A            err_msg_fmt = %22distarray context mismatch: %7B%7D %7B%7D%22%0A            assert is_self, err_msg_fmt.format(context, kwarg_value.context)%0A        else:%0A            # if not a DistArrayProxy, key it and push it to engines%0A            kwarg_keys.extend(context._key_and_push(kwarg_value))%0A    return kwarg_names, kwarg_keys%0A%0A%0A
 def 
@@ -1966,33 +1966,60 @@
      kwarg_names
-,
+ = kwargs.keys()%0A           
  kwarg_keys = ke
@@ -2029,18 +2029,16 @@
 nd_push_
-kw
 args(con
@@ -2049,16 +2049,25 @@
 , kwargs
+.values()
 )%0A%0A     
@@ -2226,32 +2226,90 @@
 en(kwarg_keys))%0A
+            fnargs_fmt = ','.join(%5Bargs_fmt, kwargs_fmt%5D)%0A
             stat
@@ -2320,16 +2320,25 @@
 t_fmt = 
+''.join(%5B
 '%7B%7D = %7B%7D
@@ -2343,19 +2343,20 @@
 %7B%7D('
- + 
+, fn
 args_fmt
  + '
@@ -2355,33 +2355,15 @@
 _fmt
- + ',' + kwargs_fmt +
+,
  ')'
+%5D)
 %0A   

Refactor @local decorator.