import numarray as num
import _combine
import operator as _operator

_PROTOTYPE = 0

if _PROTOTYPE:
    def _combine( arrays, output, nlow, nhigh, badmasks, kind):
        if arrays[0].rank == 1:
            for i in xrange(arrays[0].shape[0]):
                if badmasks is not None:
                    seq0 = [ a[i] for a,b in zip(arrays, badmasks) if not b[i] ]
                else:
                    seq0 = [ a[i] for a in arrays ]
                seq0.sort()
                if nhigh == 0:
                    seq = seq0[nlow:]
                else:
                    seq = seq0[nlow:-nhigh]
                N = len(seq)
                if kind == "median":
                    midpoint = N//2
                    if (N % 2):
                        median = seq[ midpoint ]
                    else:
                        median = (seq[ midpoint ]+seq[ midpoint-1 ])/2.0
                    output[i] = median
                elif kind == "average":
                    if N:
                        average = reduce(_operator.add, seq, 0) / (N*1.0)
                    else:
                        average = 0
                    output[i] = average
                else:
                    raise ValueError("Unknown combination type '%s'" % kind)
        else:
            for i in xrange(arrays[0].shape[0]):
                subarrays = [ a[i] for a in arrays ]
                if badmasks is not None:
                    subbadmasks = [ bm[i] for bm in badmasks ]
                else:
                    subbadmasks = None
                _combine(subarrays, output[i], nlow, nhigh, subbadmasks, kind)


else:
    _combine = _combine.combine

def median( arrays, output=None, outtype=None, nlow=0, nhigh=0, badmasks=None):
    """median() nominally computes the median pixels for a stack of
    identically shaped images.

    arrays     specifies a sequence of inputs arrays, which are nominally a
               stack of identically shaped images.

    output     may be used to specify the output array.  If none is specified,
               either arrays[0] is copied or a new array of type 'outtype'
               is created.

    outtype    specifies the type of the output array when no 'output' is
               specified.

    nlow       specifies the number of pixels to be excluded from median
               on the low end of the pixel stack.

    nhigh      specifies the number of pixels to be excluded from median
               on the high end of the pixel stack.

    badmasks   specifies boolean arrays corresponding to 'arrays', where true
               indicates that a particular pixel is not to be included in the
               median calculation.

    >>> a = num.arange(4, shape=(2,2))
    >>> arrays = [a*16, a*4, a*2, a*8]
    >>> median(arrays)
    array([[ 0,  6],
           [12, 18]])
    >>> median(arrays, nhigh=1)
    array([[ 0,  4],
           [ 8, 12]])
    >>> median(arrays, nlow=1)
    array([[ 0,  8],
           [16, 24]])
    >>> median(arrays, outtype=num.Float32)
    array([[  0.,   6.],
           [ 12.,  18.]], type=Float32)
    >>> bm = num.zeros((4,2,2), type=num.Bool)
    >>> bm[2,...] = 1
    >>> median(arrays, badmasks=bm)
    array([[ 0,  8],
           [16, 24]])
    >>> median(arrays, badmasks=threshhold(arrays, high=25))
    array([[ 0,  6],
           [ 8, 12]])

"""
    shape = arrays[0].shape
    type = arrays[0].type()
    
    if output is None:
        if outtype is not None:
            out = arrays[0].astype(outtype)
        else:
            out = arrays[0].copy()
    else:
        out = output

    for a in tuple(arrays[1:])+(out,):
        if a.shape != shape:
            raise ValueError("all arrays must have identical shapes")

    _combine(arrays, out, nlow, nhigh, badmasks, "median")

    if output is None:
        return out

