import numpy as np
from numpy import *

@test
def test_basic_indexing():
    #Single element indexing
    arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
    assert (arr[[0, 1], [1, 2]] == np.array([2, 6])).all()
    arr[[0, 2], [1, 2]] = [10, 11]
    assert arr[0, 1] == 10 and arr[2, 2] == 11
    assert arr[-1, -1] == 11
    assert (arr[0] == np.array([1, 10, 3])).all()

    #Slicing and striding
    x = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
    assert (x[1:7:2] == np.array([1, 3, 5])).all()
    assert (x[-2:10] == np.array([8, 9])).all()
    assert (x[-3:3:-1] == np.array([7, 6, 5, 4])).all()
    assert (x[5:] == np.array([5, 6, 7, 8, 9])).all()
    mx = np.array([[[1], [2], [3]], [[4], [5], [6]]])
    assert (mx[1:2] == np.array([[[4], [5], [6]]])).all()
    assert (arr[1:10:5, ::-1] == np.array([[6, 5, 4]])).all()
    selected_rows = arr[1:3]
    selected_columns = arr[:, 1]
    assert (selected_rows == np.array([[4, 5, 6], [7, 8, 11]])).all()
    assert (selected_columns == np.array([10, 5, 8])).all()

    #Dimensional indexing tools
    assert (mx[..., 0] == np.array([[1, 2, 3], [4, 5, 6]])).all()
    assert (mx[:, np.newaxis, :, :].shape == (2, 1, 3, 1))
    assert (mx[:, None, :, :].shape == (2, 1, 3, 1))
    dx = np.arange(5)
    assert (dx[:, np.newaxis] + dx[np.newaxis, :] == np.array([[0, 1, 2, 3, 4],
                                                               [1, 2, 3, 4, 5],
                                                               [2, 3, 4, 5, 6],
                                                               [3, 4, 5, 6, 7],
                                                               [4, 5, 6, 7,
                                                                8]])).all()

    arr = np.array(42)
    assert (arr[np.array(True)] == [42]).all()
    assert arr[np.array(False)].size == 0
    arr[np.array(False)] = -1
    assert arr.item() == 42
    arr[np.array(True)] = -1
    assert arr.item() == -1

test_basic_indexing()

@test
def test_advanced_indexing():
    #Integer array indexing
    x = np.arange(10, 1, -1)
    assert (x[np.array([3, 3, 1, 8])] == np.array([7, 7, 9, 2])).all()
    assert (x[np.array([3, 3, -3, 8])] == np.array([7, 7, 4, 2])).all()
    ax = np.array([[1, 2], [3, 4], [5, 6]])
    assert (ax[np.array([1, -1])] == np.array([[3, 4], [5, 6]])).all()
    try:
        ax[np.array([3, 4])]
    except IndexError as e:
        assert True
    y = np.arange(35).reshape(5, 7)
    assert (y[np.array([0, 2, 4]),
              np.array([0, 1, 2])] == np.array([0, 15, 30])).all()
    try:
        y[np.array([0, 2, 4]), np.array([0, 1])]
    except IndexError as e:
        assert True
    assert (y[np.array([0, 2, 4]), 1] == np.array([1, 15, 29])).all()
    assert (y[np.array([0, 2, 4])] == np.array([[0, 1, 2, 3, 4, 5, 6],
                                                [14, 15, 16, 17, 18, 19, 20],
                                                [28, 29, 30, 31, 32, 33,
                                                 34]])).all()
    arr = np.arange(12).reshape(4, 3)
    rows = np.array([0, 3])
    columns = np.array([0, 2])
    assert (rows[:, np.newaxis] == np.array([[0], [3]])).all()
    assert (arr[rows[:, np.newaxis], columns] == np.array([[0, 2],
                                                           [9, 11]])).all()

    #Boolean array indexing
    x = np.array([[1., 2.], [np.nan, 3.], [np.nan, np.nan]])
    assert (x[~np.isnan(x)] == np.array([1., 2., 3.])).all()
    x = np.array([1., -1., -2., 3])
    x[x < 0] += 20
    assert (x == np.array([1., 19., 18., 3.])).all()
    arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
    assert (arr[arr > 5] == np.array([6, 7, 8, 9])).all()
    mask = arr > 4
    assert (arr[mask] == np.array([5, 6, 7, 8, 9])).all()
    arr2 = np.arange(12).reshape(4, 3)
    rows = (arr2.sum(-1) % 2) == 0
    assert (rows == np.array([False, True, False, True])).all()
    a = np.arange(30).reshape(2, 3, 5)
    b = np.array([[True, True, False], [False, True, True]])
    assert (a[b] == np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9],
                              [20, 21, 22, 23, 24], [25, 26, 27, 28,
                                                     29]])).all()

test_advanced_indexing()

@test
def test_combined_indexing():
    x = np.arange(35).reshape(5, 7)
    b = x > 20
    assert (x[b[:, 5], 1:3] == np.array([[22, 23], [29, 30]])).all()
    assert (np.array([42, 32])[()] == np.array([42, 32])).all()
    assert np.array(42)[()] == 42

test_combined_indexing()