"""MA facility Copyright 1999 Regents of the University of California.
Released for unlimited redistribution; see file Legal.htm in Numerical distribution."""
__version__ = 2.0
#This version gives up on the Extension Class stuff
import Numeric
from Precision import *
import string, types

class MAError (Exception):
    def __init__ (self, args=None):
        self.args = args
    def __str__(self):
        return str(self.args)

__MaxElements = 300     #Maximum size for printing

def setMaxElements (m):
    "Set the maximum # of elements for printing arrays. "
    global __MaxElements
    n = int(m)
    if n <= 0: raise ValueError, "setMaxElements requires positive argument."
    __MaxElements = n

def getMaxElements ():
    "Get the maximum # of elements for printing arrays. "
    global __MaxElements
    return __MaxElements

def getmask (a):
    "Mask of values in a, if any; None otherwise."
    if isMA(a): 
        return a.mask()
    else: 
        return None

def getmaskarray (a):
    "Mask of values in a, ones if none."
    m = getmask(a)
    if m is None:
        return Numeric.ones(shape(a))
    else:
        return m

def mask_or (m1, m2):
    "Logical or of the masks m1 and m2, treating None as false."
    if m1 is None: return m2
    if m2 is None: return m1
    return Numeric.logical_or(m1, m2)

def filled (a, value = None):
    """
    a as a Numeric array with any masked areas replaced by value"
    No data copied if no masked area and data is already an array.
    """
    if isMA(a):
        return a.filled(value)
    else:
        return Numeric.array(a, copy=0)

class masked_unary_operation:
    def __init__ (self, aufunc, fill=0):
        self.f = aufunc
        self.fill = fill

    def __call__ (self, a):
        m = getmask(a)
        d1 = filled(a)
        if m is None:
            return self.f(d1)
        else:
            return MA(self.f(d1), m)

class masked_binary_operation:
    def __init__ (self, aufunc, fill=0):
        self.f = aufunc
        self.fill = fill

    def __call__ (self, a, b):
        d1 = filled(a, self.fill)
        d2 = filled(b, self.fill)
        m = mask_or(getmask(a), getmask(b))
        if m is None:
            return self.f(d1, d2)
        else:
            return MA(self.f(d1,d2), m)

    def reduce (self, target, axis=0):
        m = getmask(target)
        if m is None:
            return self.f.reduce (target, axis)
        else:
            t = target.filled(self.fill)
            tr = self.f.reduce(t, axis)
            m = Numeric.logical_and.reduce(m, axis)
            return MA(tr, m)

class comparison:
    "A comparison function that returns a plain mask after doing the comparison"
    def __init__ (self, f, fill1, fill2):
        self.f = f
        self.fill1 = fill1
        self.fill2 = fill2
    def __call__ (self, a, b):
        return self.f(filled(a, self.fill1), filled(b, self.fill2))


# class MA        
# the use of filled in the constructor
# makes sure we don't get MA's in data or mask 
# This prevents recursion
# and we can assume these are array objects (or None for mask).
# Note current implementation assumes mask is 1's and 0's
# so that we can use Numeric.choose in filled.

class MA:
    "Arrays with possibly masked values."
    default_real_fill_value = 0.0
    default_complex_fill_value = 0.0 + 0.0j
    default_character_fill_value = ' '
    default_integer_fill_value = 0
    default_unsigned_integer_fill_value = 0
    default_object_fill_value = None

    def __init__(self, data, mask=None, **keys):
        """MA(data, mask=None)"""
        self.__data = filled(data)
        if mask is None:
            self.__mask = None
        else:
            self.__mask = filled(mask, 1)

    def __getattr__ (self, name):
        if name == 'shape':
            return self.get_shape()
        else:
            return self.__dict__[name]

    def __setattr__ (self, name, value):
        if name == 'shape':
            self.set_shape(value)
        else:
            self.__dict__[name] = value

    def __str__(self):
        return str(self.filled())

    def __repr__(self):
        if self.__mask is None:
            return "MA(" + repr(filled(self)) + ")"
        else:
            return "MA(" + repr(filled(self)) + ", \n   " + \
                           repr(self.__mask) + ")"

    def __float__(self):
        return MA (float(self.filled(0.0)), self.__mask)

    def __int__(self):
        return MA(int(self.filled(0)), self.__mask)

# Note copy semantics here differ from Numeric        
    def __getitem__(self, i): 
        m = self.__mask
        if m is None:
            return Numeric.array(self.__data[i])
        else:
            return MA(Numeric.array(self.__data[i]), Numeric.array(m[i]))

    def __getslice__(self, i, j): 
        m = self.__mask
        if m is None:
            return Numeric.array(self.__data[i:j])
        else:
            return MA(Numeric.array(self.__data[i:j]), Numeric.array(m[i:j]))
