__all__ = ['atleast_1d','atleast_2d','atleast_3d','vstack','hstack',
           'column_stack','row_stack', 'dstack','array_split','split','hsplit',
           'vsplit','dsplit','apply_over_axes','expand_dims',
           'apply_along_axis', 'kron', 'tile', 'get_array_wrap']

import numpy.core.numeric as _nx
from numpy.core.numeric import asarray, zeros, newaxis, outer, \
     concatenate, isscalar, array, asanyarray
from numpy.core.fromnumeric import product, reshape

def apply_along_axis(func1d,axis,arr,*args):
    """ Execute func1d(arr[i],*args) where func1d takes 1-D arrays
        and arr is an N-d array.  i varies so as to apply the function
        along the given axis for each 1-d subarray in arr.
    """
    arr = asarray(arr)
    nd = arr.ndim
    if axis < 0:
        axis += nd
    if (axis >= nd):
        raise ValueError("axis must be less than arr.ndim; axis=%d, rank=%d."
            % (axis,nd))
    ind = [0]*(nd-1)
    i = zeros(nd,'O')
    indlist = range(nd)
    indlist.remove(axis)
    i[axis] = slice(None,None)
    outshape = asarray(arr.shape).take(indlist)
    i.put(indlist, ind)
    res = func1d(arr[tuple(i.tolist())],*args)
    #  if res is a number, then we have a smaller output array
    if isscalar(res):
        outarr = zeros(outshape,asarray(res).dtype)
        outarr[tuple(ind)] = res
        Ntot = product(outshape)
        k = 1
        while k < Ntot:
            # increment the index
            ind[-1] += 1
            n = -1
            while (ind[n] >= outshape[n]) and (n > (1-nd)):
                ind[n-1] += 1
                ind[n] = 0
                n -= 1
            i.put(indlist,ind)
            res = func1d(arr[tuple(i.tolist())],*args)
            outarr[tuple(ind)] = res
            k += 1
        return outarr
    else:
        Ntot = product(outshape)
        holdshape = outshape
        outshape = list(arr.shape)
        outshape[axis] = len(res)
        outarr = zeros(outshape,asarray(res).dtype)
        outarr[tuple(i.tolist())] = res
        k = 1
        while k < Ntot:
            # increment the index
            ind[-1] += 1
            n = -1
            while (ind[n] >= holdshape[n]) and (n > (1-nd)):
                ind[n-1] += 1
                ind[n] = 0
                n -= 1
            i.put(indlist, ind)
            res = func1d(arr[tuple(i.tolist())],*args)
            outarr[tuple(i.tolist())] = res
            k += 1
        return outarr


def apply_over_axes(func, a, axes):
    """Apply a function repeatedly over multiple axes, keeping the same shape
    for the resulting array.

    func is called as res = func(a, axis).  The result is assumed
    to be either the same shape as a or have one less dimension.
    This call is repeated for each axis in the axes sequence.
    """
    val = asarray(a)
    N = a.ndim
    if array(axes).ndim == 0:
        axes = (axes,)
    for axis in axes:
        if axis < 0: axis = N + axis
        args = (val, axis)
        res = func(*args)
        if res.ndim == val.ndim:
            val = res
        else:
            res = expand_dims(res,axis)
            if res.ndim == val.ndim:
                val = res
            else:
                raise ValueError, "function is not returning"\
                      " an array of correct shape"
    return val

def expand_dims(a, axis):
    """Expand the shape of a by including newaxis before given axis.
    """
    a = asarray(a)
    shape = a.shape
    if axis < 0:
        axis = axis + len(shape) + 1
    return a.reshape(shape[:axis] + (1,) + shape[axis:])


def atleast_1d(*arys):
    """ Force a sequence of arrays to each be at least 1D.

         Description:
            Force an array to be at least 1D.  If an array is 0D, the
            array is converted to a single row of values.  Otherwise,
            the array is unaltered.
         Arguments:
            *arys -- arrays to be converted to 1 or more dimensional array.
         Returns:
            input array converted to at least 1D array.
    """
    res = []
    for ary in arys:
        res.append(array(ary,copy=False,subok=True,ndmin=1))
    if len(res) == 1:
        return res[0]
    else:
        return res

