"""numarray: The big enchilada numeric module

"""
import sys as _sys
import types, math, os.path
import operator as _operator
import copy as _copy
import warnings as _warnings
from math import pi, e

import memory
import generic as _gen
import _bytes
import _numarray
import _ufunc
import _sort
import numerictypes as _nt

_PROTOTYPE = 0  # Set to 1 to switch to Python prototype code.
                # Set to 0 to inherit C code from C basetype.

# rename built-in function type so not to conflict with keyword
_type = type

MAX_LINE_WIDTH = None
PRECISION = None
SUPPRESS_SMALL = None

PyINT_TYPES = {
    bool: 1,
    int: 1,
    long: 1,
    }

PyREAL_TYPES = {
    bool: 1,
    int: 1,
    long: 1,
    float: 1,
    }

# Python numeric types with values indicating level in type hierarchy
PyNUMERIC_TYPES = {
    bool: 0,
    int: 1,
    long: 2,
    float: 3,
    complex: 4
    }

# Mapping back from level to type
PyLevel2Type = {}
for key, value in PyNUMERIC_TYPES.items():
    PyLevel2Type[value] = key
del key, value

# Mapping from Python to Numeric types
Py2NumType = {
    bool:  _nt.Bool,
    int:   _nt.Long,
    long:  _nt.Int64,
    float: _nt.Float,
    complex: _nt.Complex
    }

def array2list(arr):
    return arr.tolist()

# array factory functions

def _all_arrays(args):
    for x in args:
        if not isinstance(x, NumArray):
            return 0
    return len(args) > 0

def _maxtype(args):
    """Find the maximum scalar numeric type in the arguments.
    
    An exception is raised if the types are not python numeric types.
    """
    if not len(args):
        return None
    elif isinstance(args, NumArray):
        return args.type()    
    elif _all_arrays(args):
        temp = args[0].type()
        for x in args[1:]:
            if temp < x.type():
                temp = x.type()
        if isinstance(temp, _nt.BooleanType):
            return bool
        elif isinstance(temp, _nt.IntegralType):
            return int
        elif isinstance(temp, _nt.FloatingType):
            return float
        elif isinstance(temp, _nt.ComplexType):
            return complex
    else:
        return PyLevel2Type[_numarray._maxtype(args)]

def _storePyValueInBuffer(buffer, Ctype, index, value):
    """Store a python value in a buffer, index is in element units, not bytes"""
    # Do not use for complex scalars!
    Ctype._conv.fromPyValue(value, buffer._data,
                            index*Ctype.bytes, Ctype.bytes, 0)

def _storePyValueListInBuffer(buffer, Ctype, valuelist):
    # Do not use for complex values!
    for i in xrange(len(valuelist)):
        _storePyValueInBuffer(buffer, Ctype, i, valuelist[i])

def _fillarray(size, start, delta, type=None):
    ptype = _maxtype((start, delta))
    if ptype == long:
        ptype = _nt.Int64
    elif PyINT_TYPES.has_key(ptype):
        ptype = _nt.Long
    elif PyREAL_TYPES.has_key(ptype):
        ptype = _nt.Float
    else:
        ptype = _nt.Complex
    if type:
        outtype = _nt.getType(type)
        if (isinstance(ptype, _nt.ComplexType)
            and not isinstance( outtype, _nt.ComplexType)):
            raise TypeError("outtype must be a complex type")
    else:
        outtype = ptype
        
    if outtype > ptype: # Hack for Int64/UInt64 on 32-bit platforms.
        ptype = outtype
        
    if isinstance(outtype, _nt.ComplexType):
        # Not memory efficient at the moment
        real = _fillarray(size, complex(start).real, complex(delta).real,
                type = _realtype(ptype))
        image = _fillarray(size, complex(start).imag, complex(delta).imag,
                type = _realtype(ptype))
        outarr = NumArray((size,), outtype, real=real, imag=image)
    else:
        # save parameters in a buffer
        parbuffer = ufunc._bufferPool.getBuffer()
        _storePyValueListInBuffer(parbuffer, ptype, [start, delta])
        cfunction = _sort.functionDict[repr((ptype.name, 'fillarray'))]
        outarr = NumArray((size,), outtype)
        if ptype == outtype:
            # no conversion necessary, simple case
            _ufunc.CheckFPErrors()
            cfunction(size, 1, 1, ((outarr._data, 0), (parbuffer._data, 0)))
            errorstatus = _ufunc.CheckFPErrors()
            if errorstatus:
                ufunc.handleError(errorstatus, " in fillarray")
        else:
            # use buffer loop
            convbuffer = ufunc._bufferPool.getBuffer()
            convfunction = ptype._conv.astype[outtype.name]
            bsize = len(convbuffer._data)/ptype.bytes
            iters, lastbsize = divmod(size, bsize)
            _ufunc.CheckFPErrors()

            outoff = 0
            for i in xrange(iters + (lastbsize>0)):
                if i == iters:
                    bsize = lastbsize
                cfunction(bsize, 1, 1,
                          ((convbuffer._data, 0), (parbuffer._data, 0)))
                convfunction(bsize, 1, 1, 
                             ((convbuffer._data, 0), (outarr._data, outoff)))
                outoff += bsize*outtype.bytes
                start += delta * bsize
                _storePyValueListInBuffer(parbuffer, ptype, [start, delta])
            errorstatus = _ufunc.CheckFPErrors()
            if errorstatus:
                ufunc.handleError(errorstatus, " in fillarray")
    return outarr

def _frontseqshape(seq):
    """Find the length of all the first elements, return as a list"""
    if not len(seq):
        return (0,)
    if isinstance(seq, types.StringType):
        return (len(seq),)
    try:
        shape = []
        while 1:
            shape.append(len(seq))
            try:
                seq = seq[0]
                if isinstance(seq, types.StringType):
                    return shape
            except IndexError:
                return shape
    except TypeError:
        return shape
    except ValueError:
        if isinstance(seq, NumArray) and seq.rank == 0:
            return shape

