def test_array_maskna_diagonal():
    # ndarray.diagonal
    a = np.arange(6, maskna=True)
    a.shape = (2,3)
    a[0,1] = np.NA

    # Should produce a view into a
    res = a.diagonal()
    assert_(res.base is a)
    assert_(res.flags.maskna)
    assert_(not res.flags.ownmaskna)
    assert_equal(res, [0, 4])

    res = a.diagonal(-1)
    assert_equal(res, [3])

    res = a.diagonal(-2)
    assert_equal(res, [])

    # This diagonal has the NA
    res = a.diagonal(1)
    assert_equal(np.isna(res), [1,0])
    assert_equal(res[~np.isna(res)], [5])

    res = a.diagonal(2)
    assert_equal(res, [2])

    res = a.diagonal(3)
    assert_equal(res, [])

def test_array_maskna_reshape():
    # Simple reshape 1D -> 2D
    a = np.arange(6, maskna=True)
    a[1] = np.NA
    a[5] = np.NA

    # Reshape from 1D to C order
    b = a.reshape(2,3)
    assert_(b.base is a)
    assert_equal(b.shape, (2,3))
    assert_(b.flags.maskna)
    assert_(not b.flags.ownmaskna)
    assert_equal(np.isna(b), [[0,1,0],[0,0,1]])

    # Reshape from 1D to F order
    b = a.reshape(2,3,order='F')
    assert_(b.base is a)
    assert_equal(b.shape, (2,3))
    assert_(b.flags.maskna)
    assert_(not b.flags.ownmaskna)
    assert_equal(np.isna(b), [[0,0,0],[1,0,1]])

    # Add a new axis using 'newaxis'
    a = np.array(np.NA, maskna=True)
    assert_equal(np.isna(a[np.newaxis]), [True])

def test_array_maskna_conjugate_method():
    # ndarray.conjugate
    a = np.array([1j, 2+4j, np.NA, 2-1.5j, np.NA], maskna=True)

    b = a.conjugate()
    assert_equal(np.isna(b), [0,0,1,0,1])
    assert_equal(b[~np.isna(b)], [-1j, 2-4j, 2+1.5j])

def check_ufunc_max_1D(max_func):
    a_orig = np.array([0, 3, 2, 10, -1, 5, 7, -2])
    a = a_orig.view(maskna=True)

    # Straightforward reduce with no NAs
    b = max_func(a)
    assert_equal(b, 10)

    # Set the biggest value to NA
    a[3] = np.NA
    b = max_func(a)
    assert_(np.isna(b))

    # Skip the NA
    b = max_func(a, skipna=True)
    assert_(not b.flags.maskna)
    assert_(not np.isna(b))
    assert_equal(b, 7)

    # Set the first value to NA
    a[0] = np.NA
    b = max_func(a, skipna=True)
    assert_(not b.flags.maskna)
    assert_(not np.isna(b))
    assert_equal(b, 7)

    # Set all the values to NA - should raise the same error as
    # for an empty array
    a[...] = np.NA
    assert_raises(ValueError, max_func, a, skipna=True)

def test_array_maskna_column_stack():
    # np.column_stack
    a = np.array((1,2,3), maskna=True)
    b = np.array((2,3,4), maskna=True)
    b[2] = np.NA
    res = np.column_stack((a,b))
    assert_equal(np.isna(res), [[0,0], [0,0], [0,1]])
    assert_equal(res[~np.isna(res)], [1,2,2,3,3])

def test_maskna_ufunc_1D():
    a_orig = np.arange(3)
    a = a_orig.view(maskna=True)
    b_orig = np.array([5,4,3])
    b = b_orig.view(maskna=True)
    c_orig = np.array([0,0,0])
    c = c_orig.view(maskna=True)

    # An NA mask is produced if an operand has one
    res = a + b_orig
    assert_(res.flags.maskna)
    assert_equal(res, [5,5,5])

    res = b_orig + a
    assert_(res.flags.maskna)
    assert_equal(res, [5,5,5])

    # Can still output to a non-NA array if there are no NAs
    np.add(a, b, out=c_orig)
    assert_equal(c_orig, [5,5,5])

    # Should unmask everything if the output has NA support but
    # the inputs don't
    c_orig[...] = 0
    c[...] = np.NA
    np.add(a_orig, b_orig, out=c)
    assert_equal(c, [5,5,5])

    # If the input has NA support but an output parameter doesn't,
    # should work as long as the inputs contain no NAs
    c_orig[...] = 0
    np.add(a, b, out=c_orig)
    assert_equal(c_orig, [5,5,5])

    # An NA is produced if either operand has one
    a[0] = np.NA
    b[1] = np.NA
    res = a + b
    assert_equal(np.isna(res), [1,1,0])
    assert_equal(res[2], 5)

    # If the output contains NA, can't have out= parameter without
    # NA support
    assert_raises(ValueError, np.add, a, b, out=c_orig)

    # Divide in-place with NA
    a_orig = np.array([[3], [12.]])
    a = a_orig.view(maskna=True)
    a[0,0] = np.NA
    a /= 3
    # Shouldn't have touched the masked element
    assert_array_equal(a_orig, [[3], [4.]])
    assert_array_equal(a, [[np.NA], [4.]])
    # double-check assertions
    assert_equal(np.isna(a), [[1], [0]])
    assert_equal(a[~np.isna(a)], [4.])