def atleast_2d(*arys):
    """ Force a sequence of arrays to each be at least 2D.

         Description:
            Force an array to each be at least 2D.  If the array
            is 0D or 1D, the array is converted to a single
            row of values.  Otherwise, the array is unaltered.
         Arguments:
            arys -- arrays to be converted to 2 or more dimensional array.
         Returns:
            input array converted to at least 2D array.
    """
    res = []
    for ary in arys:
        res.append(array(ary,copy=False,subok=True,ndmin=2))
    if len(res) == 1:
        return res[0]
    else:
        return res

def atleast_3d(*arys):
    """ Force a sequence of arrays to each be at least 3D.

         Description:
            Force an array each be at least 3D.  If the array is 0D or 1D,
            the array is converted to a single 1xNx1 array of values where
            N is the orginal length of the array. If the array is 2D, the
            array is converted to a single MxNx1 array of values where MxN
            is the orginal shape of the array. Otherwise, the array is
            unaltered.
         Arguments:
            arys -- arrays to be converted to 3 or more dimensional array.
         Returns:
            input array converted to at least 3D array.
    """
    res = []
    for ary in arys:
        ary = asarray(ary)
        if len(ary.shape) == 0:
            result = ary.reshape(1,1,1)
        elif len(ary.shape) == 1:
            result = ary[newaxis,:,newaxis]
        elif len(ary.shape) == 2:
            result = ary[:,:,newaxis]
        else:
            result = ary
        res.append(result)
    if len(res) == 1:
        return res[0]
    else:
        return res


def vstack(tup):
    """ Stack arrays in sequence vertically (row wise)

        Description:
            Take a sequence of arrays and stack them vertically
            to make a single array.  All arrays in the sequence
            must have the same shape along all but the first axis.
            vstack will rebuild arrays divided by vsplit.
        Arguments:
            tup -- sequence of arrays.  All arrays must have the same
                   shape.
        Examples:
            >>> import numpy
            >>> a = array((1,2,3))
            >>> b = array((2,3,4))
            >>> numpy.vstack((a,b))
            array([[1, 2, 3],
                   [2, 3, 4]])
            >>> a = array([[1],[2],[3]])
            >>> b = array([[2],[3],[4]])
            >>> numpy.vstack((a,b))
            array([[1],
                   [2],
                   [3],
                   [2],
                   [3],
                   [4]])

    """
    return _nx.concatenate(map(atleast_2d,tup),0)

def hstack(tup):
    """ Stack arrays in sequence horizontally (column wise)

        Description:
            Take a sequence of arrays and stack them horizontally
            to make a single array.  All arrays in the sequence
            must have the same shape along all but the second axis.
            hstack will rebuild arrays divided by hsplit.
        Arguments:
            tup -- sequence of arrays.  All arrays must have the same
                   shape.
        Examples:
            >>> import numpy
            >>> a = array((1,2,3))
            >>> b = array((2,3,4))
            >>> numpy.hstack((a,b))
            array([1, 2, 3, 2, 3, 4])
            >>> a = array([[1],[2],[3]])
            >>> b = array([[2],[3],[4]])
            >>> numpy.hstack((a,b))
            array([[1, 2],
                   [2, 3],
                   [3, 4]])

    """
    return _nx.concatenate(map(atleast_1d,tup),1)

row_stack = vstack

def column_stack(tup):
    """ Stack 1D arrays as columns into a 2D array

        Description:
            Take a sequence of 1D arrays and stack them as columns
            to make a single 2D array.  All arrays in the sequence
            must have the same first dimension.  2D arrays are
            stacked as-is, just like with hstack.  1D arrays are turned
            into 2D columns first.

        Arguments:
            tup -- sequence of 1D or 2D arrays.  All arrays must have the same
                   first dimension.
        Examples:
            >>> import numpy
            >>> a = array((1,2,3))
            >>> b = array((2,3,4))
            >>> numpy.column_stack((a,b))
            array([[1, 2],
                   [2, 3],
                   [3, 4]])

    """
    arrays = []
    for v in tup:
        arr = array(v,copy=False,subok=True)
        if arr.ndim < 2:
            arr = array(arr,copy=False,subok=True,ndmin=2).T
        arrays.append(arr)
    return _nx.concatenate(arrays,1)

