from IPython.core.debugger import Tracer
import afnumpy
import private_utils as pu

def outufunc(func):
    def wrapper(*args, **kws):
        out = kws.pop('out', None)
        ret = func(*args, **kws)
        if out is not None:
            if(out.ndim):
                out[:] = ret
            else:
                out[()] = ret
        return ret
    return wrapper

def iufunc(func):
    def wrapper(*args, **kws):
        if all(isinstance(A, afnumpy.ndarray) for A in args):
            bcast_args = afnumpy.broadcast_arrays(*args)
            if(bcast_args[0].shape is not args[0].shape):
                raise ValueError("non-broadcastable output operand with"
                                 " shape %s doesn't match the broadcast"
                                 " shape %s" % (args[0].shape, bcast_args[0].shape))
            args = bcast_args
            
        return func(*args, **kws)
    return wrapper

def ufunc(func):
    def wrapper(*args, **kws):
        if all(isinstance(A, afnumpy.ndarray) for A in args):
            args = afnumpy.broadcast_arrays(*args)            
        return func(*args, **kws)
    return wrapper

def reductufunc(func):
    def wrapper(a, axis=None, dtype=None, keepdims=False):        
        if not isinstance(axis, tuple):
            axis = (axis,)
        if(a.d_array):
            if axis[0] is None:
                a = a.flat
                axis = (0,)
            shape = list(a.shape)
            s = a.d_array
            # Do in decreasing axis order to avoid problems with the pop
            for ax in sorted(axis)[::-1]:
                s = func(a, s, pu.c2f(a.shape, ax))
                if keepdims:
                    shape[ax] = 1
                else:
                    shape.pop(ax)
            ret = afnumpy.ndarray(tuple(shape), dtype=pu.typemap(s.type()), 
                                  af_array=s)
            if(dtype is not None):
                ret = ret.astype(dtype)
            if(len(shape) == 0):
                ret = ret[()]
            return ret
        else:
            # Call the host array function instead
            return getattr(a.h_array,func.__name__)(axis=None, dtype=dtype, keepdims=keepdims)
    return wrapper

@@ -1318,107 +1318,403 @@
   if
-(a.d_array):%0A            if axis%5B0%5D is None:%0A                a = a.flat%0A                axis = (0,)
+ axis%5B0%5D is None:%0A            # Special case for speed%0A            a_flat = a.flat%0A            s = func(a_flat, a_flat.d_array, pu.c2f(a_flat.shape, 0))%0A            if keepdims:%0A                shape = (1,)*a.ndim%0A            else:%0A                shape = ()%0A            ret = afnumpy.ndarray(tuple(shape), dtype=pu.typemap(s.type()), %0A                                  af_array=s)%0A        else:
 %0A   
@@ -1885,16 +1885,88 @@
 %5B::-1%5D:%0A
+                # Optimization%0A                if(a.shape%5Bax%5D %3E 1):%0A    
         
@@ -2009,16 +2009,16 @@
 e, ax))%0A
-
         
@@ -2259,20 +2259,16 @@
         
-    
 if(dtype
@@ -2278,28 +2278,24 @@
  not None):%0A
-    
             
@@ -2314,28 +2314,24 @@
 type(dtype)%0A
-    
         if(l
@@ -2351,36 +2351,32 @@
 0):%0A            
-    
 ret = ret%5B()%5D%0A  
@@ -2377,28 +2377,24 @@
 ()%5D%0A        
-    
 return ret%0A 
@@ -2396,168 +2396,8 @@
 ret%0A
-        else:%0A            # Call the host array function instead%0A            return getattr(a.h_array,func.__name__)(axis=None, dtype=dtype, keepdims=keepdims)%0A
     

correct reduction decorator