# encoding: utf-8
# ---------------------------------------------------------------------------
#  Copyright (C) 2008-2014, IPython Development Team and Enthought, Inc.
#  Distributed under the terms of the BSD License.  See COPYING.rst.
# ---------------------------------------------------------------------------

"""
Distributed unfuncs for distributed arrays.
"""

import numpy

from distarray.error import ContextError
from distarray.client import DistArray


__all__ = []  # unary_names and binary_names added to __all__ below.

# numpy unary operations to wrap
unary_names = ('absolute', 'arccos', 'arccosh', 'arcsin', 'arcsinh', 'arctan',
               'arctanh', 'conjugate', 'cos', 'cosh', 'exp', 'expm1', 'invert',
               'log', 'log10', 'log1p', 'negative', 'reciprocal', 'rint',
               'sign', 'sin', 'sinh', 'sqrt', 'square', 'tan', 'tanh')

# numpy binary operations to wrap
binary_names = ('add', 'arctan2', 'bitwise_and', 'bitwise_or', 'bitwise_xor',
                'divide', 'floor_divide', 'fmod', 'hypot', 'left_shift', 'mod',
                'multiply', 'power', 'remainder', 'right_shift', 'subtract',
                'true_divide', 'less', 'less_equal', 'equal', 'not_equal',
                'greater', 'greater_equal',)

for func_name in unary_names + binary_names:
    __all__.append(func_name)


def unary_proxy(name):
    def proxy_func(a, *args, **kwargs):
        context = determine_context(a)
        new_key = context._generate_key()
        if 'casting' in kwargs:
            exec_str = "%s = distarray.local.%s(%s, casting='%s')"
            exec_str %= (new_key, name, a.key, kwargs['casting'])
        else:
            exec_str = '%s = distarray.local.%s(%s)'
            exec_str %= (new_key, name, a.key)

        context._execute(exec_str)
        return DistArray.from_localarrays(new_key,
                                          distribution=a.distribution)
    return proxy_func


def binary_proxy(name):
    def proxy_func(a, b, *args, **kwargs):
        context = determine_context(a, b)
        is_a_dap = isinstance(a, DistArray)
        is_b_dap = isinstance(b, DistArray)
        if is_a_dap and is_b_dap:
            a_key = a.key
            b_key = b.key
        elif is_a_dap and numpy.isscalar(b):
            a_key = a.key
            b_key = context._key_and_push(b)[0]
        elif is_b_dap and numpy.isscalar(a):
            a_key = context._key_and_push(a)[0]
            b_key = b.key
        else:
            raise TypeError('only DistArray or scalars are accepted')
        new_key = context._generate_key()

        if 'casting' in kwargs:
            exec_str = "%s = distarray.local.%s(%s,%s, casting='%s')"
            exec_str %= (new_key, name, a_key, b_key, kwargs['casting'])
        else:
            exec_str = '%s = distarray.local.%s(%s,%s)'
            exec_str %= (new_key, name, a_key, b_key)

        context._execute(exec_str)
        return DistArray.from_localarrays(new_key, context=context)
    return proxy_func


def determine_context(*args):
    """ Determine a context from a functions arguments."""

    contexts = []
    # inspect args for a context
    for arg in args:
        if isinstance(arg, DistArray):
            contexts.append(arg.context)

    # check the args had a context
    if contexts == []:
        raise TypeError('Function must take DistArray or Context objects.')

    # check that all contexts are equal
    if not contexts.count(contexts[0]) == len(contexts):
        msg = ("Arguments must use the same Context (given arguments of "
               "type %r)")
        msg %= (tuple(set(contexts)),)
        raise ContextError(msg)

    return contexts[0]

# Define the functions dynamically at the module level.
for name in unary_names:
    globals()[name] = unary_proxy(name)

for name in binary_names:
    globals()[name] = binary_proxy(name)

@@ -2214,32 +2214,74 @@
   b_key = b.key%0A
+            distribution = a.distribution%0A
         elif is_
@@ -2309,16 +2309,16 @@
 lar(b):%0A
-
         
@@ -2379,24 +2379,66 @@
 _push(b)%5B0%5D%0A
+            distribution = a.distribution%0A
         elif
@@ -2514,24 +2514,24 @@
 _push(a)%5B0%5D%0A
-
             
@@ -2536,32 +2536,74 @@
   b_key = b.key%0A
+            distribution = b.distribution%0A
         else:%0A  
@@ -3103,23 +3103,33 @@
 ey, 
-context=context
+distribution=distribution
 )%0A  

Use new version of from_localarrays in binary_proxy