def dstack(tup):
    """ Stack arrays in sequence depth wise (along third dimension)

    Description:
        Take a sequence of arrays and stack them along the third axis.
        All arrays in the sequence must have the same shape along all
        but the third axis.  This is a simple way to stack 2D arrays
        (images) into a single 3D array for processing.
        dstack will rebuild arrays divided by dsplit.
    Arguments:
        tup -- sequence of arrays.  All arrays must have the same
               shape.
    Examples:
        >>> import numpy
        >>> a = array((1,2,3))
        >>> b = array((2,3,4))
        >>> numpy.dstack((a,b))
        array([[[1, 2],
                [2, 3],
                [3, 4]]])
        >>> a = array([[1],[2],[3]])
        >>> b = array([[2],[3],[4]])
        >>> numpy.dstack((a,b))
        array([[[1, 2]],
        <BLANKLINE>
               [[2, 3]],
        <BLANKLINE>
               [[3, 4]]])

    """
    return _nx.concatenate(map(atleast_3d,tup),2)

def _replace_zero_by_x_arrays(sub_arys):
    for i in range(len(sub_arys)):
        if len(_nx.shape(sub_arys[i])) == 0:
            sub_arys[i] = _nx.array([])
        elif _nx.sometrue(_nx.equal(_nx.shape(sub_arys[i]),0)):
            sub_arys[i] = _nx.array([])
    return sub_arys

def array_split(ary,indices_or_sections,axis = 0):
    """ Divide an array into a list of sub-arrays.

        Description:
           Divide ary into a list of sub-arrays along the
           specified axis.  If indices_or_sections is an integer,
           ary is divided into that many equally sized arrays.
           If it is impossible to make an equal split, each of the
           leading arrays in the list have one additional member.  If
           indices_or_sections is a list of sorted integers, its
           entries define the indexes where ary is split.

        Arguments:
           ary -- N-D array.
              Array to be divided into sub-arrays.
           indices_or_sections -- integer or 1D array.
              If integer, defines the number of (close to) equal sized
              sub-arrays.  If it is a 1D array of sorted indices, it
              defines the indexes at which ary is divided.  Any empty
              list results in a single sub-array equal to the original
              array.
           axis -- integer. default=0.
              Specifies the axis along which to split ary.
        Caveats:
           Currently, the default for axis is 0.  This
           means a 2D array is divided into multiple groups
           of rows.  This seems like the appropriate default,
    """
    try:
        Ntotal = ary.shape[axis]
    except AttributeError:
        Ntotal = len(ary)
    try: # handle scalar case.
        Nsections = len(indices_or_sections) + 1
        div_points = [0] + list(indices_or_sections) + [Ntotal]
    except TypeError: #indices_or_sections is a scalar, not an array.
        Nsections = int(indices_or_sections)
        if Nsections <= 0:
            raise ValueError, 'number sections must be larger than 0.'
        Neach_section,extras = divmod(Ntotal,Nsections)
        section_sizes = [0] + \
                        extras * [Neach_section+1] + \
                        (Nsections-extras) * [Neach_section]
        div_points = _nx.array(section_sizes).cumsum()

    sub_arys = []
    sary = _nx.swapaxes(ary,axis,0)
    for i in range(Nsections):
        st = div_points[i]; end = div_points[i+1]
        sub_arys.append(_nx.swapaxes(sary[st:end],axis,0))

    # there is a wierd issue with array slicing that allows
    # 0x10 arrays and other such things.  The following cluge is needed
    # to get around this issue.
    sub_arys = _replace_zero_by_x_arrays(sub_arys)
    # end cluge.

    return sub_arys

def split(ary,indices_or_sections,axis=0):
    """ Divide an array into a list of sub-arrays.

        Description:
           Divide ary into a list of sub-arrays along the
           specified axis.  If indices_or_sections is an integer,
           ary is divided into that many equally sized arrays.
           If it is impossible to make an equal split, an error is
           raised.  This is the only way this function differs from
           the array_split() function. If indices_or_sections is a
           list of sorted integers, its entries define the indexes
           where ary is split.

        Arguments:
           ary -- N-D array.
              Array to be divided into sub-arrays.
           indices_or_sections -- integer or 1D array.
              If integer, defines the number of (close to) equal sized
              sub-arrays.  If it is a 1D array of sorted indices, it
              defines the indexes at which ary is divided.  Any empty
              list results in a single sub-array equal to the original
              array.
           axis -- integer. default=0.
              Specifies the axis along which to split ary.
        Caveats:
           Currently, the default for axis is 0.  This
           means a 2D array is divided into multiple groups
           of rows.  This seems like the appropriate default
    """
    try: len(indices_or_sections)
    except TypeError:
        sections = indices_or_sections
        N = ary.shape[axis]
        if N % sections:
            raise ValueError, 'array split does not result in an equal division'
    res = array_split(ary,indices_or_sections,axis)
    return res