def fromlist(seq, type=None, shape=None, check_overflow=0, typecode=None):
    """fromlist creates a NumArray from  the sequence 'seq' which must be
    a list or tuple of python numeric types.  If type is specified, it
    is as the type of  the resulting NumArray.  If shape is specified,
    it becomes the  shape of the result and must  have the same number
    of elements as seq.
    """
    type = _typeFromTypeAndTypecode(type, typecode)
    if not len(seq) and type is None:
        type = _nt.Long

    if _all_arrays(seq):
        a = _gen.concatenate(seq)
        if shape is not None:
            a.shape = shape
        else:
            a.shape = (len(seq),a._shape[0]/len(seq)) + a._shape[1:]
        if type is not None and a.type() != type:
            a = a.astype(type)
        return a
        
    if type is None:
        highest_type = _maxtype(seq)
        
    tshape = _frontseqshape(seq)
    if shape is not None and _gen.product(shape) != _gen.product(tshape):
        raise ValueError("shape incompatible with sequence")
    ndim = len(tshape)
    if ndim <= 0:
        raise TypeError("Argument must be a sequence")
    if type is None:
        type = Py2NumType.get(highest_type)
    if type is None:
        raise TypeError("Cannot create array of type %s" % highest_type.__name__)
    tshape = tuple(tshape)
    arr = NumArray(shape=tshape, type=type)
    arr._check_overflow = check_overflow
    arr.fromlist(seq)
    # _numarray._setarray(arr, seq)
    if shape is not None:
        arr.setshape(shape)
    arr._check_overflow = 0
    return arr

def _typeFromTypeAndTypecode(type, typecode):
    """returns a type object from a type or typecode specifier (keyword)
    or returns the type() of any sequence which is an NDArray.
    """
    if type is not None and typecode is not None and type != typecode:
        raise ValueError("Can't define both 'type' and 'typecode' for an array.")
    elif type is not None:              # Still might be a string or typecode
        return _nt.getType(type)
    elif typecode is not None:
        return _nt.getType(typecode)
    else:
        return None
    
def getTypeObject(sequence, type, typecode):
    """getTypeObject computes the typeObject for 'sequence' if both 'type' and
    'typecode' are unspecified.  Otherwise,  it returns the typeObject specified by
    'type' or 'typecode'.
    """
    rtype = _typeFromTypeAndTypecode(type, typecode)
    if rtype is not None:
        return rtype
    elif isinstance(sequence, _gen.NDArray):  # handle array([])
        return sequence.type()
    elif hasattr(sequence, "typecode"): # for Numeric/MA
        return sequence.typecode()
    elif (isinstance(sequence, (types.ListType, types.TupleType)) and
          len(sequence) == 0):
        return _nt.Long
    else:
        if isinstance(sequence, (types.IntType, types.LongType,
                                 types.FloatType, types.ComplexType)):
            sequence = [sequence]
        try:
            return Py2NumType[ _maxtype(sequence) ]
        except KeyError:
            raise TypeError("Can't determine a reasonable type from sequence")

def array(sequence=None, typecode=None, copy=1, savespace=0,
          type=None, shape=None):
    """array() constructs a NumArray by calling NumArray, one of its
    factory functions (fromstring, fromfile, fromlist), or by making a
    copy of an existing array.  If copy=0, array() will create a new
    array only if

    sequence             specifies the contents or storage for the array

    type and typecode    are interchangeable and define the array type
                         using either a type object or string.

    copy                  when sequence is an array, copy determines
                          whether a copy or the original is returned.

    savespace            is ignored by numarray.
    
    shape                defines the shape of the array object and is
                         necessary when not implied by the sequence
                         parameter.
    """
    if sequence is None and shape is None:
        return None

    if isinstance(shape, types.IntType):
        shape = (shape,)
        
    if sequence is None and (shape is None or type is None):
        raise ValueError("Must define shape and type if sequence is None")
    
    if isinstance(sequence, _gen.NDArray):
        if type is None and typecode is None:
            if copy:
                a = sequence.copy()
            else:
                a = sequence
        else:
            type = getTypeObject(sequence, type, typecode)
            if not copy and type is sequence.type():
                a = sequence
            else:
                a = sequence.astype(type) # make a re-typed copy
        if shape is not None:
            a.setshape(shape)
        return a

    type = getTypeObject(sequence, type, typecode) 
        
    if sequence is None or _gen.SuitableBuffer(sequence):
        return NumArray(buffer=sequence, shape=shape, type=type)
    elif isinstance(sequence, types.StringType):
        return fromstring(sequence, shape=shape, type=type)
    elif isinstance(sequence, (types.ListType, types.TupleType)):
        return fromlist(sequence, type, shape)
    elif PyNUMERIC_TYPES.has_key(_type(sequence)):
        if shape and shape != ():
            raise ValueError("Only shape () is valid for a rank 0 array.")
        return fromlist([sequence], shape=(), type=type or
                        Py2NumType[_type(sequence)])
    elif isinstance(sequence, types.FileType):
        return fromfile(sequence, type=type, shape=shape)
    else:
        try:
            return sequence.__array__(type)
        except AttributeError:
            try:
                sequence[0]
            except:
                raise ValueError("Unknown input type")
            else:
                return fromlist(sequence, type, shape)

def asarray(seq, type=None, typecode=None):
    """converts scalars, lists and tuples to numarray if possible.

    passes NumArrays thru making copies only to convert types.
    """
    if isinstance(seq, _gen.NDArray) and type is None and typecode is None:
        return seq
    return array(seq, type=type, typecode=typecode, copy=0)

inputarray = asarray   # Obsolete synonym

