import pytest
import numpy as np
from numpy.testing import assert_array_equal

from landlab import RasterModelGrid
from landlab.grid.base import BAD_INDEX_VALUE as X


def test_all_active_neighbors():
    rmg = RasterModelGrid(5, 4)
    expected = np.array([
        [X, X, X, X], [X, 5, X, X], [X, 6, X, X], [X, X, X, X],
        [5, X, X, X], [6, 9, 4, 1], [7, 10, 5, 2], [X, X, 6, X],
        [9, X, X, X], [10, 13, 8, 5], [11, 14, 9, 6], [X, X, 10, X],
        [13, X, X, X], [14, 17, 12, 9], [15, 18, 13, 10], [X, X, 14, X],
        [X, X, X, X], [X, X, X, 13], [X, X, X, 14], [X, X, X, X],
    ])
    assert_array_equal(rmg.active_adjacent_nodes_at_node, expected)


def test_all_neighbors():
    rmg = RasterModelGrid((5, 4))
    expected = np.array([
        [ 1,  4, X,  X], [ 2,  5,  0,  X], [ 3,  6,  1,  X], [X,  7,  2,  X],
        [ 5,  8, X,  0], [ 6,  9,  4,  1], [ 7, 10,  5,  2], [X, 11,  6,  3],
        [ 9, 12, X,  4], [10, 13,  8,  5], [11, 14,  9,  6], [X, 15, 10,  7],
        [13, 16, X,  8], [14, 17, 12,  9], [15, 18, 13, 10], [X, 19, 14, 11],
        [17,  X, X, 12], [18,  X, 16, 13], [19,  X, 17, 14], [X,  X, 18, 15],
    ])
    assert_array_equal(rmg.adjacent_nodes_at_node, expected)


def test_active_neighbor_list_with_scalar_arg():
    rmg = RasterModelGrid(5, 4)

    assert_array_equal(rmg.active_adjacent_nodes_at_node[6],
                       np.array([7, 10, 5, 2]))
    assert_array_equal(rmg.active_adjacent_nodes_at_node[-1],
                       np.array([X, X, X, X]))
    assert_array_equal(rmg.active_adjacent_nodes_at_node[-2],
                       np.array([X, X, X, 14]))


def test_neighbor_list_with_scalar_arg():
    rmg = RasterModelGrid((5, 4))

    assert_array_equal(rmg.adjacent_nodes_at_node[6], np.array([7, 10, 5, 2]))
    assert_array_equal(rmg.adjacent_nodes_at_node[-1], np.array([X, X, 18, 15]))
    assert_array_equal(rmg.adjacent_nodes_at_node[-2], np.array([19, X, 17, 14]))


def test_active_neighbor_list_with_array_arg():
    rmg = RasterModelGrid(5, 4)
    assert_array_equal(rmg.active_adjacent_nodes_at_node[[6, -1]],
                       np.array([[7, 10, 5, 2], [X, X, X, X]]))


def test_neighbor_list_with_array_arg():
    rmg = RasterModelGrid(5, 4)
    assert_array_equal(rmg.adjacent_nodes_at_node[(6, -1), :],
                       np.array([[7, 10, 5, 2], [X, X, 18, 15]]))


def test_neighbor_list_is_read_only():
    rmg = RasterModelGrid(5, 4)
    with pytest.raises(ValueError):
        rmg.adjacent_nodes_at_node[0] = [1, 2, 3, 4]


def test_neighbors_is_contiguous():
    rmg = RasterModelGrid(5, 4)
    assert rmg.adjacent_nodes_at_node.flags['C_CONTIGUOUS']


def test_active_neighbor_list_boundary():
    """All of the neighbor IDs for a boundary cell are -1."""
    rmg = RasterModelGrid(5, 4)
    import landlab.utils.structured_grid as sgrid
    rmg.set_closed_nodes([0, 1, 2, 3, 4, 7, 8, 11, 12, 15, 16, 17, 18, 19])

    for node_id in sgrid.boundary_iter(rmg.shape):
        assert_array_equal(rmg.active_adjacent_nodes_at_node[node_id],
                           np.array([X, X, X, X]))


def test_all_diagonals():
    rmg = RasterModelGrid((5, 4))
    expected = np.array([
        [ 5,  X, X,  X], [ 6,  4,  X,  X], [ 7,  5,  X,  X], [X,  6,  X,  X],
        [ 9,  X, X,  1], [10,  8,  0,  2], [11,  9,  1,  3], [X, 10,  2,  X],
        [13,  X, X,  5], [14, 12,  4,  6], [15, 13,  5,  7], [X, 14,  6,  X],
        [17,  X, X,  9], [18, 16,  8, 10], [19, 17,  9, 11], [X, 18, 10,  X],
        [ X,  X, X, 13], [ X,  X, 12, 14], [ X,  X, 13, 15], [X,  X, 14,  X],
    ])
    assert_array_equal(rmg.diagonal_adjacent_nodes_at_node, expected)


def test_diagonal_list_with_scalar_arg():
    rmg = RasterModelGrid((5, 4))

    assert_array_equal(rmg.diagonal_adjacent_nodes_at_node[6], np.array([11, 9, 1, 3]))
    assert_array_equal(rmg.diagonal_adjacent_nodes_at_node[-1], np.array([X, X, 14, X]))
    assert_array_equal(rmg.diagonal_adjacent_nodes_at_node[-2], np.array([X, X, 13, 15]))


def test_diagonal_list_with_array_arg():
    rmg = RasterModelGrid(5, 4)
    assert_array_equal(rmg.diagonal_adjacent_nodes_at_node[(6, -1), :],
                       np.array([[11, 9, 1, 3], [X, X, 14, X]]))


def test_diagonal_list_is_read_only():
    rmg = RasterModelGrid(5, 4)
    with pytest.raises(ValueError):
        rmg.diagonal_adjacent_nodes_at_node[0] = [1, 2, 3, 4]


def test_diagonals_is_contiguous():
    rmg = RasterModelGrid(5, 4)
    assert rmg.diagonal_adjacent_nodes_at_node.flags['C_CONTIGUOUS']
