def test_columnwise_iteration_with_structured_array_and_mask():
    input = np.arange(12).reshape((4, 3))
    m = LazyArray(input, shape=(4, 3))  # 4 rows, 3 columns
    mask = np.array([False, True, True])
    cols = [col for col in m.by_column(mask=mask)]
    assert_array_equal(cols[0], input[:, 1])
    assert_array_equal(cols[1], input[:, 2])

def test_columnwise_iteration_with_function():
    input = lambda i, j: 2 * i + j
    m = LazyArray(input, shape=(4, 3))
    cols = [col for col in m.by_column()]
    assert_array_equal(cols[0], np.array([0, 2, 4, 6]))
    assert_array_equal(cols[1], np.array([1, 3, 5, 7]))
    assert_array_equal(cols[2], np.array([2, 4, 6, 8]))

def test_multiple_operations_with_structured_array():
    input = np.arange(12).reshape((4, 3))
    m0 = LazyArray(input, shape=(4, 3))
    m1 = (m0 + 2) < 5
    m2 = (m0 < 5) + 2
    assert_array_equal(m1.evaluate(simplify=True), (input + 2) < 5)
    assert_array_equal(m2.evaluate(simplify=True), (input < 5) + 2)
    assert_array_equal(m0.evaluate(simplify=True), input)

def test_evaluate_with_functional_array():
    input = lambda i,j: 2*i + j
    m = LazyArray(input, shape=(4,3))
    assert_array_equal(m.evaluate(),
                        np.array([[0, 1, 2],
                                     [2, 3, 4],
                                     [4, 5, 6],
                                     [6, 7, 8]]))

def test_columnwise_iteration_with_random_array_parallel_safe_no_mask():
    random.mpi_rank=0; random.num_processes=2
    input = random.RandomDistribution(rng=MockRNG(parallel_safe=True))
    copy_input = random.RandomDistribution(rng=MockRNG(parallel_safe=True))
    m = LazyArray(input, shape=(4,3))
    cols = [col for col in m.by_column()]
    assert_array_equal(cols[0], copy_input.next(4, mask_local=False))
    assert_array_equal(cols[1], copy_input.next(4, mask_local=False))
    assert_array_equal(cols[2], copy_input.next(4, mask_local=False))

def test_columnwise_iteration_with_random_array_parallel_safe_with_mask():
    random.mpi_rank=0; random.num_processes=2
    input = random.RandomDistribution(rng=MockRNG(parallel_safe=True))
    copy_input = random.RandomDistribution(rng=MockRNG(parallel_safe=True))
    m = LazyArray(input, shape=(4,3))
    mask = np.array([False, False, True])
    cols = [col for col in m.by_column(mask=mask)]
    assert_equal(len(cols), 1)
    assert_array_almost_equal(cols[0], copy_input.next(12, mask_local=False)[8:], 15)

def test_columnwise_iteration_with_random_array_parallel_safe_with_mask():
    orig_get_mpi_config = random.get_mpi_config
    random.get_mpi_config = lambda: (0, 2)
    input = random.RandomDistribution("uniform", (0, 1), rng=MockRNG(parallel_safe=True))
    copy_input = random.RandomDistribution("gamma", (2, 3), rng=MockRNG(parallel_safe=True))
    m = LazyArray(input, shape=(4, 3))
    mask = np.array([False, False, True])
    cols = [col for col in m.by_column(mask=mask)]
    assert_equal(len(cols), 1)
    assert_array_almost_equal(cols[0], copy_input.next(12, mask_local=False)[8:], 15)
    random.get_mpi_config = orig_get_mpi_config

def test_columnwise_iteration_with_random_array_parallel_safe_no_mask():
    orig_get_mpi_config = random.get_mpi_config
    random.get_mpi_config = lambda: (0, 2)
    input = random.RandomDistribution("uniform", (0, 1), rng=MockRNG(parallel_safe=True))
    copy_input = random.RandomDistribution("normal", (0, 1), rng=MockRNG(parallel_safe=True))
    m = LazyArray(input, shape=(4, 3))
    cols = [col for col in m.by_column()]
    assert_array_equal(cols[0], copy_input.next(4, mask_local=False))
    assert_array_equal(cols[1], copy_input.next(4, mask_local=False))
    assert_array_equal(cols[2], copy_input.next(4, mask_local=False))
    random.get_mpi_config = orig_get_mpi_config