def check_maskna_ufunc_sum_1D(sum_func):
    a = np.arange(3.0, maskna=True)
    b = np.array(0.5)
    c_orig = np.array(0.5)
    c = c_orig.view(maskna=True)

    # Since 'a' has no NA values, this should work
    sum_func(a, out=b)
    assert_equal(b, 3.0)
    b[...] = 7
    sum_func(a, skipna=True, out=b)
    assert_equal(b, 3.0)

    ret = sum_func(a)
    assert_equal(ret, 3.0)
    ret = sum_func(a, skipna=True)
    assert_equal(ret, 3.0)

    # With an NA value, the reduce should throw with the non-NA output param
    a[1] = np.NA
    assert_raises(ValueError, sum_func, a, out=b)

    # With an NA value, the output parameter can still be an NA-array
    c_orig[...] = 0.5
    sum_func(a, out=c)
    assert_equal(c_orig, 0.5)
    assert_(np.isna(c))

    # Should not touch the out= element when assigning NA
    b[...] = 1.0
    d = b.view(maskna=True)
    sum_func(a, out=d)
    assert_(np.isna(d))
    assert_equal(b, 1.0)

    # Without an output parameter, return NA
    ret = sum_func(a)
    assert_(np.isna(ret))

    # With 'skipna=True'
    ret = sum_func(a, skipna=True)
    assert_equal(ret, 2.0)

    # With 'skipna=True', and out= parameter
    b[...] = 0.5
    sum_func(a, skipna=True, out=b)
    assert_equal(b, 2.0)

    # With 'skipna=True', and out= parameter with a mask
    c[...] = 0.5
    c[...] = np.NA
    sum_func(a, skipna=True, out=c)
    assert_(not np.isna(c))
    assert_equal(c, 2.0)

def test_array_maskna_squeeze():
    # np.squeeze
    a = np.zeros((1,3,1,1,4,2,1), maskna=True)
    a[0,1,0,0,3,0,0] = np.NA

    res = np.squeeze(a)
    assert_equal(res.shape, (3,4,2))
    assert_(np.isna(res[1,3,0]))

    res = np.squeeze(a, axis=(0,2,6))
    assert_equal(res.shape, (3,1,4,2))
    assert_(np.isna(res[1,0,3,0]))

def test_array_maskna_setasflat():
    # Copy from a C to a F array with some NAs
    a_orig = np.empty((2,3), order='C')
    b_orig = np.empty((3,2), order='F')
    a = a_orig.view(maskna=True)
    b = b_orig.view(maskna=True)
    a[...] = 1
    a[0,1] = np.NA
    a[1,2] = np.NA
    b[...] = 2
    b.setasflat(a)
    assert_equal(np.isna(a), [[0,1,0],[0,0,1]])
    assert_equal(b_orig, [[1,2],[1,1],[1,2]])
    assert_equal(np.isna(b), [[0,1],[0,0],[0,1]])

def test_isna():
    # Objects which are not np.NA or ndarray all return False
    assert_equal(np.isna(True), False)
    assert_equal(np.isna("abc"), False)
    assert_equal(np.isna([1,2,3]), False)
    assert_equal(np.isna({3:5}), False)
    # Various NA values return True
    assert_equal(np.isna(np.NA), True)
    assert_equal(np.isna(np.NA()), True)
    assert_equal(np.isna(np.NA(5)), True)
    assert_equal(np.isna(np.NA(dtype='f4')), True)
    assert_equal(np.isna(np.NA(12,dtype='f4')), True)