def hsplit(ary,indices_or_sections):
    """ Split ary into multiple columns of sub-arrays

        Description:
            Split a single array into multiple sub arrays.  The array is
            divided into groups of columns.  If indices_or_sections is
            an integer, ary is divided into that many equally sized sub arrays.
            If it is impossible to make the sub-arrays equally sized, the
            operation throws a ValueError exception. See array_split and
            split for other options on indices_or_sections.
        Arguments:
           ary -- N-D array.
              Array to be divided into sub-arrays.
           indices_or_sections -- integer or 1D array.
              If integer, defines the number of (close to) equal sized
              sub-arrays.  If it is a 1D array of sorted indices, it
              defines the indexes at which ary is divided.  Any empty
              list results in a single sub-array equal to the original
              array.
        Returns:
            sequence of sub-arrays.  The returned arrays have the same
            number of dimensions as the input array.
        Related:
            hstack, split, array_split, vsplit, dsplit.
        Examples:
            >>> import numpy
            >>> a= array((1,2,3,4))
            >>> numpy.hsplit(a,2)
            [array([1, 2]), array([3, 4])]
            >>> a = array([[1,2,3,4],[1,2,3,4]])
            >>> hsplit(a,2)
            [array([[1, 2],
                   [1, 2]]), array([[3, 4],
                   [3, 4]])]

    """
    if len(_nx.shape(ary)) == 0:
        raise ValueError, 'hsplit only works on arrays of 1 or more dimensions'
    if len(ary.shape) > 1:
        return split(ary,indices_or_sections,1)
    else:
        return split(ary,indices_or_sections,0)

def vsplit(ary,indices_or_sections):
    """ Split ary into multiple rows of sub-arrays

        Description:
            Split a single array into multiple sub arrays.  The array is
            divided into groups of rows.  If indices_or_sections is
            an integer, ary is divided into that many equally sized sub arrays.
            If it is impossible to make the sub-arrays equally sized, the
            operation throws a ValueError exception. See array_split and
            split for other options on indices_or_sections.
        Arguments:
           ary -- N-D array.
              Array to be divided into sub-arrays.
           indices_or_sections -- integer or 1D array.
              If integer, defines the number of (close to) equal sized
              sub-arrays.  If it is a 1D array of sorted indices, it
              defines the indexes at which ary is divided.  Any empty
              list results in a single sub-array equal to the original
              array.
        Returns:
            sequence of sub-arrays.  The returned arrays have the same
            number of dimensions as the input array.
        Caveats:
           How should we handle 1D arrays here?  I am currently raising
           an error when I encounter them.  Any better approach?

           Should we reduce the returned array to their minium dimensions
           by getting rid of any dimensions that are 1?
        Related:
            vstack, split, array_split, hsplit, dsplit.
        Examples:
            import numpy
            >>> a = array([[1,2,3,4],
            ...            [1,2,3,4]])
            >>> numpy.vsplit(a,2)
            [array([[1, 2, 3, 4]]), array([[1, 2, 3, 4]])]

    """
    if len(_nx.shape(ary)) < 2:
        raise ValueError, 'vsplit only works on arrays of 2 or more dimensions'
    return split(ary,indices_or_sections,0)

def dsplit(ary,indices_or_sections):
    """ Split ary into multiple sub-arrays along the 3rd axis (depth)

        Description:
            Split a single array into multiple sub arrays.  The array is
            divided into groups along the 3rd axis.  If indices_or_sections is
            an integer, ary is divided into that many equally sized sub arrays.
            If it is impossible to make the sub-arrays equally sized, the
            operation throws a ValueError exception. See array_split and
            split for other options on indices_or_sections.
        Arguments:
           ary -- N-D array.
              Array to be divided into sub-arrays.
           indices_or_sections -- integer or 1D array.
              If integer, defines the number of (close to) equal sized
              sub-arrays.  If it is a 1D array of sorted indices, it
              defines the indexes at which ary is divided.  Any empty
              list results in a single sub-array equal to the original
              array.
        Returns:
            sequence of sub-arrays.  The returned arrays have the same
            number of dimensions as the input array.
        Caveats:
           See vsplit caveats.
        Related:
            dstack, split, array_split, hsplit, vsplit.
        Examples:
            >>> a = array([[[1,2,3,4],[1,2,3,4]]])
            >>> dsplit(a,2)
            [array([[[1, 2],
                    [1, 2]]]), array([[[3, 4],
                    [3, 4]]])]

    """
    if len(_nx.shape(ary)) < 3:
        raise ValueError, 'vsplit only works on arrays of 3 or more dimensions'
    return split(ary,indices_or_sections,2)