def fromstring(datastring, type=None, shape=None, typecode=None):
    """Create an array from binary data contained in a string (by copying)"""
    type = _typeFromTypeAndTypecode(type, typecode)
    if not shape:
        size, rem = divmod(len(datastring), type.bytes)
        if rem:
            raise ValueError("Type size inconsistent with string length")
        else:
            shape = (size,) # default to a 1-d array
    elif _type(shape) is types.IntType:
        shape = (shape,)
    if len(datastring) != (_gen.product(shape)*type.bytes):
        raise ValueError("Specified shape and type inconsistent with string length")
    arr = NumArray(shape=shape, type=type)
    strbuff = buffer(datastring)
    nelements = arr.nelements()
    # Currently uses only the byte-by-byte copy, should be optimized to use
    # larger element copies when possible.
    cfunc = _bytes.functionDict["copyNbytes"]
    cfunc((nelements,), strbuff, 0, (type.bytes,),
          arr._data, 0, (type.bytes,), type.bytes)
    if arr._type == _nt.Bool:
        arr = ufunc.not_equal(arr, 0)
    return arr

def fromfile(file, type, shape=None):
    """Create an array from binary file data

    If file is a string then that file is opened, else it is assumed
    to be a file object. No options at the moment, all file positioning
    must be done prior to this function call with a file object

    """
    type = _nt.getType(type)
    name =  0
    if _type(file) == _type(""):
        name = 1
        file = open(file, 'rb')
    size = os.path.getsize(file.name) - file.tell()
    if not shape:
        nelements, rem = divmod(size, type.bytes)
        if rem:
            raise ValueError(
                "Type size inconsistent with shape or remaining bytes in file")
        shape = (nelements,)
    elif _type(shape) is types.IntType:
        shape=(shape,)
    nbytes = _gen.product(shape)*type.bytes
    if nbytes > size:
        raise ValueError(
                "Not enough bytes left in file for specified shape and type")
    # create the array
    arr = NumArray(shape=shape, type=type)
    # Use of undocumented file method! XXX
    nbytesread = file.readinto(arr._data)
    if nbytesread != nbytes:
        raise IOError("Didn't read as many bytes as expected")
    if name:
        file.close()
    if arr._type == _nt.Bool:
        arr = ufunc.not_equal(arr, 0)
    return arr

class UsesOpPriority(object):
    """Classes can subclass from UsesOpPriority to signal to numarray
    that perhaps the class' r-operator hook (e.g. __radd__) should be
    given preference over NumArray's l-operator hook (e.g. __add__).
    This would be done so that when different object types are used in
    an operation (e.g. NumArray + MaskedArray) the type of the result
    is well defined and independent of the order of the operands
    (e.g. MaskedArray).

    Before altering the "normal" behavior of an operator, this scheme
    (implemented in the operator hook functions of NumArray) first
    checks to see if the other operand subclasses UsesOpPriority.  If
    it does, the op priorities of both operands are compared, and an
    appropriate hook function from the one with the highest priority
    is called.

    Thus, a subclass of NumArray which wants to ensure that its type
    dominates in mixed type operations should define a class level
    op_priority > 0.  If several subclasses wind up doing this,
    op_priority will determine how they relate to one another as well.
    """
    op_priority = 0.0

