import sys
import warnings
import functools
import operator

import numpy as np
from numpy.core._multiarray_tests import array_indexing
from itertools import product
from numpy.testing import (
    assert_, assert_equal, assert_raises, assert_array_equal, assert_warns,
    HAS_REFCOUNT,
    )


class TestIndexing:
    def test_index_no_floats(self):
        a = np.array([[[5]]])

        assert_raises(IndexError, lambda: a[0.0])
        assert_raises(IndexError, lambda: a[0, 0.0])
        assert_raises(IndexError, lambda: a[0.0, 0])
        assert_raises(IndexError, lambda: a[0.0,:])
        assert_raises(IndexError, lambda: a[:, 0.0])
        assert_raises(IndexError, lambda: a[:, 0.0,:])
        assert_raises(IndexError, lambda: a[0.0,:,:])
        assert_raises(IndexError, lambda: a[0, 0, 0.0])
        assert_raises(IndexError, lambda: a[0.0, 0, 0])
        assert_raises(IndexError, lambda: a[0, 0.0, 0])
        assert_raises(IndexError, lambda: a[-1.4])
        assert_raises(IndexError, lambda: a[0, -1.4])
        assert_raises(IndexError, lambda: a[-1.4, 0])
        assert_raises(IndexError, lambda: a[-1.4,:])
        assert_raises(IndexError, lambda: a[:, -1.4])
        assert_raises(IndexError, lambda: a[:, -1.4,:])
        assert_raises(IndexError, lambda: a[-1.4,:,:])
        assert_raises(IndexError, lambda: a[0, 0, -1.4])
        assert_raises(IndexError, lambda: a[-1.4, 0, 0])
        assert_raises(IndexError, lambda: a[0, -1.4, 0])
        assert_raises(IndexError, lambda: a[0.0:, 0.0])
        assert_raises(IndexError, lambda: a[0.0:, 0.0,:])

    def test_slicing_no_floats(self):
        a = np.array([[5]])

        # start as float.
        assert_raises(TypeError, lambda: a[0.0:])
        assert_raises(TypeError, lambda: a[0:, 0.0:2])
        assert_raises(TypeError, lambda: a[0.0::2, :0])
        assert_raises(TypeError, lambda: a[0.0:1:2,:])
        assert_raises(TypeError, lambda: a[:, 0.0:])
        # stop as float.
        assert_raises(TypeError, lambda: a[:0.0])
        assert_raises(TypeError, lambda: a[:0, 1:2.0])
        assert_raises(TypeError, lambda: a[:0.0:2, :0])
        assert_raises(TypeError, lambda: a[:0.0,:])
        assert_raises(TypeError, lambda: a[:, 0:4.0:2])
        # step as float.
        assert_raises(TypeError, lambda: a[::1.0])
        assert_raises(TypeError, lambda: a[0:, :2:2.0])
        assert_raises(TypeError, lambda: a[1::4.0, :0])
        assert_raises(TypeError, lambda: a[::5.0,:])
        assert_raises(TypeError, lambda: a[:, 0:4:2.0])
        # mixed.
        assert_raises(TypeError, lambda: a[1.0:2:2.0])
        assert_raises(TypeError, lambda: a[1.0::2.0])
        assert_raises(TypeError, lambda: a[0:, :2.0:2.0])
        assert_raises(TypeError, lambda: a[1.0:1:4.0, :0])
        assert_raises(TypeError, lambda: a[1.0:5.0:5.0,:])
        assert_raises(TypeError, lambda: a[:, 0.4:4.0:2.0])
        # should still get the DeprecationWarning if step = 0.
        assert_raises(TypeError, lambda: a[::0.0])

    def test_index_no_array_to_index(self):
        # No non-scalar arrays.
        a = np.array([[[1]]])

        assert_raises(TypeError, lambda: a[a:a:a])

    def test_none_index(self):
        # `None` index adds newaxis
        a = np.array([1, 2, 3])
        assert_equal(a[None], a[np.newaxis])
        assert_equal(a[None].ndim, a.ndim + 1)

    def test_empty_tuple_index(self):
        # Empty tuple index creates a view
        a = np.array([1, 2, 3])
        assert_equal(a[()], a)
        assert_(a[()].base is a)
        a = np.array(0)
        assert_(isinstance(a[()], np.int_))

    def test_void_scalar_empty_tuple(self):
        s = np.zeros((), dtype='V4')
        assert_equal(s[()].dtype, s.dtype)
        assert_equal(s[()], s)
        assert_equal(type(s[...]), np.ndarray)

    def test_same_kind_index_casting(self):
        # Indexes should be cast with same-kind and not safe, even if that
        # is somewhat unsafe. So test various different code paths.
        index = np.arange(5)
        u_index = index.astype(np.uintp)
        arr = np.arange(10)

        assert_array_equal(arr[index], arr[u_index])
        arr[u_index] = np.arange(5)
        assert_array_equal(arr, np.arange(10))

        arr = np.arange(10).reshape(5, 2)
        assert_array_equal(arr[index], arr[u_index])

        arr[u_index] = np.arange(5)[:,None]
        assert_array_equal(arr, np.arange(5)[:,None].repeat(2, axis=1))

        arr = np.arange(25).reshape(5, 5)
        assert_array_equal(arr[u_index, u_index], arr[index, index])

    def test_empty_fancy_index(self):
        # Empty list index creates an empty array
        # with the same dtype (but with weird shape)
        a = np.array([1, 2, 3])
        assert_equal(a[[]], [])
        assert_equal(a[[]].dtype, a.dtype)

        b = np.array([], dtype=np.intp)
        assert_equal(a[[]], [])
        assert_equal(a[[]].dtype, a.dtype)

        b = np.array([])
        assert_raises(IndexError, a.__getitem__, b)

    def test_ellipsis_index(self):
        a = np.array([[1, 2, 3],
                      [4, 5, 6],
                      [7, 8, 9]])
        assert_(a[...] is not a)
        assert_equal(a[...], a)
        # `a[...]` was `a` in numpy <1.9.
        assert_(a[...].base is a)

        # Slicing with ellipsis can skip an
        # arbitrary number of dimensions
        assert_equal(a[0, ...], a[0])
        assert_equal(a[0, ...], a[0,:])
        assert_equal(a[..., 0], a[:, 0])

        # Slicing with ellipsis always results
        # in an array, not a scalar
        assert_equal(a[0, ..., 1], np.array(2))

        # Assignment with `(Ellipsis,)` on 0-d arrays
        b = np.array(1)
        b[(Ellipsis,)] = 2
        assert_equal(b, 2)

    def test_single_int_index(self):
        # Single integer index selects one row
        a = np.array([[1, 2, 3],
                      [4, 5, 6],
                      [7, 8, 9]])

        assert_equal(a[0], [1, 2, 3])
        assert_equal(a[-1], [7, 8, 9])

        # Index out of bounds produces IndexError
        assert_raises(IndexError, a.__getitem__, 1 << 30)
        # Index overflow produces IndexError
        assert_raises(IndexError, a.__getitem__, 1 << 64)

    def test_single_bool_index(self):
        # Single boolean index
        a = np.array([[1, 2, 3],
                      [4, 5, 6],
                      [7, 8, 9]])

        assert_equal(a[np.array(True)], a[None])
        assert_equal(a[np.array(False)], a[None][0:0])

    def test_boolean_shape_mismatch(self):
        arr = np.ones((5, 4, 3))

        index = np.array([True])
        assert_raises(IndexError, arr.__getitem__, index)

        index = np.array([False] * 6)
        assert_raises(IndexError, arr.__getitem__, index)

        index = np.zeros((4, 4), dtype=bool)
        assert_raises(IndexError, arr.__getitem__, index)

        assert_raises(IndexError, arr.__getitem__, (slice(None), index))

    def test_boolean_indexing_onedim(self):
        # Indexing a 2-dimensional array with
        # boolean array of length one
        a = np.array([[ 0.,  0.,  0.]])
        b = np.array([ True], dtype=bool)
        assert_equal(a[b], a)
        # boolean assignment
        a[b] = 1.
        assert_equal(a, [[1., 1., 1.]])

    def test_boolean_assignment_value_mismatch(self):
        # A boolean assignment should fail when the shape of the values
        # cannot be broadcast to the subscription. (see also gh-3458)
        a = np.arange(4)

        def f(a, v):
            a[a > -1] = v

        assert_raises(ValueError, f, a, [])
        assert_raises(ValueError, f, a, [1, 2, 3])
        assert_raises(ValueError, f, a[:1], [1, 2, 3])

    def test_boolean_assignment_needs_api(self):
        # See also gh-7666
        # This caused a segfault on Python 2 due to the GIL not being
        # held when the iterator does not need it, but the transfer function
        # does
        arr = np.zeros(1000)
        indx = np.zeros(1000, dtype=bool)
        indx[:100] = True
        arr[indx] = np.ones(100, dtype=object)

        expected = np.zeros(1000)
        expected[:100] = 1
        assert_array_equal(arr, expected)

    def test_boolean_indexing_twodim(self):
        # Indexing a 2-dimensional array with
        # 2-dimensional boolean array
        a = np.array([[1, 2, 3],
                      [4, 5, 6],
                      [7, 8, 9]])
        b = np.array([[ True, False,  True],
                      [False,  True, False],
                      [ True, False,  True]])
        assert_equal(a[b], [1, 3, 5, 7, 9])
        assert_equal(a[b[1]], [[4, 5, 6]])
        assert_equal(a[b[0]], a[b[2]])

        # boolean assignment
        a[b] = 0
        assert_equal(a, [[0, 2, 0],
                         [4, 0, 6],
                         [0, 8, 0]])

    def test_boolean_indexing_list(self):
        # Regression test for #13715. It's a use-after-free bug which the
        # test won't directly catch, but it will show up in valgrind.
        a = np.array([1, 2, 3])
        b = [True, False, True]
        # Two variants of the test because the first takes a fast path
        assert_equal(a[b], [1, 3])
        assert_equal(a[None, b], [[1, 3]])

    def test_reverse_strides_and_subspace_bufferinit(self):
        # This tests that the strides are not reversed for simple and
        # subspace fancy indexing.
        a = np.ones(5)
        b = np.zeros(5, dtype=np.intp)[::-1]
        c = np.arange(5)[::-1]

        a[b] = c
        # If the strides are not reversed, the 0 in the arange comes last.
        assert_equal(a[0], 0)

        # This also tests that the subspace buffer is initialized:
        a = np.ones((5, 2))
        c = np.arange(10).reshape(5, 2)[::-1]
        a[b, :] = c
        assert_equal(a[0], [0, 1])

    def test_reversed_strides_result_allocation(self):
        # Test a bug when calculating the output strides for a result array
        # when the subspace size was 1 (and test other cases as well)
        a = np.arange(10)[:, None]
        i = np.arange(10)[::-1]
        assert_array_equal(a[i], a[i.copy('C')])

        a = np.arange(20).reshape(-1, 2)

    def test_uncontiguous_subspace_assignment(self):
        # During development there was a bug activating a skip logic
        # based on ndim instead of size.
        a = np.full((3, 4, 2), -1)
        b = np.full((3, 4, 2), -1)

        a[[0, 1]] = np.arange(2 * 4 * 2).reshape(2, 4, 2).T
        b[[0, 1]] = np.arange(2 * 4 * 2).reshape(2, 4, 2).T.copy()

        assert_equal(a, b)

    def test_too_many_fancy_indices_special_case(self):
        # Just documents behaviour, this is a small limitation.
        a = np.ones((1,) * 32)  # 32 is NPY_MAXDIMS
        assert_raises(IndexError, a.__getitem__, (np.array([0]),) * 32)

    def test_scalar_array_bool(self):
        # NumPy bools can be used as boolean index (python ones as of yet not)
        a = np.array(1)
        assert_equal(a[np.bool_(True)], a[np.array(True)])
        assert_equal(a[np.bool_(False)], a[np.array(False)])

        # After deprecating bools as integers:
        #a = np.array([0,1,2])
        #assert_equal(a[True, :], a[None, :])
        #assert_equal(a[:, True], a[:, None])
        #
        #assert_(not np.may_share_memory(a, a[True, :]))

    def test_everything_returns_views(self):
        # Before `...` would return a itself.
        a = np.arange(5)

        assert_(a is not a[()])
        assert_(a is not a[...])
        assert_(a is not a[:])

    def test_broaderrors_indexing(self):
        a = np.zeros((5, 5))
        assert_raises(IndexError, a.__getitem__, ([0, 1], [0, 1, 2]))
        assert_raises(IndexError, a.__setitem__, ([0, 1], [0, 1, 2]), 0)

    def test_trivial_fancy_out_of_bounds(self):
        a = np.zeros(5)
        ind = np.ones(20, dtype=np.intp)
        ind[-1] = 10
        assert_raises(IndexError, a.__getitem__, ind)
        assert_raises(IndexError, a.__setitem__, ind, 0)
        ind = np.ones(20, dtype=np.intp)
        ind[0] = 11
        assert_raises(IndexError, a.__getitem__, ind)
        assert_raises(IndexError, a.__setitem__, ind, 0)

    def test_trivial_fancy_not_possible(self):
        # Test that the fast path for trivial assignment is not incorrectly
        # used when the index is not contiguous or 1D, see also gh-11467.
        a = np.arange(6)
        idx = np.arange(6, dtype=np.intp).reshape(2, 1, 3)[:, :, 0]
        assert_array_equal(a[idx], idx)

        # this case must not go into the fast path, note that idx is
        # a non-contiuguous none 1D array here.