# --------
    def __setitem__(self, index, value):
        self.__data[index] = filled(value)
        self.__mask = mask_or(self.__mask, getmask(value))
 
    def __setslice__(self, i, j, value):
        self.__data[i:j] = filled(value)
        m = getmask(value)
        if m is not None:
            sm = getmaskarray(self)
            sm[i:j] = mask_or(sm[i:j], m)
            self.__mask = sm

    def __len__ (self):
        "Return total number of elements in array."
        return self.size()

    def __cmp__ (self, other):
        raise MAError, "Cannot use comparison operators on MA instances."
 
    def __abs__(self): 
        return masked_unary_operation(Numeric.absolute)(self)

    def __neg__(self): 
        return masked_unary_operation(Numeric.negative)(self)

    def __add__(self, other): 
        return masked_binary_operation(Numeric.add)(self, other)
                        
    __radd__ = __add__

    def __sub__(self, other): 
        return masked_binary_operation(Numeric.subtract,0)(self, other)

    def __rsub__(self, other): 
        return masked_binary_operation(Numeric.subtract,0)(other, self)

    def __mul__(self, other): 
        return masked_binary_operation(Numeric.multiply,0)(self, other)
    
    __rmul__ = __mul__

    def __div__(self, other): 
        return divide(self, other)

    def __rdiv__(self, other): 
        return divide(other, self)

    def __pow__(self,other, third=None): 
        if third is not None: 
            raise MAException, "3 arg power unspoorted"
        return masked_binary_operation(Numeric.power, 1)(self, other)

    def __sqrt__(self): 
        return masked_unary_operation(Numeric.sqrt,0)(self)

    def default_fill_value (self):
        "Function to calculate default fill value for an array."
        x = self.typecode()
        if x in typecodes['Float']:
            return MA.default_real_fill_value
        if x in typecodes['Integer']:
            return MA.default_integer_fill_value
        if x in typecodes['Complex']:
            return MA.default_complex_fill_value
        if x in typecodes['Character']:
            return MA.default_character_fill_value
        if x in typecodes['UnsignedInteger']:
            return Numeric.absolute(MA.default_integer_fill_value)
        else:
            return MA.default_object_fill_value

    def set_fill_value (self, v):
        "Set the fill value to v."
        self.fill_value = v

    def filled (self, value=None):
        """ Self with masked data replaced by value.
            If value is None, self.fill_value used.
            self.default_fill_value () used if value is None.
            Not a copy if no masked data.
            Result is romised to be a Numeric.array.
        """
        if self.__mask is None:
            return self.__data
        else:
            if value is None:
                if hasattr(self, 'fill_value'):
                    v = self.fill_value
                else:
                    v = self.default_fill_value ()
                    self.fill_value = v
            else:
                v = value
            return Numeric.choose(self.__mask, (self.__data, v))

    def get_shape(self):
        "Return the tuple giving the current shape."
        return self.__data.shape

    def ids (self):
        """Return the ids of the data and mask areas"""
        return (id(self.__data), id(self.__mask))

    def mask(self):
        "Return the data mask, or None."
        return self.__mask
            
    def set_shape (self, *sizes):
        """shape(n, m, ...) sets the shape."""
        self.__data.shape = sizes
        if self.__mask is not None:
            self.__mask.shape = sizes
            
    def size (self, axis = None):
        "Number of elements in array, or in a particular axis."
        s = self.shape
        if axis is None:
            return reduce(Numeric.multiply, s, 1)
        else:
            return s[axis]
            
    def typecode(self):
        return self.__data.typecode()

# these are disgusting
    def is_contiguous (self):
        return self.data.is_contiguous()

    def byte_swapped(self):
        return self.data.byte_swapped()
    
def isMA (x):
    "Is x an instance of MA?"
    return isinstance(x, MA)

def masked_array (data, masked_value, typecode=None, copy=1, atol=1.e-8, rtol=1.e-8, notest=0):
    """Create a masked array for numerical data; no mask unless necessary
       or if notest != 0
    """
    abs = Numeric.absolute
    d = Numeric.array(data, typecode=typecode, copy=copy)
    dm = Numeric.less(abs(d-masked_value), atol+rtol*abs(d))
    if notest or Numeric.sometrue(Numeric.ravel(dm)):
        result = MA(d, dm)
    else:
        result = MA(d)
    result.set_fill_value(masked_value)
    return result
    

def masked_object (data, masked_value, copy=1):
    "Create a masked array of places where exactly equal to masked value"
    dm = Numeric.equal(data, masked_value, copy=copy)
    return MA(data, dm)
    
def rank (a):
    "Get the rank of a"
    return len(shape(a))
        
def shape (a):
    "Get the shape of a"
    if isMA (a):
        return a.get_shape()
    else:
        return Numeric.shape(a)

def size (a, axis=None):
    "Get the number of elements in a, or along a certain axis."
    if isMA (a):
        return a.size(axis)
    else:
        s = Numeric.shape(a)
        if axis is None:
            return reduce(Numeric.multiply, s, 1)
        else:
            return s[axis]

def count (a, axis = None):
    "Count of the non-masked elements."   
    m = getmask(a)
    if m is None:
        return size(a, axis)
    else:
        n1 = size(a, axis)
        n2 = sum (m, axis)
        return n1 - n2