class NumArray(_numarray._numarray, _gen.NDArray, UsesOpPriority):
    """Fundamental Numeric Array
    
    type       The type of each data element, e.g. Int32  
    byteorder  The actual ordering of bytes in buffer: "big" or "little".
    """
    if _PROTOTYPE:
        def __init__(self, shape=None, type=None, buffer=None,
                     byteoffset=0, bytestride=None, byteorder=_sys.byteorder,
                     aligned=1, real=None, imag=None):
            type = _nt.getType(type)
            itemsize = type.bytes
            _gen.NDArray.__init__(self, shape, itemsize, buffer,
                                  byteoffset, bytestride)
            self._type = type
            if byteorder in ["little", "big"]:
                self._byteorder = byteorder
            else:
                raise ValueError("byteorder must be 'little' or 'big'")
            if real is not None:
                self.real = real
            if imag is not None:
                self.imag = imag

        def _copyFrom(self, arr):
            """Copy elements from another array.

            This overrides the _generic NDArray version
            """
            # Test for simple case first
            if isinstance(arr, NumArray):
                if (arr.nelements() == 0 and self.nelements() == 0):
                    return
                if (arr._type == self._type and
                    self._shape == arr._shape and
                    arr._byteorder == self._byteorder and
                    _gen.product(arr._strides) != 0 and
                    arr.isaligned() and self.isaligned()):
                    name = 'copy'+`self._itemsize`+'bytes'
                    cfunc = ( _bytes.functionDict.get(name) or
                              _bytes.functionDict["copyNbytes"])
                    cfunc(self._shape, arr._data,  arr._byteoffset,  arr._strides,
                          self._data, self._byteoffset, self._strides,
                          self._itemsize)
                    return
            elif PyNUMERIC_TYPES.has_key(_type(arr)):
                # Convert scalar to a one element array for broadcasting
                arr = array([arr])
            else:
                raise TypeError('argument is not array or number')
            barr = self._broadcast(arr)
            ufunc._copyFromAndConvert(barr, self)

        def view(self):
            """Returns a new array object which refers to the same data as the
            original array.  The new array object can be manipulated (reshaped,
            restrided, new attributes, etc.) without affecting the original array.
            Modifications of the array data *do* affect the original array.
            """
            v = _gen.NDArray.view(self)
            v._type = self._type
            v._byteorder = self._byteorder
            return v

    def __del__(self):
        if self._shadows != None:
            self._shadows._copyFrom(self)
            self._shadows = None
 
    def __getstate__(self):
        """returns state of NumArray for pickling."""
        # assert not hasattr(self, "_shadows") # Not a good idea for pickling.
        state = _gen.NDArray.__getstate__(self)
        state["_byteorder"] = self._byteorder
        state["_type"]      = self._type.name
        return state

    def __setstate__(self, state):
        """sets state of NumArray after unpickling."""
        _gen.NDArray.__setstate__(self, state)
        self._byteorder = state["_byteorder"]
        self._type      = _nt.getType(state["_type"])

    def _put(self, indices, values, **keywds):
        ufunc._put(self, indices, values, **keywds)

    def _take(self, indices, **keywds):
        return ufunc._take(self, indices, **keywds)

    def getreal(self):
        if isinstance(self._type, _nt.ComplexType):
            t = _realtype(self._type)
            arr = NumArray(self._shape, t, buffer=self._data,
                           byteoffset=self._byteoffset,
                           bytestride=self._bytestride,
                           byteorder=self._byteorder)
            arr._strides = self._strides[:]
            return arr
        elif isinstance(self._type, _nt.FloatingType):
            return self
        else:
            return self.astype(_nt.Float64)            

    def setreal(self, value):
        if isinstance(self._type, (_nt.ComplexType, _nt.FloatingType)):
            self.getreal()[:] = value
        else:
            raise TypeError("Can't setreal() on a non-floating-point array")
                    
    real = property(getreal, setreal,
                    doc="real component of a non-complex numarray")

    def getimag(self):
        if isinstance(self._type, _nt.ComplexType):
            t = _realtype(self._type)
            arr = NumArray(self._shape, t, buffer=self._data,
                           byteoffset=self._byteoffset+t.bytes,
                           bytestride=self._bytestride,
                           byteorder=self._byteorder)
            arr._strides = self._strides[:]
            return arr
        else:
            zeros = self.new(_nt.Float64)
            zeros[:] = 0.0
            return zeros

    def setimag(self, value):
        if isinstance(self._type, _nt.ComplexType):
            self.getimag()[:] = value
        else:
            raise TypeError("Can't set imaginary component of a non-comlex type")

    imag = property(getimag, setimag,
                        doc="imaginary component of complex array")

    real = property(getreal, setreal,
                        doc="real component of complex array")

    setimaginary = setimag
    getimaginary = getimag
    imaginary = imag

    def conjugate(self):
        """Returns the conjugate a - bj of complex array a + bj."""
        a = self.copy()
        ufunc.minus(a.getimag(), a.getimag())
        return a

    def togglebyteorder(self):
        """reverses the state of the _byteorder attribute:  little <-> big."""
        self._byteorder = {"little":"big","big":"little"}[self._byteorder]

    def byteswap(self):
        """Byteswap data in place, leaving the _byteorder attribute untouched.
        """
        if self._itemsize == 1:
            return
        if isinstance(self._type, _nt.ComplexType):
            fname = "byteswap" + self._type.name
        else:
            fname = "byteswap"+str(self._itemsize)+"bytes"
            
        cfunc = _bytes.functionDict[fname]
        cfunc(self._shape, 
              self._data, self._byteoffset, self._strides,
              self._data, self._byteoffset, self._strides)
        
    _byteswap = byteswap  # alias to keep old code working.

    def byteswapped(self):
        """returns a byteswapped copy of self, with adjusted _byteorder."""
        b = self.copy()
        b._byteswap()
        return b
        
    def info(self):
        """info() prints out the key attributes of a numarray."""
        _gen.NDArray.info(self)
        print "byteorder:", self._byteorder
        print "byteswap:", self.isbyteswapped()
        print "type:", repr(self._type)

    def astype(self, type=None):
        """Return a copy of the array converted to the given type"""
        if type is None:
            type = self._type
        type = _nt.getType(type)
        if type == self._type:
            # always return a copy even if type is same
            return self.copy()
        if type._conv:
            retarr = self.__class__(buffer=None, shape=self._shape, type=type)
            if retarr.nelements() == 0:
                return retarr
            if self.is_c_array():
                _ufunc.CheckFPErrors()
                cfunc = self._type._conv.astype[type.name]
                cfunc(self.nelements(), 1, 1,
                      ((self._data, self._byteoffset), (retarr._data, 0)))
                errorstatus = _ufunc.CheckFPErrors()
                if errorstatus:
                    ufunc.handleError(errorstatus, " during type conversion")
            else:
                ufunc._copyFromAndConvert(self, retarr)
        elif type.fromtype:
            retarr = type.fromtype(self)
        else:
            raise TypeError("Don't know how to convert from %s to %s" %
                            (self._type.name, type.name))
        return retarr

    def __tonumtype__(self, type):
        """__tonumtype__ supports C-API NA_setFromPythonScalar permitting a rank-0
        array to be used in lieu of a numerical scalar value."""
        if self.rank != 0:
            raise ValueError("Can't use non-rank-0 array as a scalar.")
        if type is not self.type():
            s = self.astype(type)
        else:
            s = self
        return s[()]

    def is_c_array(self):
        """returns 1 iff an array is c-contiguous, aligned, and not
        byteswapped."""
        return (self.iscontiguous() and self.isaligned() and not
                self.isbyteswapped()) != 0

    def is_f_array(self):
        """returns 1 iff an array is fortan-contiguous, aligned, and not
        byteswapped."""
        return (self.is_fortran_contiguous() and self.isaligned() and not
                self.isbyteswapped())

    def new(self, type=None):
        """Return a new array of given type with same shape as this array

        Note this only creates the array; it does not copy the data.
        """
        if type is None:
            type = self._type
        return self.__class__(shape=self._shape, type=type)
        
    def type(self):
        """Return the type object for the array"""
        return self._type

    def copy(self):
        """Returns a native byte order copy of the array."""
        c = _gen.NDArray.copy(self)
        c._byteorder = self._byteorder
        c._type = self._type
        if self.isbyteswapped():
            c.byteswap()
            c.togglebyteorder()
        return c

    def _stdtype(self, t):
        return t in [_nt.Long, _nt.Float, _nt.Complex]

    def __str__(self):
        return array_str(self)

    def __repr__(self):
        return array_repr(self)

    def __int__(self):
        if len(self._shape) == 0:
            return int(self[()])
        else:
            raise TypeError, "Only rank-0 numarray can be cast to integers."

    def __float__(self):
        if len(self._shape) == 0:
            return float(self[()])
        else:
            raise TypeError, "Only rank-0 numarray can be cast to floats."

    def __complex__(self):
        if len(self._shape) == 0:
            return complex(self[()])
        else:
            raise TypeError, "Only rank-0 numarray can be cast to complex."

    def __abs__(self): return ufunc.abs(self)
    def __neg__(self): return ufunc.minus(self)
    def __invert__(self): return ufunc.bitwise_not(self)
    def __pos__(self): return self
    
    def __add__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__radd__(self)
        else:
            return ufunc.add(self, operand)
    
    def __radd__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__add__(self)
        else:
            r = ufunc.add(operand, self)
            if isinstance(r, NumArray):
                r.__class__ = self.__class__
            return r

    def __sub__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__rsub__(self)
        else:
            return ufunc.subtract(self, operand)
    
    def __rsub__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__sub__(self)
        else:
            r = ufunc.subtract(operand, self)
            if isinstance(r, NumArray):
                r.__class__ = self.__class__
            return r

    def __mul__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__rmul__(self)
        else:
            return ufunc.multiply(self, operand)
    
    def __rmul__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__mul__(self)
        else:
            r = ufunc.multiply(operand, self)
            if isinstance(r, NumArray):
                r.__class__ = self.__class__
            return r

    def __div__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__rdiv__(self)
        else:
            return ufunc.divide(self, operand)
    
    def __truediv__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__rtruediv__(self)
        else:
            return ufunc.true_divide(self, operand)
    
    def __floordiv__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__rfloordiv__(self)
        else:
            return ufunc.floor_divide(self, operand)
    
    def __rdiv__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__div__(self)
        else:
            r = ufunc.divide(operand, self)
            if isinstance(r, NumArray):
                r.__class__ = self.__class__
            return r

    def __rtruediv__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__truediv__(self)
        else:
            r = ufunc.true_divide(operand, self)
            if isinstance(r, NumArray):
                r.__class__ = self.__class__
            return r

    def __rfloordiv__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__floordiv__(self)
        else:
            r = ufunc.floor_divide(operand, self)
            if isinstance(r, NumArray):
                r.__class__ = self.__class__
            return r

    def __mod__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__rmod__(self)
        else:
            return ufunc.remainder(self, operand)
    
    def __rmod__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__mod__(self)
        else:
            r = ufunc.remainder(operand, self)
            if isinstance(r, NumArray):
                r.__class__ = self.__class__
            return r

    def __divmod__(self,operand):
        """returns the tuple (self/operand, self%operand)."""
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__rdivmod__(self)
        else:
            return ufunc.divide_remainder(self, operand)

    def __rdivmod__(self,operand):
        """returns the tuple (operand/self, operand%self)."""
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__divmod__(self)
        else:
            r = ufunc.divide_remainder(operand, self)
            if isinstance(r, NumArray):
                r.__class__ = self.__class__
            return r

    def __pow__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__rpow__(self)
        else:
            return ufunc.power(self, operand)
    
    def __rpow__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__pow__(self)
        else:
            r = ufunc.power(operand, self)
            if isinstance(r, NumArray):
                r.__class__ = self.__class__
            return r

    def __and__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__rand__(self)
        else:
            return ufunc.bitwise_and(self, operand)
    
    def __rand__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__and__(self)
        else:
            r = ufunc.bitwise_and(operand, self)
            if isinstance(r, NumArray):
                r.__class__ = self.__class__
            return r

    def __or__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__ror__(self)
        else:
            return ufunc.bitwise_or(self, operand)
    
    def __ror__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__or__(self)
        else:
            r = ufunc.bitwise_or(operand, self)
            if isinstance(r, NumArray):
                r.__class__ = self.__class__
            return r

    def __xor__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__rxor__(self)
        else:
            return ufunc.bitwise_xor(self, operand)
    
    def __rxor__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__xor__(self)
        else:
            r = ufunc.bitwise_xor(operand, self)
            if isinstance(r, NumArray):
                r.__class__ = self.__class__
            return r

    def __rshift__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__rrshift__(self)
        else:
            return ufunc.rshift(self, operand)
    
    def __rrshift__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__rshift__(self)
        else:
            r = ufunc.rshift(operand, self)
            if isinstance(r, NumArray):
                r.__class__ = self.__class__
            return r

    def __lshift__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__rlshift__(self)
        else:
            return ufunc.lshift(self, operand)
    
    def __rlshift__(self, operand):
        if (isinstance(operand, UsesOpPriority) and
            self.op_priority < operand.op_priority):
            return operand.__lshift__(self)
        else:
            r = ufunc.lshift(operand, self)
            if isinstance(r, NumArray):
                r.__class__ = self.__class__
            return r

    # augmented assignment operators

    def __iadd__(self, operand):
        ufunc.add(self, operand, self)
        return self
    
    def __isub__(self, operand):
        ufunc.subtract(self, operand, self)
        return self
    
    def __imul__(self, operand):
        ufunc.multiply(self, operand, self)
        return self
    
    def __idiv__(self, operand):
        ufunc.divide(self, operand, self)
        return self
    
    def __ifloordiv__(self, operand):
        ufunc.floor_divide(self, operand, self)
        return self
    
    def __itruediv__(self, operand):
        ufunc.true_divide(self, operand, self)
        return self
    
    def __imod__(self, operand):
        ufunc.remainder(self, operand, self)
        return self
    
    def __ipow__(self, operand):
        ufunc.power(self, operand, self)
        return self
    
    def __iand__(self, operand):
        ufunc.bitwise_and(self, operand, self)
        return self
    
    def __ior__(self, operand):
        ufunc.bitwise_or(self, operand, self)
        return self
    
    def __ixor__(self, operand):
        ufunc.bitwise_xor(self, operand, self)
        return self
    
    def __irshift__(self, operand):
        ufunc.rshift(self, operand, self)
        return self
    
    def __ilshift__(self, operand):
        ufunc.lshift(self, operand, self)
        return self

    # rich comparisons (only works in Python 2.1 and later)

    def __lt__(self, operand):
        if isinstance(self._type, _nt.ComplexType):
            raise TypeError("Complex numarray don't support < comparison")
        else:
            return ufunc.less(self, operand)
        
    def __gt__(self, operand):
        if isinstance(self._type, _nt.ComplexType):
            raise TypeError("Complex numarray don't support > comparison")
        else:
            return ufunc.greater(self, operand)

    def __le__(self, operand):
        if isinstance(self._type, _nt.ComplexType):
            raise TypeError("Complex numarray don't support <= comparison")
        else:
            return ufunc.less_equal(self, operand)

    def __ge__(self, operand):
        if isinstance(self._type, _nt.ComplexType):
            raise TypeError("Complex numarray don't support >= comparison")
        else:
            return ufunc.greater_equal(self, operand)

    def __eq__(self, operand):
        if operand is None:
            return 0
        else:
            return ufunc.equal(self, operand)
    def __ne__(self, operand):
        if operand is None:
            return 1
        else:
            return ufunc.not_equal(self, operand)

    def sort(self, axis=-1):
        """sorts an array in-place along the specified axis, returning None."""
        if axis==-1:
            ufunc._sortN(self)
        else:
            self.swapaxes(axis,-1)
            ufunc._sortN(self)
            self.swapaxes(axis,-1)

    def _argsort(self, axis=-1):
        if axis==-1:
            ashape = self.getshape()
            w = array(shape=ashape, type=_nt.Long)
            w[...,:] = arange(ashape[-1], type=_nt.Long)
            ufunc._argsortN(self,w)
            return w
        else:
            self.swapaxes(axis,-1)
            w = self._argsort()
            w.swapaxes(axis, -1)
            return w

    def argsort(self, axis=-1):
        """returns the indices of 'self' which if taken from self would return
        a copy of 'self' sorted along the specified 'axis'"""
        return self.copy()._argsort(axis)

    def argmax(self, axis=-1):
        """returns the index(es) of the greatest element(s) of 'self' along the
        specified 'axis'."""
        import numeric
        return numeric.argmax(self, axis)

    def argmin(self, axis=-1):
        """returns the index(es) of the least element(s) of 'self' along the
        specified 'axis'."""
        import numeric
        return numeric.argmin(self, axis)

    def diagonal(self, *args, **keywords):
        """returns the diagonal elements of the array."""
        return diagonal(self, *args, **keywords)

    def trace(self, *args, **keywords):
        """returns the sum of the diagonal elements of the array."""
        return trace(self, *args, **keywords)

    def typecode(self):
        """returns the Numeric typecode of the array."""
        return _nt.typecode[self._type]

    def min(self):
        """Returns the minimum element in the array."""
        return ufunc.minimum.reduce(ufunc.minimum.areduce(self).flat)

    def max(self):
        """Returns the maximum element in the array."""
        return ufunc.maximum.reduce(ufunc.maximum.areduce(self).flat)

    def sum(self, type=None):
        """Returns the sum of all elements in the array."""
        if type is None:
            type = _nt.MaximumType(self._type)
        return ufunc.add.reduce(ufunc.add.areduce(self, type=type).flat, type=type)

    def mean(self):
        """Returns the average of all elements in the array."""
        return self.sum()/(self.nelements()*1.0)

    def stddev(self):
        """Returns the standard deviation of the array."""
        m = self.mean()
        N = self.nelements()
        return math.sqrt(((self - m)**2).sum()/(N-1))

    def spacesaver(self):
        """Always False.  Supported for Numeric compatibility."""
        return 0

