from MA import *
def put(*stuff):
    for x in stuff:
        print x,
    print

def test (va, mv):
    ua = masked_array(va, mv)
    put("shape(ua), ua.shape, ua.get_shape()", shape(ua), ua.shape, ua.get_shape())
    put("rank(ua), size(ua), size(ua,0), size(ua,1)\n", 
         rank(ua), size(ua), size(ua,0), size(ua,1))
    put("mask(ua)\n", getmask(ua))
    put("repr(ua)\n", repr(ua))
    put("va\n", va)
    put("ua\n", ua)
    put ("ua.filled(-1)\n", ua.filled(-1))
    put("isMA(ua)", isMA(ua))
    put("isMA(va)", isMA(va))
    put("ua * va", type(ua*va), '\n', ua*va)
    put("va * ua", type(va*ua), '\n', va*ua)

    put("ua[:,2]\n", ua[:,2])
    u1 = ravel(ua)
    put("u1", u1.filled(), u1.mask())
    
    put("u1[7]\n", u1[7])
    put("u1[4:12]\n", u1[4:12])

    put("cos(ua)\n", cos(ua))
    put("abssolute(-ua)\n", absolute(-ua))
    put("ua+ua\n", ua+ua)
    put("ua**2\n", ua**2)
    put("sum(ua)\n", sum(ua))
    put("sum(ua, 1)\n", sum(ua, 1))
    put("equal(ua, va)\n", equal(ua, va))
    put("average(ua)\n", average(ua))
    put("average(ua, 1)\n", average(ua, 1))
    put("take(ua, (1,2))\n", take(ua, (1,2)))
    put("take(ua, (1,2), 1)\n", take(ua, (1,2), 1))

va=reshape(arange(12), (4,3))
mv = 99
va[0,1]=mv
va[2,2] = mv
test (va, mv)

va = reshape(arange(12)*1.0, (4,3))
mv = 1.e20
va[0,1] = mv
va[2,2] = mv
test (va, mv)

# now try pathology
va=reshape(arange(12)*1.0, (4,3))
mv = 1.e20
va[0,1] = mv
va[2,2] = mv
va[:, 2] = mv
test (va, mv)

a = arange(4, typecode=Float)
b = arange(4, typecode=Float) -2.0
am = MA(a, [1,0,0,0])
bm = MA(b, [0,0,0,1])
put("am\n", am)
put("bm\n", bm)
put("a\n", a)
put("b\n", b)
put("divide(a, b)\n", divide(a, b))
put("am / b\n", am / b)
put("a / bm\n", a / bm)
put("am / bm\n", am / bm)

ua = masked_array(va, mv, copy=0)
ub = MA(va, ua.mask()) #shares a mask and data
ua[1,0] = 22.0
uc = MA(ua, ua.mask()) #shares only the mask
ua.mask()[0,0] = 1
put("ids\n", map(lambda x: x.ids(), (ua,ub,uc)))
put("3 tests of sharing, should be all zeros\n")
put("ub[1,0] - ua[1,0]\n", ub[1,0] - ua[1, 0])
put("ua.filled(10.) - ub.filled(10.)\n", ua.filled(10.) - ub.filled(10.))
put("ua - uc\n", ua - uc)
wo = where([1,1,0,0], am, am*10.)
wo.missing_value = 1.e20
put("wo\n", wo)
ns = arange(12, typecode=Float)
ms = masked_array(ns, 5.0)
qs = masked_array(ns, 100.0)
ns = reshape(ns, (4,3))
ms.set_shape(4,3)
qs.set_shape(4,3)
put("Should be (4,3)\n", shape(ns))
put("Should be (4,3)\n", shape(ms))
put("Should be (4,3)\n", shape(qs))