def get_array_wrap(*args):
    """Find the wrapper for the array with the highest priority.

    In case of ties, leftmost wins. If no wrapper is found, return None
    """
    wrappers = [(getattr(x, '__array_priority__', 0), -i,
                 x.__array_wrap__) for i, x in enumerate(args)
                                   if hasattr(x, '__array_wrap__')]
    wrappers.sort()
    if wrappers:
        return wrappers[-1][-1]
    return None

def kron(a,b):
    """kronecker product of a and b

    Kronecker product of two arrays is block array
    [[ a[ 0 ,0]*b, a[ 0 ,1]*b, ... , a[ 0 ,n-1]*b  ],
     [ ...                                   ...   ],
     [ a[m-1,0]*b, a[m-1,1]*b, ... , a[m-1,n-1]*b  ]]
    """
    wrapper = get_array_wrap(a, b)
    b = asanyarray(b)
    a = array(a,copy=False,subok=True,ndmin=b.ndim)
    ndb, nda = b.ndim, a.ndim
    if (nda == 0 or ndb == 0):
        return _nx.multiply(a,b)
    as_ = a.shape
    bs = b.shape
    if not a.flags.contiguous:
        a = reshape(a, as_)
    if not b.flags.contiguous:
        b = reshape(b, bs)
    nd = ndb
    if (ndb != nda):
        if (ndb > nda):
            as_ = (1,)*(ndb-nda) + as_
        else:
            bs = (1,)*(nda-ndb) + bs
            nd = nda
    result = outer(a,b).reshape(as_+bs)
    axis = nd-1
    for _ in xrange(nd):
        result = concatenate(result, axis=axis)
    if wrapper is not None:
        result = wrapper(result)
    return result


def tile(A, reps):
    """Repeat an array the number of times given in the integer tuple, reps.

    If reps has length d, the result will have dimension of max(d, A.ndim).
    If reps is scalar it is treated as a 1-tuple.

    If A.ndim < d, A is promoted to be d-dimensional by prepending new axes.
    So a shape (3,) array is promoted to (1,3) for 2-D replication,
    or shape (1,1,3) for 3-D replication.
    If this is not the desired behavior, promote A to d-dimensions manually
    before calling this function.

    If d < A.ndim, tup is promoted to A.ndim by pre-pending 1's to it.  Thus
    for an A.shape of (2,3,4,5), a tup of (2,2) is treated as (1,1,2,2)


    Examples:
    >>> a = array([0,1,2])
    >>> tile(a,2)
    array([0, 1, 2, 0, 1, 2])
    >>> tile(a,(1,2))
    array([[0, 1, 2, 0, 1, 2]])
    >>> tile(a,(2,2))
    array([[0, 1, 2, 0, 1, 2],
           [0, 1, 2, 0, 1, 2]])
    >>> tile(a,(2,1,2))
    array([[[0, 1, 2, 0, 1, 2]],
    <BLANKLINE>
           [[0, 1, 2, 0, 1, 2]]])

    See Also:
       repeat
    """
    try:
        tup = tuple(reps)
    except TypeError:
        tup = (reps,)
    d = len(tup)
    c = _nx.array(A,copy=False,subok=True,ndmin=d)
    shape = list(c.shape)
    n = max(c.size,1)
    if (d < c.ndim):
        tup = (1,)*(c.ndim-d) + tup
    for i, nrep in enumerate(tup):
        if nrep!=1:
            c = c.reshape(-1,n).repeat(nrep,0)
        dim_in = shape[i]
        dim_out = dim_in*nrep
        shape[i] = dim_out
        n /= max(dim_in,1)
    return c.reshape(shape)