def test_array_maskna_astype():
    dtsrc = [np.dtype(d) for d in '?bhilqpBHILQPefdgFDGSUO']
    #dtsrc.append(np.dtype([('b', np.int, (1,))]))
    dtsrc.append(np.dtype('datetime64[D]'))
    dtsrc.append(np.dtype('timedelta64[s]'))

    dtdst = [np.dtype(d) for d in '?bhilqpBHILQPefdgFDGSUO']
    #dtdst.append(np.dtype([('b', np.int, (1,))]))
    dtdst.append(np.dtype('datetime64[D]'))
    dtdst.append(np.dtype('timedelta64[s]'))

    warn_ctx = WarningManager()
    warn_ctx.__enter__()
    try:
        warnings.simplefilter("ignore", np.ComplexWarning)
        for dt1 in dtsrc:
            a = np.ones(2, dt1, maskna=1)
            a[1] = np.NA
            for dt2 in dtdst:
                msg = 'type %s to %s conversion' % (dt1, dt2)
                b = a.astype(dt2)
                assert_(b.flags.maskna, msg)
                assert_(b.flags.ownmaskna, msg)
                assert_(np.isna(b[1]), msg)
    finally:
        warn_ctx.__exit__()

def test_copyto_fromscalar():
    a = np.arange(6, dtype='f4').reshape(2,3)

    # Simple copy
    np.copyto(a, 1.5)
    assert_equal(a, 1.5)
    np.copyto(a.T, 2.5)
    assert_equal(a, 2.5)

    # Where-masked copy
    mask = np.array([[0,1,0],[0,0,1]], dtype='?')
    np.copyto(a, 3.5, where=mask)
    assert_equal(a, [[2.5,3.5,2.5],[2.5,2.5,3.5]])
    mask = np.array([[0,1],[1,1],[1,0]], dtype='?')
    np.copyto(a.T, 4.5, where=mask)
    assert_equal(a, [[2.5,4.5,4.5],[4.5,4.5,3.5]])

    # Simple copy to NA-masked
    a_orig = a
    a = a_orig.view(maskna=True)
    a[...] = np.NA
    np.copyto(a, 0.5)
    assert_equal(a, 0.5)
    a[...] = np.NA
    np.copyto(a.T, 1.5)
    assert_equal(a, 1.5)

    # Where-masked copy to NA-masked
    a[0,0] = np.NA
    a[1,1] = np.NA
    mask = np.array([[1,0,1],[0,0,1]], dtype='?')
    np.copyto(a, 2.5, where=mask)
    assert_equal(np.isna(a), [[0,0,0],[0,1,0]])
    assert_equal(a_orig, [[2.5,1.5,2.5],[1.5,1.5,2.5]])

    # Simple preservena=True copy
    a[0,0] = np.NA
    np.copyto(a, 3.5, preservena=True)
    assert_equal(np.isna(a), [[1,0,0],[0,1,0]])
    assert_equal(a_orig, [[2.5,3.5,3.5],[3.5,1.5,3.5]])

    # Where-masked preservena=True copy
    mask = np.array([[1,0,1],[0,0,1]], dtype='?')
    np.copyto(a, 4.5, where=mask, preservena=True)
    assert_equal(np.isna(a), [[1,0,0],[0,1,0]])
    assert_equal(a_orig, [[2.5,3.5,4.5],[3.5,1.5,4.5]])

def test_array_maskna_construction():
    # Construction with NA inputs
    a = np.array([1.0, 2.0, np.NA, 7.0], maskna=True)
    assert_equal(a.dtype, np.dtype('f8'))
    assert_(a.flags.maskna)
    assert_equal(type(a[2]), np.NAType)
    # Without the 'maskna=True', still produces an NA mask if NA is there
    a = np.array([1.0, 2.0, np.NA, 7.0])
    assert_equal(a.dtype, np.dtype('f8'))
    assert_(a.flags.maskna)
    assert_equal(type(a[2]), np.NAType)
    # Without any NAs, does not produce an NA mask
    a = np.array([1.0, 2.0, 4.0, 7.0])
    assert_equal(a.dtype, np.dtype('f8'))
    assert_(not a.flags.maskna)

    # From np.NA as a straight scalar
    a = np.array(np.NA, maskna=True)
    assert_equal(type(a), np.ndarray)
    assert_(np.isna(a))

    # As a special case, converting np.NA to an array produces
    # a zero-dimensional masked array
    a = np.array(np.NA)
    assert_equal(type(a), np.ndarray)
    assert_(np.isna(a))

    # The data type defaults to the same as an empty array if all is NA
    a = np.array([np.NA], maskna=True)
    b = np.array([])
    assert_equal(a.dtype, b.dtype)
    assert_(np.isna(a))

    a = np.zeros((3,))
    assert_(not a.flags.maskna)
    a = np.zeros((3,), maskna=True)
    assert_(a.flags.maskna)
    assert_equal(np.isna(a), False)

    a = np.ones((3,))
    assert_(not a.flags.maskna)
    a = np.ones((3,), maskna=True)
    assert_(a.flags.maskna)
    assert_equal(np.isna(a), False)

    # np.empty returns all NAs if maskna is set to True
    a = np.empty((3,))
    assert_(not a.flags.maskna)
    a = np.empty((3,), maskna=True)
    assert_(a.flags.maskna)
    assert_equal(np.isna(a), True)

    # np.empty_like returns all NAs if maskna is set to True
    tmp = np.arange(3)
    a = np.empty_like(tmp)
    assert_(not a.flags.maskna)
    a = np.empty_like(tmp, maskna=True)
    assert_(a.flags.maskna)
    assert_equal(np.isna(a), True)