def average( arrays, output=None, outtype=None, nlow=0, nhigh=0, badmasks=None):
    """average() nominally computes the average pixel value for a stack of
    identically shaped images.

    arrays     specifies a sequence of inputs arrays, which are nominally a
               stack of identically shaped images.

    output     may be used to specify the output array.  If none is specified,
               either arrays[0] is copied or a new array of type 'outtype'
               is created.

    outtype    specifies the type of the output array when no 'output' is
               specified.

    nlow       specifies the number of pixels to be excluded from average
               on the low end of the pixel stack.

    nhigh      specifies the number of pixels to be excluded from average
               on the high end of the pixel stack.

    badmasks   specifies boolean arrays corresponding to 'arrays', where true
               indicates that a particular pixel is not to be included in the
               average calculation.

    >>> a = num.arange(4, shape=(2,2))
    >>> arrays = [a*16, a*4, a*2, a*8]
    >>> average(arrays)
    array([[ 0,  7],
           [15, 22]])
    >>> average(arrays, nhigh=1)
    array([[ 0,  4],
           [ 9, 14]])
    >>> average(arrays, nlow=1)
    array([[ 0,  9],
           [18, 28]])
    >>> average(arrays, outtype=num.Float32)
    array([[  0. ,   7.5],
           [ 15. ,  22.5]], type=Float32)
    >>> bm = num.zeros((4,2,2), type=num.Bool)
    >>> bm[2,...] = 1
    >>> average(arrays, badmasks=bm)
    array([[ 0,  9],
           [18, 28]])
    >>> average(arrays, badmasks=threshhold(arrays, high=25))
    array([[ 0,  7],
           [ 9, 14]])

    """
    shape = arrays[0].shape
    type = arrays[0].type()
    
    if output is None:
        if outtype is not None:
            out = arrays[0].astype(outtype)
        else:
            out = arrays[0].copy()
    else:
        out = output

    for a in tuple(arrays[1:])+(out,):
        if a.shape != shape:
            raise ValueError("all arrays must have identical shapes")

    _combine(arrays, out, nlow, nhigh, badmasks, "average")

    if output is None:
        return out

def threshhold(arrays, low=None, high=None, outputs=None):
    """threshhold() computes a boolean array 'outputs' with
    corresponding elements for each element of arrays.  The
    boolean value is true where each of the arrays values
    is < the low or >= the high threshholds.

    >>> a=num.arange(100, shape=(10,10))
    >>> threshhold(a, 1, 50)
    array([[1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], type=Bool)
    >>> threshhold([ range(10)]*10, 3, 7)
    array([[1, 1, 1, 0, 0, 0, 0, 1, 1, 1],
           [1, 1, 1, 0, 0, 0, 0, 1, 1, 1],
           [1, 1, 1, 0, 0, 0, 0, 1, 1, 1],
           [1, 1, 1, 0, 0, 0, 0, 1, 1, 1],
           [1, 1, 1, 0, 0, 0, 0, 1, 1, 1],
           [1, 1, 1, 0, 0, 0, 0, 1, 1, 1],
           [1, 1, 1, 0, 0, 0, 0, 1, 1, 1],
           [1, 1, 1, 0, 0, 0, 0, 1, 1, 1],
           [1, 1, 1, 0, 0, 0, 0, 1, 1, 1],
           [1, 1, 1, 0, 0, 0, 0, 1, 1, 1]], type=Bool)
    >>> threshhold(a, high=50)
    array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], type=Bool)
    >>> threshhold(a, low=50)
    array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], type=Bool)
    """

    if not isinstance(arrays[0],  num.NumArray):
        return threshhold( num.inputarray(arrays), low, high, outputs)
    
    if outputs is None:
        outs = num.zeros(shape=(len(arrays),)+arrays[0].shape,
                         type=num.Bool)
    else:
        outs = outputs

    for i in range(len(arrays)):
        a, out = arrays[i], outs[i]
        out[:] = 0

        if high is not None:
            num.greater_equal(a, high, out)
            if low is not None:
                num.logical_or(out, a < low, out)
        else:
            if low is not None:
                num.less(a, low, out)

    if outputs is None:
        return outs

def _bench():
    """time a 10**6 element median"""
    import time
    a = num.arange(10**6, shape=(1000, 1000))
    arrays = [a*2, a*64, a*16, a*8]
    t0 = time.clock()
    median(arrays)
    print "maskless:", time.clock()-t0
    
    a = num.arange(10**6, shape=(1000, 1000))
    arrays = [a*2, a*64, a*16, a*8]
    t0 = time.clock()
    median(arrays, badmasks=num.zeros((1000,1000), type=num.Bool))
    print "masked:", time.clock()-t0