def test_columnwise_iteration_with_random_array_parallel_safe_no_mask():
    orig_get_mpi_config = random.get_mpi_config

    # first, with a single MPI node
    random.get_mpi_config = lambda: (0, 2)
    input = random.RandomDistribution('uniform', (0, 1), rng=MockRNG(parallel_safe=True))
    m = LazyArray(input, shape=(4, 3))
    cols_np1 = [col for col in m.by_column()]

    # now, on one node of two
    random.get_mpi_config = lambda: (1, 2)
    input = random.RandomDistribution('uniform', (0, 1), rng=MockRNG(parallel_safe=True))
    m = LazyArray(input, shape=(4, 3))
    cols_np2_1 = [col for col in m.by_column()]

    for i in range(3):
        assert_array_equal(cols_np1[i], cols_np2_1[i])

    random.get_mpi_config = orig_get_mpi_config

def test_columnwise_iteration_with_random_array_parallel_safe_with_mask():
    orig_get_mpi_config = random.get_mpi_config

    mask = np.array([False, False, True])

    # first, with a single MPI node
    random.get_mpi_config = lambda: (0, 2)
    input = random.RandomDistribution('uniform', (0, 1), rng=MockRNG(parallel_safe=True))
    m = LazyArray(input, shape=(4, 3))
    cols_np1 = [col for col in m.by_column(mask=mask)]

    # now, on one node of two
    random.get_mpi_config = lambda: (0, 2)
    input = random.RandomDistribution('uniform', (0, 1), rng=MockRNG(parallel_safe=True))
    m = LazyArray(input, shape=(4, 3))
    cols_np2_0 = [col for col in m.by_column(mask=mask)]

    # now, on the other node of two
    random.get_mpi_config = lambda: (1, 2)
    input = random.RandomDistribution('uniform', (0, 1), rng=MockRNG(parallel_safe=True))
    m = LazyArray(input, shape=(4, 3))
    cols_np2_1 = [col for col in m.by_column(mask=mask)]

    assert_equal(len(cols_np1), 1)
    assert_equal(len(cols_np2_0), 1)
    assert_equal(len(cols_np2_1), 1)
    assert_array_equal(cols_np1[0], cols_np2_0[0])

    random.get_mpi_config = orig_get_mpi_config

def test_columnwise_iteration_with_structured_array():
    input = np.arange(12).reshape((4, 3))
    m = LazyArray(input, shape=(4, 3))  # 4 rows, 3 columns
    cols = [col for col in m.by_column()]
    assert_array_equal(cols[0], input[:, 0])
    assert_array_equal(cols[2], input[:, 2])

def test_evaluate_with_structured_array():
    input = np.arange(12).reshape((4, 3))
    m = LazyArray(input, shape=(4, 3))
    assert_array_equal(m.evaluate(), input)

def test_setitem_nonexpanded_different_value():
    A = LazyArray(3, shape=(5,))
    assert A.evaluate(simplify=True) == 3
    A[0] = 4
    A[4] = 5
    assert_array_equal(A.evaluate(simplify=True), np.array([4, 3, 3, 3, 5]))

def test_columnwise_iteration_with_flat_array():
    m = LazyArray(5, shape=(4, 3))  # 4 rows, 3 columns
    cols = [col for col in m.by_column()]
    assert_equal(cols, [5, 5, 5])

def test_create_with_array():
    A = LazyArray(np.array([1, 2, 3]), shape=(3,))
    assert A.shape == (3,)
    assert_array_equal(A.evaluate(simplify=True), np.array([1, 2, 3]))

def test_setitem_nonexpanded_same_value():
    A = LazyArray(3, shape=(5,))
    assert A.evaluate(simplify=True) == 3
    A[0] = 3
    assert A.evaluate(simplify=True) == 3

def test_create_with_float():
    A = LazyArray(3.0, shape=(5,))
    assert A.shape == (5,)
    assert A.evaluate(simplify=True) == 3.0

def test_columnwise_iteration_with_flat_array_and_mask():
    m = LazyArray(5, shape=(4, 3))  # 4 rows, 3 columns
    mask = np.array([True, False, True])
    cols = [col for col in m.by_column(mask=mask)]
    assert_equal(cols, [5, 5])

def test_lt_with_flat_array():
    m0 = LazyArray(5, shape=(4, 3))
    m1 = m0 < 10
    assert_equal(m1.evaluate(simplify=True), True)
    assert_equal(m0.evaluate(simplify=True), 5)

def test_iadd_with_flat_array():
    m = LazyArray(5, shape=(4, 3))
    m += 2
    assert_array_equal(m.evaluate(), 7 * np.ones((4, 3)))
    assert_equal(m.base_value, 5)
    assert_equal(m.evaluate(simplify=True), 7)