def test_maskna_take_1D():
    a = np.arange(5, maskna=True)
    b = np.arange(3)
    c = b.view(maskna=True)

    # Take without any NAs
    assert_equal(a.take([0,2,4]), [0,2,4])

    # Take without any NAs, into non-NA output parameter
    a.take([0,2,4], out=b)
    assert_equal(b, [0,2,4])

    # Take without any NAs, into NA output parameter
    b[...] = 1
    c[...] = np.NA
    a.take([0,2,4], out=c)
    assert_equal(c, [0,2,4])

    # Take with some NAs
    a[2] = np.NA
    a[3] = np.NA
    ret = a.take([0,2,4])
    assert_equal([ret[0], ret[2]], [0,4])
    assert_equal(np.isna(ret), [0,1,0])

    # Take with some NAs, into NA output parameter
    b[...] = 1
    c[...] = np.NA
    a.take([0,2,4], out=c)
    assert_equal(b, [0,1,4])
    assert_equal([c[0], c[2]], [0,4])
    assert_equal(np.isna(c), [0,1,0])

    c[...] = 1
    a.take([0,2,4], out=c)
    assert_equal(b, [0,1,4])
    assert_equal([c[0], c[2]], [0,4])
    assert_equal(np.isna(c), [0,1,0])

    # Take with an NA just at the start
    a = np.arange(5, maskna=True)
    a[0] = np.NA
    res = a.take([1,2,3,4])
    assert_equal(res, [1,2,3,4])

def test_array_maskna_view_NA_assignment_1D():
    a = np.arange(10)
    a_ref = a.copy()

    # Make sure that assigning NA doesn't affect the original data
    b = a.view(maskna=True)
    b[...] = np.NA
    assert_equal(np.isna(b), True)
    assert_equal(a, a_ref)

    b = a.view(maskna=True)
    b[:] = np.NA
    assert_equal(np.isna(b), True)
    assert_equal(a, a_ref)

    b = a.view(maskna=True)
    b[3:5] = np.NA
    assert_equal(np.isna(b), [0,0,0,1,1,0,0,0,0,0])
    assert_equal(a, a_ref)

    b = a.view(maskna=True)
    b[3:10:3] = np.NA
    assert_equal(np.isna(b), [0,0,0,1,0,0,1,0,0,1])
    assert_equal(a, a_ref)

    b = a.view(maskna=True)
    b[3] = np.NA
    assert_equal(np.isna(b), [0,0,0,1,0,0,0,0,0,0])
    assert_equal(a, a_ref)

    b = a.view(maskna=True)
    mask = np.array([0,1,0,1,1,0,0,0,1,1], dtype='?')
    b[mask] = np.NA
    assert_equal(np.isna(b), mask)
    assert_equal(a, a_ref)

def test_array_maskna_ravel():
    # From a C array
    a = np.zeros((2,3), maskna=True, order='C')
    a[0,1] = np.NA
    a[1,2] = np.NA

    # Ravel in C order returns a view
    b = np.ravel(a)
    assert_(b.base is a)
    assert_equal(b.shape, (6,))
    assert_(b.flags.maskna)
    assert_(not b.flags.ownmaskna)
    assert_equal(np.isna(b), [0,1,0,0,0,1])

    # Ravel in F order returns a copy
    b = np.ravel(a, order='F')
    assert_(b.base is None)
    assert_equal(b.shape, (6,))
    assert_(b.flags.maskna)
    assert_(b.flags.ownmaskna)
    assert_equal(np.isna(b), [0,0,1,0,0,1])

    a = np.arange(12, maskna=True).reshape(2,3,2).swapaxes(1,2)
    assert_equal(a.ravel(order='K'), np.arange(12))

