import pytest
from aarecommon.math.raster_grid import grid_to_image_id, image_id_to_grid


def test_grid_to_image_id_single_cell():
    assert grid_to_image_id(grid_x=1, grid_y=1, number_of_cols=1) == 0


def test_grid_to_image_id_first_row_left_to_right():
    assert grid_to_image_id(grid_x=1, grid_y=1, number_of_cols=5) == 0
    assert grid_to_image_id(grid_x=2, grid_y=1, number_of_cols=5) == 1
    assert grid_to_image_id(grid_x=3, grid_y=1, number_of_cols=5) == 2
    assert grid_to_image_id(grid_x=4, grid_y=1, number_of_cols=5) == 3
    assert grid_to_image_id(grid_x=5, grid_y=1, number_of_cols=5) == 4


def test_grid_to_image_id_second_row_right_to_left():
    assert grid_to_image_id(grid_x=5, grid_y=2, number_of_cols=5) == 5
    assert grid_to_image_id(grid_x=4, grid_y=2, number_of_cols=5) == 6
    assert grid_to_image_id(grid_x=3, grid_y=2, number_of_cols=5) == 7
    assert grid_to_image_id(grid_x=2, grid_y=2, number_of_cols=5) == 8
    assert grid_to_image_id(grid_x=1, grid_y=2, number_of_cols=5) == 9


def test_image_id_to_grid_first_row_left_to_right():
    assert image_id_to_grid(0, 5) == (1, 1)
    assert image_id_to_grid(1, 5) == (2, 1)
    assert image_id_to_grid(2, 5) == (3, 1)
    assert image_id_to_grid(3, 5) == (4, 1)
    assert image_id_to_grid(4, 5) == (5, 1)


def test_image_id_to_grid_second_row_right_to_left():
    assert image_id_to_grid(5, 5) == (5, 2)
    assert image_id_to_grid(6, 5) == (4, 2)
    assert image_id_to_grid(7, 5) == (3, 2)
    assert image_id_to_grid(8, 5) == (2, 2)
    assert image_id_to_grid(9, 5) == (1, 2)


def test_round_trip_conversion_for_multiple_grid_shapes():
    for number_of_cols in (1, 2, 3, 5, 8):
        for grid_y in range(1, 7):
            for grid_x in range(1, number_of_cols + 1):
                image_id = grid_to_image_id(
                    grid_x=grid_x, grid_y=grid_y, number_of_cols=number_of_cols
                )
                assert image_id_to_grid(image_id, number_of_cols) == (grid_x, grid_y)


def test_known_3x3_serpentine_mapping():
    expected = {
        (1, 1): 0,
        (2, 1): 1,
        (3, 1): 2,
        (3, 2): 3,
        (2, 2): 4,
        (1, 2): 5,
        (1, 3): 6,
        (2, 3): 7,
        (3, 3): 8,
    }
    for grid_pos, image_id in expected.items():
        assert grid_to_image_id(grid_pos[0], grid_pos[1], 3) == image_id
        assert image_id_to_grid(image_id, 3) == grid_pos


def test_grid_to_image_id_rejects_zero_columns():
    with pytest.raises(ValueError, match="number_of_cols must be >= 1"):
        grid_to_image_id(grid_x=1, grid_y=1, number_of_cols=0)


def test_grid_to_image_id_rejects_zero_grid_x():
    with pytest.raises(ValueError, match="grid_x must be >= 1"):
        grid_to_image_id(grid_x=0, grid_y=1, number_of_cols=5)


def test_grid_to_image_id_rejects_zero_grid_y():
    with pytest.raises(ValueError, match="grid_y must be >= 1"):
        grid_to_image_id(grid_x=1, grid_y=0, number_of_cols=5)


def test_grid_to_image_id_rejects_grid_x_larger_than_number_of_cols():
    with pytest.raises(ValueError, match="grid_x cannot be greater than number_of_cols"):
        grid_to_image_id(grid_x=6, grid_y=1, number_of_cols=5)


def test_image_id_to_grid_rejects_zero_columns():
    with pytest.raises(ValueError, match="number_of_cols must be >= 1"):
        image_id_to_grid(image_id=0, number_of_cols=0)


def test_image_id_to_grid_rejects_negative_image_id():
    with pytest.raises(ValueError, match="image_id must be >= 0"):
        image_id_to_grid(image_id=-1, number_of_cols=5)