class ComplexArray(NumArray):   # Deprecated
    pass

def Complex32_fromtype(arr):
    """Used for converting other types to Complex32.

    This is used to set an fromtype attribute of the ComplexType objects"""
    rarr = arr.astype(Float32)
    retarr = ComplexArray(arr._shape, _nt.Complex32)
    retarr.getreal()[:] = rarr
    retarr.getimag()[:] = 0.
    return retarr

def Complex64_fromtype(arr):
    """Used for converting other types to Complex64.

    This is used to set an fromtype attribute of the ComplexType objects"""
    rarr = arr.astype(Float64)
    retarr = ComplexArray(arr._shape, _nt.Complex64)
    retarr.getreal()[:] = rarr
    retarr.getimag()[:] = 0.
    return retarr

# Check whether byte order is big endian or little endian.

from sys import byteorder
isBigEndian = (byteorder == "big")
del byteorder

# Add fromtype function to Complex types

_nt.Complex32.fromtype = Complex32_fromtype
_nt.Complex64.fromtype = Complex64_fromtype

# Return type of complex type components

def _isComplexType(type):  # Deprecated
    return type in [_nt.Complex32, _nt.Complex64]

def _realtype(complextype):
    if complextype == _nt.Complex32:
        return _nt.Float32
    else:
        return _nt.Float64