def sum (a, axis = None):
    "Sum of all elements, or along a certain axis"
    if axis is None:
        return Numeric.add.reduce(Numeric.ravel(filled(a, 0)))
    else:
        return Numeric.add.reduce(filled(a, 0), axis)
 
def product (a, axis = None):
    "Product of all elements, or along a certain axis."
    if axis is None:
        return Numeric.add.reduce(Numeric.ravel(filled(a, 0)))
    else:
        return Numeric.add.reduce(filled(a, 0), axis)

def average(a, axis=None):
    "Average of the nonmasked elements of a"
    c = 1.0 * count(a, axis)
    a = sum(a, axis) 
    return divide(a, c)
    
NewAxis = Numeric.NewAxis
arange = Numeric.arange
ones = Numeric.ones
zeros = Numeric.zeros
array = Numeric.array

#minimum, maximum?
#searchsorted?

sqrt = masked_unary_operation(Numeric.sqrt, 0.0)
sin = masked_unary_operation(Numeric.sin, 0.0)
cos = masked_unary_operation(Numeric.cos, 0.0)
absolute = masked_unary_operation(Numeric.absolute, 0)
negative = masked_unary_operation(Numeric.negative, 0)
absolute = masked_unary_operation(Numeric.absolute, 0)
nonzero = masked_unary_operation(Numeric.nonzero, 0)

add = masked_binary_operation(Numeric.add, 0)
subtract = masked_binary_operation(Numeric.subtract, 0)
multiply = masked_binary_operation(Numeric.multiply, 1)
power = masked_binary_operation(Numeric.power, 1)

sometrue = masked_unary_operation(Numeric.sometrue, 0)
alltrue = masked_unary_operation(Numeric.alltrue, 1)

# all these comparisons return false where masks are true       
equal = comparison(Numeric.equal, 0, 1)
not_equal = comparison(Numeric.not_equal, 0, 0)
less_equal = comparison(Numeric.less_equal, 1, 0)
greater_equal = comparison(Numeric.greater_equal, 0, 1)
less = comparison(Numeric.less, 1, 0)
greater = comparison(Numeric.greater, 0, 1)

def choose (condition, t):
    "Shaped like condition, values t[0] where condition true, t[1] otherwise"
    d = Numeric.choose(filled(condition,0), map(filled, t))
    m = Numeric.choose(filled(condition,0), map(getmaskarray, t))
    return MA(d, m)

def where (condition, x, y):
    return choose(Numeric.not_equal(condition, 0), (y, x))

def reshape (a, newshape):
    "Copy of a with a new shape."
    m = getmask(a)
    d = Numeric.reshape(filled(a), newshape)
    if m is None:
        return d
    else:
        return MA(d, Numeric.reshape(m, newshape))

def ravel (a):
    "a as one-dimensional, may share data and mask"
    m = getmask(a)
    d = Numeric.ravel(filled(a))   
    if m is None:
        return d
    else:
        return MA(d, Numeric.ravel(m))

def concatenate (somearrays, axis=0):
    "concatenate the arrays along the given axis"
    d = Numeric.concatenate(map(filled, somearrays), axis)
    dm = Numeric.concatenate(map(getmaskarray, somearrays), axis)
    return MA(d, dm)

def take (a, indices, axis=0):
    m = getmask(a)
    d = filled(a)
    if m is None:
        return Numeric.take(d, indices, axis)
    else:
        return MA(Numeric.take(d, indices, axis), 
                  Numeric.take(m, indices, axis))
                           
def divide (x, y):
    "a/b, masked where b was masked or zero"
    a = filled(x, 0)
    b = filled(y, 1)
    bad_elements = equal(b, 0)
    c_mask = mask_or(mask_or(getmask(a), getmask(b)), bad_elements)
    return MA(Numeric.divide(a, b), c_mask)
    
# This section is stolen from a post about how to limit array printing.

def limitedArrayRepr(a, max_line_width = None, precision = None, suppress_small = None):
    global __MaxElements
    s = a.shape
    elems =  Numeric.multiply.reduce(s)
    if elems > __MaxElements:
        if len(s) > 1:
            return 'array (%s) , type = %s, has %d elements' % \
                 (string.join(map(str, s), ","), a.typecode(), elems)
        else:
            return Numeric.array2string (a[:__MaxElements], max_line_width, precision,
                 suppress_small,',',0) + \
               ('\n + %d more elements' % (elems - __MaxElements))
    else:
        return Numeric.array2string (a, max_line_width, precision, 
                suppress_small,',',0)


# replace the default with a smarter function for huge arrays
original_array_repr = Numeric.array_repr
original_array_str = Numeric.array_str
Numeric.array_repr = limitedArrayRepr
Numeric.array_str = limitedArrayRepr

# fixup multiarray, as well.
import multiarray
multiarray.set_string_function(limitedArrayRepr, 0)
multiarray.set_string_function(limitedArrayRepr, 1)

if __name__ == '__main__':
        import MAtest