def test_maskna_nonzero_1D():
    a = np.zeros((5,), maskna=True)

    # The nonzeros without any NAs
    assert_equal(np.count_nonzero(a), 0)
    assert_equal(np.nonzero(a)[0], [])
    a[2] = 3
    assert_equal(np.count_nonzero(a), 1)
    assert_equal(np.nonzero(a)[0], [2])
    a[3:] = 2
    assert_equal(np.count_nonzero(a), 3)
    assert_equal(np.nonzero(a)[0], [2,3,4])

    # The nonzeros with an NA
    a[2] = np.NA
    assert_(np.isna(np.count_nonzero(a)))
    assert_raises(ValueError, np.nonzero, a)

def test_array_maskna_view_NA_assignment_2D():
    a = np.arange(6).reshape(2,3)
    a_ref = a.copy()

    # Make sure that assigning NA doesn't affect the original data
    b = a.view(maskna=True)
    b[...] = np.NA
    assert_equal(np.isna(b), True)
    assert_equal(a, a_ref)

    b = a.view(maskna=True)
    b[:] = np.NA
    assert_equal(np.isna(b), True)
    assert_equal(a, a_ref)

    b = a.view(maskna=True)
    b[0,:] = np.NA
    assert_equal(np.isna(b[0]), True)
    assert_equal(np.isna(b[1]), False)
    assert_equal(a, a_ref)

    b = a.view(maskna=True)
    b[1:,1:3] = np.NA
    assert_equal(np.isna(b), [[0,0,0],[0,1,1]])
    assert_equal(a, a_ref)

    b = a.view(maskna=True)
    b[1,::2] = np.NA
    assert_equal(np.isna(b), [[0,0,0],[1,0,1]])
    assert_equal(a, a_ref)

    b = a.view(maskna=True)
    b[0,2] = np.NA
    assert_equal(np.isna(b), [[0,0,1],[0,0,0]])
    assert_equal(a, a_ref)

    b = a.view(maskna=True)
    mask = np.array([[1,0,1],[1,1,0]], dtype='?')
    b[mask] = np.NA
    assert_equal(np.isna(b), mask)
    assert_equal(a, a_ref)

def test_array_maskna_item():
    # With a zero-dimensional array
    a = np.array(np.NA, maskna=True)

    # Should return NA as the item
    assert_equal(type(a.item()), np.NAType)

    # Should be able to set the item
    a.itemset(1.5)
    assert_(not np.isna(a))
    assert_equal(a, 1.5)
    a.itemset(np.NA)
    assert_(np.isna(a))

    # With a one-dimensional array
    a = np.array([1, np.NA, 2, np.NA], maskna=True)

    # Should return the scalar or NA as the item
    assert_(not np.isna(a.item(0)))
    assert_equal(type(a.item(1)), np.NAType)

    # Should be able to set the items
    a.itemset(0, np.NA)
    assert_(np.isna(a[0]))
    a.itemset(1, 12)
    assert_(not np.isna(a[1]))
    assert_equal(a[1], 12)

    # With a two-dimensional array
    a = np.arange(6, maskna=True).reshape(2,3)
    a[0,1] = np.NA
    # Should return the scalar or NA as the item
    assert_(not np.isna(a.item((0,0))))
    assert_equal(type(a.item((0,1))), np.NAType)

    # Should be able to set the items
    a.itemset((0,1), 8)
    assert_(not np.isna(a[0,1]))
    assert_equal(a[0,1], 8)
    a.itemset((1,1), np.NA)
    assert_(np.isna(a[1,1]))

def test_array_maskna_std_mean_methods():
    # ndarray.std, ndarray.mean
    a = np.array([[2, np.NA, 10],
                  [4, 8, 7],
                  [12, 4, np.NA]], maskna=True)

    res = a.mean(axis=0)
    assert_equal(np.isna(res), [0,1,1])
    assert_equal(res[~np.isna(res)], [np.array([2,4,12]).mean()])

    res = a.mean(axis=-1)
    assert_equal(np.isna(res), [1,0,1])
    assert_equal(res[~np.isna(res)], [np.array([4,8,7]).mean()])

    res = a.std(axis=0)
    assert_equal(np.isna(res), [0,1,1])
    assert_equal(res[~np.isna(res)], [np.array([2,4,12]).std()])

    res = a.std(axis=-1)
    assert_equal(np.isna(res), [1,0,1])
    assert_equal(res[~np.isna(res)], [np.array([4,8,7]).std()])