def conjugate(a):
    """conjugate(a) returns the complex conjugate of 'a'"""
    a = asarray(a)
    if not isinstance(a._type, _nt.ComplexType):
        a = a.astype(_nt.Complex64)
    return a.conjugate()

def zeros(shape, type=None, typecode=None):
    """Constructs a zero filled array of the specified shape and type."""
    type = _typeFromTypeAndTypecode(type, typecode)
    if type is None:
        type = _nt.Long
    retarr = NumArray(shape=shape, type=type)
    retarr._data.clear()
    return retarr

def ones(shape, type=None, typecode=None):
    """Constructs an array of the specified shape and type filled with ones."""
    type = _typeFromTypeAndTypecode(type, typecode)
    shape = _gen.getShape(shape)
    retarr = _fillarray(_gen.product(shape), 1, 0, type)
    retarr.setshape(shape)
    return retarr

def arange(a1, a2=None, stride=1, type=None, shape=None, typecode=None):
    """Returns an array of numbers in sequence over the specified range."""
    # Return empty range of correct type for single negative paramter.
    type = _typeFromTypeAndTypecode(type, typecode)
    if not isinstance(a1, types.ComplexType) and a1 < 0 and a2 is None:
        t = __builtins__["type"](a1)
        return zeros(shape=(0,), type=Py2NumType[t])
    if a2 == None:
        start = 0 + (0*a1) # to make it same type as stop
        stop  = a1
    else:
        start = a1 +(0*a2)
        stop  = a2
    delta = (stop-start)/stride ## xxx int divide issue
    if _type(delta) == types.ComplexType:
        # xxx What make sense here?
        size = int(math.ceil(delta.real))
    else:
        size = int(math.ceil((stop-start)/float(stride)))
    if size < 0:
        size = 0
    r = _fillarray(size, start, stride, type)
    if shape is not None:
        r.setshape(shape)
    return r

