import numpy as np
from numpy.testing import assert_array_equal

from landlab import NodeStatus, RasterModelGrid

CB = NodeStatus.CLOSED
FV = NodeStatus.FIXED_VALUE


def test_get_status():
    """Test getting the node status array."""
    grid = RasterModelGrid((4, 5))
    assert_array_equal(
        grid.status_at_node,
        [FV, FV, FV, FV, FV, FV, 0, 0, 0, FV, FV, 0, 0, 0, FV, FV, FV, FV, FV, FV],
    )


def test_set_status_with_scalar():
    """Test setting status with a scalar."""
    grid = RasterModelGrid((4, 5))
    grid.status_at_node[6] = 2

    assert_array_equal(
        grid.status_at_node,
        [FV, FV, FV, FV, FV, FV, 2, 0, 0, FV, FV, 0, 0, 0, FV, FV, FV, FV, FV, FV],
    )


def test_set_status_with_slice():
    """Test setting status with an array slice."""
    grid = RasterModelGrid((4, 5))
    grid.status_at_node[6:9] = 2

    assert_array_equal(
        grid.status_at_node,
        [FV, FV, FV, FV, FV, FV, 2, 2, 2, FV, FV, 0, 0, 0, FV, FV, FV, FV, FV, FV],
    )


def test_set_status_with_array_bool():
    """Test setting node status with boolean array."""
    grid = RasterModelGrid((4, 5))
    inds = np.full((20,), False, dtype=bool)
    inds[6] = True
    inds[7] = True
    inds[13] = True

    grid.status_at_node[inds] = 2
    assert_array_equal(
        grid.status_at_node,
        [FV, FV, FV, FV, FV, FV, 2, 2, 0, FV, FV, 0, 0, 2, FV, FV, FV, FV, FV, FV],
    )


def test_set_with_itemset():
    grid = RasterModelGrid((4, 5))
    grid.status_at_node.itemset(7, 2)

    assert_array_equal(
        grid.status_at_node,
        [FV, FV, FV, FV, FV, FV, 0, 2, 0, FV, FV, 0, 0, 0, FV, FV, FV, FV, FV, FV],
    )


def test_set_status_with_array():
    """Test that active links are reset after changing the node status."""
    grid = RasterModelGrid((4, 5))

    assert_array_equal(
        grid.active_links,
        [5, 6, 7, 9, 10, 11, 12, 14, 15, 16, 18, 19, 20, 21, 23, 24, 25],
    )

    grid.status_at_node[:5] = CB
    assert_array_equal(
        grid.status_at_node,
        [CB, CB, CB, CB, CB, FV, 0, 0, 0, FV, FV, 0, 0, 0, FV, FV, FV, FV, FV, FV],
    )

    assert_array_equal(
        grid.active_links, [9, 10, 11, 12, 14, 15, 16, 18, 19, 20, 21, 23, 24, 25]
    )