arrayrange = arange  # alias arange as arrayrange.

def identity(n, type=None, typecode=None):
    """Returns an array resembling and identity matrix."""
    type = _typeFromTypeAndTypecode(type, typecode)
    a = zeros(shape=(n,n), type=type)
    i = arange(n)
    a.put((i, i), 1, axis=(0,))
    return a

if _PROTOTYPE:
    def dot(array1, array2):
        """dot matrix-multiplies array1 by array2.
        """
        return ufunc.innerproduct(array1, _gen.swapaxes(array2, -1, -2))
else:
    from _numarray import dot

matrixmultiply = dot  # Deprecated in Numeric

def outerproduct(array1, array2):
    """outerproduct(array1, array2) computes the NxM outerproduct of N vector
    'array1' and M vector 'array2', where result[i,j] = array1[i]*array2[j].
    """
    array1=_gen.reshape(asarray(array1), (-1,1))  # ravel array1 into an Nx1
    array2=_gen.reshape(asarray(array2), (1,-1))  # ravel array2 into a 1xM
    return matrixmultiply(array1,array2)   # return NxM result

def allclose (array1, array2, rtol=1.e-5, atol=1.e-8): # From Numeric 20.3
    """allclose returns true if all components of array1 and array2 are
    equal subject to given tolerances.  The relative error rtol must be
    positive and << 1.0.  The absolute error atol comes into play for those
    elements of y that are very small or zero; it says how small x must be
    also.
    """
    x, y = asarray(array1), asarray(array2)
    d = ufunc.less(ufunc.abs(x-y), atol + rtol * ufunc.abs(y))
    return bool(alltrue(_gen.ravel(d)))

# JC's revised diagonal for supporting dimensions > 2 correctly.
def diagonal(a, offset= 0, axis1=0, axis2=1):
    """diagonal(a, offset=0, axis1=0, axis2=1) returns the given diagonals
    defined by the last two dimensions of the array.
    """
    a = inputarray(a)
    ###  do not swap axis1 and axis2, so the offset sign is meaningful
    ###if axis2 < axis1: axis1, axis2 = axis2, axis1
    ###if axis2 > 1:
    if 1: ###
        new_axes = range(len(a._shape))
        ###del new_axes[axis2]; del new_axes[axis1]
        try: ###
            new_axes.remove(axis1)  ###
            new_axes.remove(axis2)  ###
        except: ###
            raise ValueError, "axis1(=%d) and axis2(=%d) must be different and within range." % (axis1, axis2) ###
        ###new_axes [0:0] = [axis1, axis2]
        new_axes = new_axes + [axis1, axis2] ### insert at the end, not the beginning
        a = _gen.transpose(a, new_axes)
    s = a._shape
    rank = len(s) ###
    if rank == 2: ###
        n1 = s[0]
        n2 = s[1]
        n = n1 * n2
        s = (n,)
        a = _gen.reshape(a, s)
        if offset < 0:
            return _gen.take(a, range(- n2 * offset, min(n2, n1+offset) *
                                      (n2+1) - n2 * offset, n2+1), axis=0)
        else:
            return _gen.take(a, range(offset, min(n1, n2-offset) *
                                      (n2+1) + offset, n2+1), axis=0)
    else :
        my_diagonal = []
        for i in range(s[0]):
            my_diagonal.append(diagonal(a[i], offset, rank-3, rank-2)) ###
        return array(my_diagonal)

# From Numeric-21.0
def trace(array, offset=0, axis1=0, axis2=1):
    """trace returns the sum along diagonals (defined by the last
    two dimenions) of the array.
    """
    return ufunc.add.reduce(diagonal(array, offset, axis1, axis2))

def rank(array):
    """rank returns the number of dimensions in an array."""
    return len(shape(array))

def shape(array):
    """shape returns the shape tuple of an array."""
    try:
        return array.shape
    except AttributeError:
        return asarray(array).getshape()

def size(array, axis=None):
    """size  returns the number of elements in an array or
    along the specified axis."""
    array = asarray(array)
    if axis is None:
        return array.nelements()
    else:
        s = array.shape
        if axis < 0:
            axis += len(s)
        return s[axis]

def array_str(array, max_line_width=MAX_LINE_WIDTH, precision=PRECISION,
              suppress_small=SUPPRESS_SMALL):
    """Formats and array as a string."""
    array = asarray(array)
    return arrayprint.array2string(
        array, max_line_width, precision, suppress_small, ' ', "")

def array_repr(array, max_line_width=MAX_LINE_WIDTH, precision=PRECISION,
               suppress_small=SUPPRESS_SMALL):
    """Returns the repr string of an array."""
    array = asarray(array)
    lst = arrayprint.array2string(
        array, max_line_width, precision, suppress_small, ', ', "array(")
    typeless = array._stdtype(array._type)
    if array.__class__ is not NumArray:
        cName= array.__class__.__name__
    else:
        cName = "array"
    if typeless and not hasattr(array, "_explicit_type"):
        return cName + "(%s)" % lst
    else:
        return cName + "(%s, type=%s)" % (lst, array._type.name)

def around(array, digits=0, output=None):
    """rounds 'array'  to 'digits' of precision, storing the result in
    'output', or returning the result as new array if output=None"""
    array = asarray(array)
    scale = 10.0**digits
    if output is None:
        wout = array.copy()
    else:
        wout = output
    if digits != 0:
        wout *= scale  # bug in 2.2.1 and earlier causes fail as bad sequence op
    ufunc._round(wout, wout)
    if digits != 0:
        wout /= scale
    if output is None:
        return wout

def round(*args, **keys):
    _warnings.warn("round() is deprecated.  Switch to around().",
                   DeprecationWarning)
    return ufunc._round(*args, **keys)

def explicit_type(x):
    """explicit_type(x) returns a view of x which will always show it's type in it's repr.
    This is useful when the same test is run in two places, one where the type used *is* the
    default and hence not normally repr'ed, and one where the type used *is not* the default
    and so is displayed.
    """
    y = x.view()
    y._explicit_type = 1
    return y

ArrayType = NumArray  # Alias for backwards compatability with Numeric

def array_equiv(array1, array2):

    """array_equiv returns True if 'a' and 'b' are shape consistent
    (mutually broadcastable) and have all elements equal and False
    otherwise."""
    
    try:
        array1, array2 = asarray(array1), asarray(array2)
    except TypeError:
        return 0
    if not isinstance(array1, NumArray) or not isinstance(array2, NumArray):
        return 0
    try:
        array1, array2 = array1._dualbroadcast(array2)
    except ValueError:
        return 0
    return ufunc.logical_and.reduce(_gen.ravel(array1 == array2))

def array_equal(array1, array2):
    
    """array_equal returns True if array1 and array2 have identical shapes
    and all elements equal and False otherwise."""

    try:
        array1, array2 = asarray(array1), asarray(array2)
    except TypeError:
        return 0
    if not isinstance(array1, NumArray) or not isinstance(array2, NumArray):
        return 0
    if array1._shape != array2._shape:
        return 0
    return ufunc.logical_and.reduce(_gen.ravel(array1 == array2))


class _UBuffer(NumArray):
    """_UBuffer is used to hold a single "block" of ufunc data during
    the block-wise processing of all elements in an array.

    Subclassing the buffer object from numnumarray simplifies (and speeds!)
    their usage at the C level.  They are not intended to be used as
    public array objects, hence they are private!
    """

    def __init__(self, pybuffer):
        NumArray.__init__(self, (len(pybuffer),), _nt.Int8, pybuffer)
        self._strides    = None   # how it is distinguished from a real array

    def isbyteswapped(self):           return 0    
    def isaligned(self):               return 1
    def iscontiguous(self):            return 1
    def is_c_array(self):              return 1
    def _getByteOffset(self, shape):   return 0
    
    def __del__(self):
        """On deletion return the data to the buffer pool"""
        if self._data is not None:
            if ufunc is not None and ufunc._bufferPool is not None:
                ufunc._bufferPool.buffers.append(self._data)

    def __repr__(self):
        return "<_UBuffer of size:%d>" % self.shape[0]

import ufunc
import arrayprint

sum = ufunc.add.reduce
cumsum = ufunc.add.accumulate
product = ufunc.multiply.reduce
cumproduct = ufunc.multiply.accumulate
absolute = ufunc.abs  
negative = ufunc.minus
fmod = ufunc.remainder

def alltrue(array, axis=0):
    """For 1D arrays, returns True IFF all the elements of the array are nonzero.
    For higher dimensional arrays, returns a reduction.
    """
    array = asarray(array)
    if array.rank == 0:
        return array[()] == 1
    else:
        return ufunc.logical_and.reduce(array, axis=axis)
    
def sometrue(array, axis=0):
    """For 1D arrays, returns True IFF any one of the elements of the array is nonzero.
    For higher dimensional arrays, returns a reduction.
    """
    array = asarray(array)
    if array.rank == 0:
        return array[()] == 1
    else:
        return ufunc.logical_or.reduce(array, axis=axis)
        

NewArray = NumArray  # unnecessary following merger of real/complex Arrays 

def tensormultiply(array1, array2):
    """tensormultiply returns the product for any rank >=1 arrays, defined as:

    r_{xxx, yyy} = \sum_k array1_{xxx, k} array2_{k, yyyy}

    where xxx, yyy denote the rest of the a and b dimensions.
    """
    array1, array2 = asarray(array1), asarray(array2)
    if array1.shape[-1] != array2.shape[0]:
        raise ValueError, "Unmatched dimensions"
    shape = array1.shape[:-1] + array2.shape[1:]
    return _gen.reshape(dot(_gen.reshape(array1, (-1, array1.shape[-1])),
                            _gen.reshape(array2, (array2.shape[0], -1))), shape)

def kroneckerproduct(a,b):
	'''Computes a otimes b where otimes is the Kronecker product operator.
	
	Note: the Kronecker product is also known as the matrix direct product
	or tensor product.  It is defined as follows for 2D arrays a and b
	where shape(a)=(m,n) and shape(b)=(p,q):
	c = a otimes b  => cij = a[i,j]*b  where cij is the ij-th submatrix of c.
	So shape(c)=(m*p,n*q).

	>>> print kroneckerproduct([[1,2]],[[3],[4]])
	[[3 6]
	 [4 8]]
	>>> print kroneckerproduct([[1,2]],[[3,4]])
	[ [3 4 6 8]]
	>>> print kroneckerproduct([[1],[2]],[[3],[4]])
	[[3]
	 [4]
	 [6]
	 [8]]
	'''
	a, b = asarray(a), asarray(b)
	if not (len(shape(a))==2 and len(shape(b))==2):
        	raise ValueError, 'Input must be 2D arrays.'
	if not a.iscontiguous():
		a = _gen.reshape(a, a.shape)
	if not b.iscontiguous():
		b = _gen.reshape(b, b.shape)
	o = outerproduct(a,b)
	o.shape = a.shape + b.shape
	return _gen.concatenate(_gen.concatenate(o, axis=1), axis=1)
