from copy import deepcopy
import math
from unittest import mock

from affine import Affine
import numpy as np
import pytest
from numpy.testing import assert_allclose

import rasterio
from rasterio.enums import MergeAlg
from rasterio.errors import WindowError, ShapeSkipWarning
from rasterio.features import (
    bounds, geometry_mask, geometry_window, is_valid_geom, rasterize, sieve,
    shapes)

from .conftest import MockGeoInterface, requires_gdal3_11, requires_gdal_lt_3_11, requires_gdal3_12_1, requires_gdal_lt_3_12_1

DEFAULT_SHAPE = (10, 10)


def test_bounds_point():
    g = {'type': 'Point', 'coordinates': [10, 10]}
    assert bounds(g) == (10, 10, 10, 10)
    assert bounds(MockGeoInterface(g)) == (10, 10, 10, 10)


def test_bounds_line():
    g = {'type': 'LineString', 'coordinates': [[0, 0], [10, 10]]}
    assert bounds(g) == (0, 0, 10, 10)
    assert bounds(MockGeoInterface(g)) == (0, 0, 10, 10)


def test_bounds_ring():
    g = {'type': 'LinearRing', 'coordinates': [[0, 0], [10, 10], [10, 0]]}
    assert bounds(g) == (0, 0, 10, 10)
    assert bounds(MockGeoInterface(g)) == (0, 0, 10, 10)


def test_bounds_polygon():
    g = {'type': 'Polygon', 'coordinates': [[[0, 0], [10, 10], [10, 0]]]}
    assert bounds(g) == (0, 0, 10, 10)
    assert bounds(MockGeoInterface(g)) == (0, 0, 10, 10)


def test_bounds_z():
    g = {'type': 'Point', 'coordinates': [10, 10, 10]}
    assert bounds(g) == (10, 10, 10, 10)
    assert bounds(MockGeoInterface(g)) == (10, 10, 10, 10)


@pytest.mark.parametrize('geometry', [
    {'type': 'Polygon'},
    {'type': 'Polygon', 'not_coordinates': []},
    {'type': 'bogus', 'not_coordinates': []},
    {
        'type': 'GeometryCollection',
        'geometries': [
            {'type': 'Point', 'coordinates': [1, 1]},
            {'type': 'LineString', 'not_coordinates': [[-10, -20], [10, 20]]},
        ]
    }
])
def test_bounds_invalid_obj(geometry):
    with pytest.raises(ValueError, match="geometry must be a GeoJSON-like geometry, GeometryCollection, or FeatureCollection"):
        bounds(geometry)


def test_bounds_feature_collection(basic_featurecollection):
    fc = basic_featurecollection
    assert bounds(fc) == bounds(fc['features'][0]) == (2, 2, 4.25, 4.25)


def test_bounds_geometry_collection():
    gc = {
        'type': 'GeometryCollection',
        'geometries': [
            {'type': 'Point', 'coordinates': [1, 1]},
            {'type': 'LineString', 'coordinates': [[-10, -20], [10, 20]]},
            {'type': 'Polygon', 'coordinates': [[[5, 5], [25, 50], [25, 5]]]}
        ]
    }

    assert bounds(gc) == (-10, -20, 25, 50)
    assert bounds(MockGeoInterface(gc)) == (-10, -20, 25, 50)


def test_bounds_existing_bbox(basic_featurecollection):
    """Test with existing bbox in geojson.

    Similar to that produced by rasterio.  Values specifically modified here
    for testing, bboxes are not valid as written.
    """
    fc = basic_featurecollection
    fc['bbox'] = [0, 10, 10, 20]
    fc['features'][0]['bbox'] = [0, 100, 10, 200]

    assert bounds(fc['features'][0]) == (0, 100, 10, 200)
    assert bounds(fc) == (0, 10, 10, 20)


def test_geometry_mask(basic_geometry, basic_image_2x2):
    assert np.array_equal(
        basic_image_2x2 == 0,
        geometry_mask(
            [basic_geometry],
            out_shape=DEFAULT_SHAPE,
            transform=Affine.identity()
        )
    )


def test_geometry_mask_invert(basic_geometry, basic_image_2x2):
    assert np.array_equal(
        basic_image_2x2,
        geometry_mask(
            [basic_geometry],
            out_shape=DEFAULT_SHAPE,
            transform=Affine.identity(),
            invert=True
        )
    )


@pytest.mark.parametrize(
    "geom",
    [{"type": "Invalid"}, {"type": "Point"}, {"type": "Point", "coordinates": []}],
)
def test_geometry_invalid_geom(geom):
    """An invalid geometry should fail"""
    with pytest.warns(ShapeSkipWarning):
        mask = geometry_mask(
            [geom],
            out_shape=DEFAULT_SHAPE,
            transform=Affine.identity())

    assert mask.shape == DEFAULT_SHAPE
    assert np.all(mask)


def test_geometry_mask_invalid_shape(basic_geometry):
    """A width==0 or height==0 should fail with ValueError"""

    for shape in [(0, 0), (1, 0), (0, 1)]:
        with pytest.raises(ValueError) as exc_info:
            geometry_mask(
                [basic_geometry],
                out_shape=shape,
                transform=Affine.identity())

        assert 'must be > 0' in exc_info.value.args[0]


def test_geometry_mask_no_transform(basic_geometry):
    with pytest.raises(TypeError):
        geometry_mask(
            [basic_geometry],
            out_shape=DEFAULT_SHAPE,
            transform=None)


def test_geometry_window_no_pad(basic_image_file, basic_geometry):
    with rasterio.open(basic_image_file) as src:
        window = geometry_window(src, [basic_geometry, basic_geometry])
        assert window.flatten() == (2, 2, 3, 3)


def test_geometry_window_geo_interface(basic_image_file, basic_geometry):
    with rasterio.open(basic_image_file) as src:
        window = geometry_window(src, [MockGeoInterface(basic_geometry)])
        assert window.flatten() == (2, 2, 3, 3)


def test_geometry_window_pixel_precision(basic_image_file):
    """Window offsets should be floor, width and height ceiling"""

    geom2 = {
        'type': 'Polygon',
        'coordinates': [[
            (1.99999, 2),
            (1.99999, 4.0001), (4.0001, 4.0001), (4.0001, 2),
            (1.99999, 2)
        ]]
    }

    with rasterio.open(basic_image_file) as src:
        window = geometry_window(src, [geom2], pixel_precision=6)
        assert window.flatten() == (1, 2, 4, 3)


def test_geometry_window_north_up(path_rgb_byte_tif):
    geometry = {
        'type': 'Polygon',
        'coordinates': [[
            (200000, 2700000),
            (200000, 2750000),
            (250000, 2750000),
            (250000, 2700000),
            (200000, 2700000)
        ]]
    }

    with rasterio.open(path_rgb_byte_tif) as src:
        window = geometry_window(src, [geometry])
    assert window.flatten() == (326, 256, 168, 167)


def test_geometry_window_rotated_boundless():
    """Get the right boundless window for a rotated dataset"""
    sqrt2 = math.sqrt(2.0)
    dataset = mock.MagicMock()
    dataset.transform = (
        Affine.rotation(-45.0)
        * Affine.translation(-sqrt2, sqrt2)
        * Affine.scale(sqrt2 / 2.0, -sqrt2 / 2.0)
    )
    dataset.height = 4.0
    dataset.width = 4.0

    geometry = {
        "type": "Polygon",
        "coordinates": [
            [(-2.0, -2.0), (-2.0, 2.0), (2.0, 2.0), (2.0, -2.0), (-2.0, -2.0),]
        ],
    }

    win = geometry_window(dataset, [geometry, geometry], boundless=True)
    assert win.col_off == pytest.approx(-2.0)
    assert win.row_off == pytest.approx(-2.0)
    assert win.width == pytest.approx(2.0 * dataset.width)
    assert win.height == pytest.approx(2.0 * dataset.height)


def test_geometry_window_pad(basic_image_file, basic_geometry):
    # Note: this dataset's geotransform is not a geographic one.
    # x increases with col, but y also increases with row.
    # It's flipped, not rotated like a south-up world map.
    with rasterio.open(basic_image_file) as src:
        transform = src.transform
        dataset = mock.MagicMock()
        dataset.res = src.res
        dataset.transform = src.transform
        dataset.height = src.height
        dataset.width = src.width
        window = geometry_window(dataset, [basic_geometry], pad_x=0.5, pad_y=0.5)

    assert window.flatten() == (1, 1, 4, 4)


def test_geometry_window_large_shapes(basic_image_file):
    geometry = {
        'type': 'Polygon',
        'coordinates': [[
            (-2000, -2000),
            (-2000, 2000),
            (2000, 2000),
            (2000, -2000),
            (-2000, -2000)
        ]]
    }

    with rasterio.open(basic_image_file) as src:
        window = geometry_window(src, [geometry])
        assert window.flatten() == (0, 0, src.height, src.width)


def test_geometry_window_no_overlap(path_rgb_byte_tif, basic_geometry):
    """Geometries that do not overlap raster raises WindowError"""

    with rasterio.open(path_rgb_byte_tif) as src:
        with pytest.raises(WindowError):
            geometry_window(src, [basic_geometry], north_up=False)


def test_is_valid_geo_interface(geojson_point):
    """Properly formed Point object with geo interface is valid"""
    assert is_valid_geom(MockGeoInterface(geojson_point))


def test_is_valid_geom_point(geojson_point):
    """Properly formed GeoJSON Point is valid"""
    assert is_valid_geom(geojson_point)

    # Empty coordinates are invalid
    geojson_point['coordinates'] = []
    assert not is_valid_geom(geojson_point)


def test_is_valid_geom_multipoint(geojson_multipoint):
    """Properly formed GeoJSON MultiPoint is valid"""
    assert is_valid_geom(geojson_multipoint)

    # Empty iterable is invalid
    geom = deepcopy(geojson_multipoint)
    geom['coordinates'] = []
    assert not is_valid_geom(geom)

    # Empty first coordinate is invalid
    geom = deepcopy(geojson_multipoint)
    geom['coordinates'] = [[]]


def test_is_valid_geom_line(geojson_line):
    """Properly formed GeoJSON LineString is valid"""

    assert is_valid_geom(geojson_line)

    # Empty iterable is invalid
    geom = deepcopy(geojson_line)
    geom['coordinates'] = []
    assert not is_valid_geom(geom)

    # Empty first coordinate is invalid
    geom = deepcopy(geojson_line)
    geom['coordinates'] = [[]]


def test_is_valid_geom_multiline(geojson_line):
    """Properly formed GeoJSON MultiLineString is valid"""

    assert is_valid_geom(geojson_line)

    # Empty iterables are invalid
    geom = deepcopy(geojson_line)
    geom['coordinates'] = []
    assert not is_valid_geom(geom)

    geom = deepcopy(geojson_line)
    geom['coordinates'] = [[]]
    assert not is_valid_geom(geom)

    # Empty first coordinate is invalid
    geom = deepcopy(geojson_line)
    geom['coordinates'] = [[[]]]
    assert not is_valid_geom(geom)


def test_is_valid_geom_polygon(geojson_polygon):
    """Properly formed GeoJSON Polygon is valid"""

    assert is_valid_geom(geojson_polygon)

    # Empty iterables are invalid
    geom = deepcopy(geojson_polygon)
    geom['coordinates'] = []
    assert not is_valid_geom(geom)

    geom = deepcopy(geojson_polygon)
    geom['coordinates'] = [[]]
    assert not is_valid_geom(geom)

    # Empty first coordinate is invalid
    geom = deepcopy(geojson_polygon)
    geom['coordinates'] = [[[]]]
    assert not is_valid_geom(geom)


def test_is_valid_geom_ring(geojson_polygon):
    """Properly formed GeoJSON LinearRing is valid"""
    geojson_ring = deepcopy(geojson_polygon)
    geojson_ring['type'] = 'LinearRing'
    # take first ring from polygon as sample
    geojson_ring['coordinates'] = geojson_ring['coordinates'][0]
    assert is_valid_geom(geojson_ring)

    # Empty iterables are invalid
    geom = deepcopy(geojson_ring)
    geom['coordinates'] = []
    assert not is_valid_geom(geom)

    geom = deepcopy(geojson_ring)
    geom['coordinates'] = [[]]
    assert not is_valid_geom(geom)


def test_is_valid_geom_multipolygon(geojson_multipolygon):
    """Properly formed GeoJSON MultiPolygon is valid"""

    assert is_valid_geom(geojson_multipolygon)

    # Empty iterables are invalid
    geom = deepcopy(geojson_multipolygon)
    geom['coordinates'] = []
    assert not is_valid_geom(geom)

    geom = deepcopy(geojson_multipolygon)
    geom['coordinates'] = [[]]
    assert not is_valid_geom(geom)

    geom = deepcopy(geojson_multipolygon)
    geom['coordinates'] = [[[]]]
    assert not is_valid_geom(geom)

    # Empty first coordinate is invalid
    geom = deepcopy(geojson_multipolygon)
    geom['coordinates'] = [[[[]]]]
    assert not is_valid_geom(geom)


def test_is_valid_geom_geomcollection(geojson_geomcollection):
    """Properly formed GeoJSON GeometryCollection is valid"""

    assert is_valid_geom(geojson_geomcollection)

    # Empty GeometryCollection is invalid
    geom = deepcopy(geojson_geomcollection)
    geom['geometries'] = []
    assert not is_valid_geom(geom)


@pytest.mark.parametrize("geom", [None, 1, "foo", "type", ["type"], {"type": "Invalid"}, {"type": "Point"}])
def test_is_valid_geom_invalid_inputs(geom):
    """Improperly formed GeoJSON objects should fail"""
    assert not is_valid_geom(geom)


def test_rasterize_point(geojson_point):
    expected = np.zeros(shape=DEFAULT_SHAPE, dtype='uint8')
    expected[2, 2] = 1

    assert np.array_equal(
        rasterize([geojson_point], out_shape=DEFAULT_SHAPE),
        expected
    )


def test_rasterize_to_dataset(tmp_path, geojson_point):
    expected = np.zeros(shape=DEFAULT_SHAPE, dtype="uint8")
    expected[2, 2] = 1

    with rasterio.open(
        tmp_path / "test.tif",
        "w",
        driver="GTiff",
        count=1,
        width=10,
        height=10,
        dtype="uint8",
        crs="EPSG:4326",
        transform=Affine.translation(0.0, 0.0) * Affine.scale(1.0, 1.0),
    ) as dst:
        rasterize([geojson_point], dst_path=dst)

    with rasterio.open(tmp_path / "test.tif") as dst:
        result = dst.read(indexes=1)

    assert np.array_equal(result, expected)


def test_rasterize_to_file(tmp_path, geojson_point):
    expected = np.zeros(shape=DEFAULT_SHAPE, dtype="uint8")
    expected[2, 2] = 1
    expected = np.ma.masked_equal(expected, 0)

    dst_kwds = dict(
        driver="GTiff",
        count=1,
        width=10,
        height=10,
        dtype="uint8",
        crs="EPSG:4326",
        transform=Affine.translation(0.0, 0.0) * Affine.scale(1.0, 1.0),
    )
    rasterize(
        [geojson_point], nodata=0.0, dst_path=(tmp_path / "test.tif"), dst_kwds=dst_kwds
    )

    with rasterio.open(tmp_path / "test.tif") as dst:
        result = dst.read(indexes=1, masked=True)

    assert np.array_equal(result, expected)


def test_rasterize_point_dtype_int(geojson_point):
    """Demonstrate fix of #3043."""
    expected = np.zeros(shape=DEFAULT_SHAPE, dtype=int)
    expected[2, 2] = 1

    assert np.array_equal(
        rasterize([geojson_point], out_shape=DEFAULT_SHAPE),
        expected
    )


def test_rasterize_multipoint(geojson_multipoint):
    expected = np.zeros(shape=DEFAULT_SHAPE, dtype='uint8')
    expected[2, 2] = 1
    expected[4, 4] = 1

    assert np.array_equal(
        rasterize([geojson_multipoint], out_shape=DEFAULT_SHAPE), expected
    )


def test_rasterize_multipoint_masked(geojson_multipoint):
    expected = np.zeros(shape=DEFAULT_SHAPE, dtype="uint8")
    expected[2, 2] = 1
    expected[4, 4] = 1
    expected = np.ma.masked_equal(expected, 0)

    assert np.array_equal(
        rasterize([geojson_multipoint], out_shape=DEFAULT_SHAPE, masked=True), expected
    )


def test_rasterize_line(geojson_line):
    expected = np.zeros(shape=DEFAULT_SHAPE, dtype='uint8')
    expected[2, 2] = 1
    expected[3, 3] = 1
    expected[4, 4] = 1

    assert np.array_equal(
        rasterize([geojson_line], out_shape=DEFAULT_SHAPE),
        expected
    )


def test_rasterize_multiline(geojson_multiline):
    expected = np.zeros(shape=DEFAULT_SHAPE, dtype='uint8')
    expected[2, 2] = 1
    expected[3, 3] = 1
    expected[4, 4] = 1
    expected[0, 0:5] = 1

    assert np.array_equal(
        rasterize([geojson_multiline], out_shape=DEFAULT_SHAPE),
        expected
    )


def test_rasterize_polygon(geojson_polygon, basic_image_2x2):
    assert np.array_equal(
        rasterize([geojson_polygon], out_shape=DEFAULT_SHAPE),
        basic_image_2x2
    )


def test_rasterize_multipolygon(geojson_multipolygon):
    expected = np.zeros(shape=DEFAULT_SHAPE, dtype='uint8')
    expected[0:1, 0:1] = 1
    expected[2:4, 2:4] = 1

    assert np.array_equal(
        rasterize([geojson_multipolygon], out_shape=DEFAULT_SHAPE),
        expected
    )


def test_rasterize_geomcollection(geojson_geomcollection):
    expected = np.zeros(shape=DEFAULT_SHAPE, dtype='uint8')
    expected[0:1, 0:1] = 1
    expected[2:4, 2:4] = 1

    assert np.array_equal(
        rasterize([geojson_geomcollection], out_shape=DEFAULT_SHAPE),
        expected
    )


def test_rasterize_geo_interface(geojson_polygon, basic_image_2x2):
    assert np.array_equal(
        rasterize([MockGeoInterface(geojson_polygon)], out_shape=DEFAULT_SHAPE),
        basic_image_2x2
    )


def test_rasterize_geomcollection_no_hole():
    """
    Make sure that bug reported in
    https://github.com/rasterio/rasterio/issues/1253
    does not recur.  GeometryCollections are flattened to individual parts,
    and should result in no holes where parts overlap.
    """

    geomcollection = {'type': 'GeometryCollection', 'geometries': [
        {'type': 'Polygon',
            'coordinates': (((0, 0), (0, 5), (5, 5), (5, 0), (0, 0)),)},
        {'type': 'Polygon',
            'coordinates': (((2, 2), (2, 7), (7, 7), (7, 2), (2, 2)),)}
    ]}

    expected = rasterize(geomcollection['geometries'], out_shape=DEFAULT_SHAPE)

    assert np.array_equal(
        rasterize([geomcollection], out_shape=DEFAULT_SHAPE),
        expected
    )


def test_rasterize_multipolygon_no_hole():
    """
    Make sure that bug reported in
    https://github.com/rasterio/rasterio/issues/1253
    does not recur.  MultiPolygons are flattened to individual parts,
    and should result in no holes where parts overlap.
    """

    poly1 = (((0, 0), (0, 5), (5, 5), (5, 0), (0, 0)),)
    poly2 = (((2, 2), (2, 7), (7, 7), (7, 2), (2, 2)),)

    polys = [{'type': 'Polygon', 'coordinates': poly1},
             {'type': 'Polygon', 'coordinates': poly2}]

    multipoly = {'type': 'MultiPolygon', 'coordinates': [poly1, poly2]}

    expected = rasterize(polys, out_shape=DEFAULT_SHAPE)

    assert np.array_equal(
        rasterize([multipoly], out_shape=DEFAULT_SHAPE),
        expected
    )


@pytest.mark.parametrize(
    "input",
    [
        [{"type"}],
        [{"type": "Invalid"}],
        [{"type": "Point"}],
        [{"type": "Point", "coordinates": []}],
        [{"type": "GeometryCollection", "geometries": []}],
    ],
)
def test_rasterize_invalid_geom(input):
    """Invalid GeoJSON should be skipped with a warning."""
    with pytest.warns(ShapeSkipWarning):
        out = rasterize(input, out_shape=DEFAULT_SHAPE)

    assert out.shape == DEFAULT_SHAPE
    assert np.all(out == 0)


def test_rasterize_skip_only_invalid_geom(geojson_polygon, basic_image_2x2):
    """Rasterize operation should succeed for at least one valid geometry."""
    with pytest.warns(ShapeSkipWarning, match="Invalid or empty shape"):
        out = rasterize(
            [geojson_polygon, {"type": "Polygon", "coordinates": []}],
            out_shape=DEFAULT_SHAPE,
        )

    assert np.array_equal(out, basic_image_2x2)


@pytest.mark.parametrize("dtype", [
        None,
        "int8",
        "uint8",
        "uint16",
        "int16",
        "uint32",
        "int32",
        "int64",
        pytest.param(
            "float16",
            marks=requires_gdal3_11,
        ),
        "float32",
        "float64"
])
def test_rasterize_out_image(dtype, basic_geometry, basic_image_2x2):
    """Rasterize operation should succeed for an out image."""
    out = np.zeros(DEFAULT_SHAPE, dtype=dtype)
    rasterize([basic_geometry], out=out)
    assert np.array_equal(basic_image_2x2, out)


@requires_gdal_lt_3_11
def test_rasterize_unsupported_dtype_single(basic_geometry):
    """A non-supported data type for out should raise an exception."""
    out = np.zeros(DEFAULT_SHAPE, dtype=np.float16)
    with pytest.raises(ValueError):
        rasterize([basic_geometry], out=out)


@pytest.mark.xfail(reason="shape values are always unsafely cast to the given dtype.")
def test_rasterize_shapes_out_dtype_mismatch(basic_geometry):
    """Shape values must be able to fit in data type for out."""
    out = np.zeros(DEFAULT_SHAPE, dtype=np.uint8)
    with pytest.raises(ValueError):
        rasterize([(basic_geometry, 10000000)], out=out)


def test_rasterize_missing_out(basic_geometry):
    """If both out and out_shape are missing, should raise exception."""
    with pytest.raises(ValueError):
        rasterize([basic_geometry], out=None, out_shape=None)


def test_rasterize_missing_shapes():
    """Shapes are not required for this operation."""
    out = rasterize([], out_shape=DEFAULT_SHAPE)
    assert out.shape == DEFAULT_SHAPE
    assert np.all(out == 0)


def test_rasterize_invalid_shapes_skip():
    """Invalid shapes can be skipped with a warning, the default."""
    with pytest.warns(ShapeSkipWarning):
        rasterize([{'foo': 'bar'}], out_shape=DEFAULT_SHAPE)


def test_rasterize_invalid_shapes_no_skip():
    """Invalid shapes can raise an exception rather than be skipped."""
    with pytest.raises(ValueError):
        rasterize([{'foo': 'bar'}], out_shape=DEFAULT_SHAPE, skip_invalid=False)


@pytest.mark.parametrize("shape", [(1, 10, 10), (10,)])
def test_rasterize_invalid_out_shape(basic_geometry, shape):
    """output array shape must be 2D."""
    with pytest.raises(ValueError, match="Invalid out_shape"):
        rasterize([basic_geometry], out_shape=shape)

    with pytest.raises(ValueError) as ex:
        rasterize([basic_geometry], out_shape=(10,))
    assert 'Invalid out_shape' in str(ex.value)


@pytest.mark.parametrize("shape", [(0, 0), (1, 0), (0, 1)])
def test_rasterize_invalid_out_extent(basic_geometry, shape):
    """output array shape must have nonzero extent."""
    with pytest.raises(ValueError, match="must be > 0"):
        rasterize([basic_geometry], out_shape=shape)


def test_rasterize_default_value(basic_geometry, basic_image_2x2):
    """All shapes should rasterize to the default value."""
    default_value = 2
    truth = basic_image_2x2 * default_value

    assert np.array_equal(
        truth,
        rasterize(
            [basic_geometry], out_shape=DEFAULT_SHAPE,
            default_value=default_value
        )
    )


def test_rasterize_default_value_for_none(basic_geometry, basic_image_2x2):
    """All shapes should rasterize to the default value."""
    assert np.all(
        rasterize([(basic_geometry, None)], out_shape=DEFAULT_SHAPE, fill=2) == 2
    )


def test_rasterize_int64_default_value(basic_geometry):
    """A default value that requires an int64 succeeds for GDAL >= 3.5."""
    rasterize([basic_geometry], out_shape=DEFAULT_SHAPE, default_value=1000000000000)


def test_rasterize_fill_value(basic_geometry, basic_image_2x2):
    """All pixels not covered by shapes should be given fill value."""
    default_value = 2
    assert np.array_equal(
        basic_image_2x2 + 1,
        rasterize(
            [basic_geometry], out_shape=DEFAULT_SHAPE, fill=1,
            default_value=default_value
        )
    )


def test_rasterize_invalid_fill_value(basic_geometry):
    """Fill arg no longer determines data type."""
    out = rasterize(
        [basic_geometry], out_shape=DEFAULT_SHAPE, fill=0.5, default_value=2
    )
    assert out.dtype.kind != "f"


def test_rasterize_all_touched(basic_geometry, basic_image):
    assert np.array_equal(
        basic_image,
        rasterize(
            [basic_geometry], out_shape=DEFAULT_SHAPE, all_touched=True
        )
    )


def test_rasterize_merge_alg_add(basic_geometry, basic_image_2x2x2):
    """
    Rasterizing two times the basic_geometry with the "add" merging
    option should output the shape with the value 2
    """
    with rasterio.Env():
        assert np.array_equal(
            basic_image_2x2x2,
            rasterize(
                [basic_geometry, basic_geometry], merge_alg=MergeAlg.add,
                out_shape=DEFAULT_SHAPE)
        )


def test_rasterize_value(basic_geometry, basic_image_2x2):
    """
    All shapes should rasterize to the value passed in a tuple alongside
    each shape
    """
    value = 5
    assert np.array_equal(
        basic_image_2x2 * value,
        rasterize(
            [(basic_geometry, value)], out_shape=DEFAULT_SHAPE
        )
    )


@pytest.mark.parametrize(
    "dtype",
    [
        pytest.param(
            "float16",
            marks=requires_gdal_lt_3_11,
        ),
    ],
)
def test_rasterize_unsupported_dtype(basic_geometry, dtype):
    """Unsupported types should all raise exceptions."""
    with pytest.raises(ValueError):
        rasterize([basic_geometry], out_shape=DEFAULT_SHAPE, dtype=dtype)


def test_rasterize_geometries_symmetric():
    """Make sure that rasterize is symmetric with shapes."""
    transform = (1.0, 0.0, 0.0, 0.0, -1.0, 0.0)
    truth = np.zeros(DEFAULT_SHAPE, dtype=rasterio.ubyte)
    truth[2:5, 2:5] = 1
    s = shapes(truth, transform=transform)
    result = rasterize(s, out_shape=DEFAULT_SHAPE, transform=transform)
    assert np.array_equal(result, truth)


def test_rasterize_internal_driver_manager(basic_geometry):
    """Rasterize should work without explicitly calling driver manager."""
    assert rasterize([basic_geometry], out_shape=DEFAULT_SHAPE).sum() == 4


def test_rasterize_geo_interface_2(geojson_polygon):
    """Objects that implement the geo interface should rasterize properly"""

    class GeoObj:
        @property
        def __geo_interface__(self):
            return geojson_polygon

    assert rasterize([GeoObj()], out_shape=DEFAULT_SHAPE).sum() == 4


def test_rasterize__numpy_coordinates__fail():
    # https://github.com/rasterio/rasterio/issues/2385
    shapes = [
        (
            {
                "type": "LineString",
                "coordinates": np.array(
                    [
                        [425596.0123443, 4932971.35636043],
                        [425598.03434254, 4932966.09916503],
                        [425592.56573176, 4932963.99585319],
                        [425590.54373353, 4932969.2530486],
                        [425596.0123443, 4932971.35636043],
                    ]
                ),
            },
            2,
        ),
        (
            {
                "type": "LineString",
                "coordinates": np.array(
                    [
                        [425582.9243515, 4932973.24623693],
                        [425592.85588065, 4932951.94800393],
                        [425584.24595668, 4932947.93313045],
                        [425574.31442752, 4932969.23136344],
                        [425582.9243515, 4932973.24623693],
                    ]
                ),
            },
            2,
        ),
    ]
    out = rasterio.features.rasterize(shapes=shapes, out_shape=(100, 100))
    assert out.shape == (100, 100)
    # will fail and be filled with 0
    assert (out == 0).all()


def test_shapes(basic_image):
    """Test creation of shapes from pixel values."""
    shapely = pytest.importorskip("shapely", reason="Test requires shapely.")
    results = list(shapes(basic_image))

    assert len(results) == 2

    shape, value = results[0]
    assert shape == {
        'coordinates': [
            [(2, 2), (2, 5), (5, 5), (5, 2), (2, 2)]
        ],
        'type': 'Polygon'
    }
    assert value == 1

    shape, value = results[1]
    assert shapely.geometry.shape(shape).area == 91.0
    assert value == 0


def test_shapes_2509(basic_image):
    """Test creation of shapes from pixel values, issue #2509."""
    shapely = pytest.importorskip("shapely", reason="Test requires shapely.")
    image_with_strides = np.pad(basic_image, 1)[1:-1, 1:-1]
    np.testing.assert_array_equal(basic_image, image_with_strides)
    assert image_with_strides.__array_interface__["strides"] is not None

    results = list(shapes(image_with_strides))

    assert len(results) == 2

    shape, value = results[0]
    assert shape == {
        'coordinates': [
            [(2, 2), (2, 5), (5, 5), (5, 2), (2, 2)]
        ],
        'type': 'Polygon'
    }
    assert value == 1

    shape, value = results[1]
    assert shapely.geometry.shape(shape).area == 91.0
    assert value == 0


def test_shapes_band(pixelated_image, pixelated_image_file):
    """Shapes from a band should match shapes from an array."""
    truth = list(shapes(pixelated_image))

    with rasterio.open(pixelated_image_file) as src:
        band = rasterio.band(src, 1)
        assert truth == list(shapes(band))

        # Mask band should function, but will mask out some results
        assert truth[0] == list(shapes(band, mask=band))[0]


def test_shapes_connectivity_rook(diagonal_image):
    """
    Diagonals are not connected, so there will be 1 feature per pixel plus
    background.
    """
    assert len(list(shapes(diagonal_image, connectivity=4))) == 12


def test_shapes_connectivity_queen(diagonal_image):
    """
    Diagonals are connected, so there will be 1 feature for all pixels plus
    background.
    """
    assert len(list(shapes(diagonal_image, connectivity=8))) == 2


def test_shapes_connectivity_invalid(diagonal_image):
    """Invalid connectivity should raise exception."""
    with pytest.raises(ValueError):
        assert next(shapes(diagonal_image, connectivity=12))


def test_shapes_mask(basic_image):
    """Only pixels not masked out should be converted to features."""
    mask = np.ones(basic_image.shape, dtype=rasterio.bool_)
    mask[4:5, 4:5] = False

    results = list(shapes(basic_image, mask=mask))

    assert len(results) == 2

    shape, value = results[0]
    assert shape == {
        'coordinates': [
            [(2, 2), (2, 5), (4, 5), (4, 4), (5, 4), (5, 2), (2, 2)]
        ],
        'type': 'Polygon'
    }
    assert value == 1


def test_shapes_masked_array(basic_image):
    """Only pixels not masked out should be converted to features."""
    mask = np.full(basic_image.shape, False, dtype=rasterio.bool_)
    mask[4:5, 4:5] = True

    results = list(shapes(np.ma.masked_array(basic_image, mask=mask)))

    assert len(results) == 2

    shape, value = results[0]
    assert shape == {
        'coordinates': [
            [(2, 2), (2, 5), (4, 5), (4, 4), (5, 4), (5, 2), (2, 2)]
        ],
        'type': 'Polygon'
    }
    assert value == 1


def test_shapes_blank_mask(basic_image):
    """Mask is blank so results should mask shapes without mask."""
    assert np.array_equal(
        list(shapes(
            basic_image,
            mask=np.ones(basic_image.shape, dtype=rasterio.bool_))
        ),
        list(shapes(basic_image))
    )


def test_shapes_invalid_mask_shape(basic_image):
    """A mask that is the wrong shape should fail."""
    with pytest.raises(ValueError):
        next(shapes(
            basic_image,
            mask=np.ones(
                (basic_image.shape[0] + 10, basic_image.shape[1] + 10),
                dtype=rasterio.bool_
            )
        ))


def test_shapes_invalid_mask_dtype(basic_image):
    """A mask that is the wrong dtype should fail."""
    for dtype in ('int8', 'int16', 'int32'):
        with pytest.raises(ValueError):
            next(shapes(
                basic_image,
                mask=np.ones(basic_image.shape, dtype=dtype)
            ))


@pytest.mark.parametrize(
    "dtype, test_value, rtol",
    [
        ("int8", -127, 0),
        ("int16", -32768, 0),
        ("int32", -2147483648, 0),
        ("int64", 20439845334323, 0),
        ("uint8", 255, 0),
        ("uint16", 65535, 0),
        ("uint32", 4294967295, 0),
        pytest.param("float16", -9.1, 0.1, marks=requires_gdal3_11),
        ("float32", 1.434532, 1e-6),
        pytest.param("float64", 2717538304, 0.1, marks=requires_gdal3_12_1),
    ],
)
def test_shapes_supported_dtypes(basic_image, dtype, rtol, test_value):
    """Supported data types should return valid results."""
    shape, value = next(shapes(basic_image.astype(dtype) * test_value))
    assert_allclose(value, test_value, rtol=rtol)


@pytest.mark.parametrize(
    "dtype, test_value",
    [
        ("uint64", 20439845334323),
        pytest.param("float64", 2717538304, marks=requires_gdal_lt_3_12_1),
    ],
)
def test_shapes_partially_supported_dtypes(basic_image, dtype, test_value):
    with pytest.warns(UserWarning, match="Truncation issues"):
        shape, value = next(shapes(basic_image.astype(dtype) * test_value))
        assert_allclose(value, test_value)

@pytest.mark.parametrize(
    "dtype, test_value",
    [
        pytest.param("float16", -9.1, marks=requires_gdal_lt_3_11),
    ],
)
def test_shapes_unsupported_dtypes(basic_image, dtype, test_value):
    """Unsupported data types should raise exceptions."""
    with pytest.raises(ValueError):
        next(shapes(basic_image.astype(dtype) * test_value))


def test_shapes_internal_driver_manager(basic_image):
    """Shapes should work without explicitly calling driver manager."""
    assert next(shapes(basic_image))[0]['type'] == 'Polygon'


def test_sieve_small(basic_image, pixelated_image):
    """
    Setting the size smaller than or equal to the size of the feature in the
    image should not change the image.
    """
    assert np.array_equal(
        basic_image,
        sieve(pixelated_image, basic_image.sum())
    )


def test_sieve_large(basic_image):
    """
    Setting the size larger than size of feature should leave us an empty image.
    """
    assert not np.any(sieve(basic_image, basic_image.sum() + 1))


def test_sieve_invalid_size(basic_image):
    for invalid_size in (0, 45.1234, basic_image.size + 1):
        with pytest.raises(ValueError):
            sieve(basic_image, invalid_size)


def test_sieve_connectivity_rook(diagonal_image):
    """Diagonals are not connected, so feature is removed."""
    assert not np.any(
        sieve(diagonal_image, diagonal_image.sum(), connectivity=4)
    )


def test_sieve_connectivity_queen(diagonal_image):
    """Diagonals are connected, so feature is retained."""
    assert np.array_equal(
        diagonal_image,
        sieve(diagonal_image, diagonal_image.sum(), connectivity=8)
    )


def test_sieve_connectivity_invalid(basic_image):
    with pytest.raises(ValueError):
        sieve(basic_image, 54, connectivity=12)


def test_sieve_out(basic_image):
    """Output array passed in should match the returned array."""
    output = np.zeros_like(basic_image)
    output[1:3, 1:3] = 5
    sieved_image = sieve(basic_image, basic_image.sum(), out=output)
    assert np.array_equal(basic_image, sieved_image)
    assert np.array_equal(output, sieved_image)


def test_sieve_invalid_out(basic_image):
    """Output with different dtype or shape should fail."""
    with pytest.raises(ValueError):
        sieve(
            basic_image, basic_image.sum(),
            out=np.zeros(basic_image.shape, dtype=rasterio.int32)
        )

    with pytest.raises(ValueError):
        sieve(
            basic_image, basic_image.sum(),
            out=np.zeros(
                (basic_image.shape[0] + 10, basic_image.shape[1] + 10),
                dtype=rasterio.ubyte
            )
        )


def test_sieve_mask(basic_image):
    """
    Only areas within the overlap of mask and input will be kept, so long
    as mask is a bool or uint8 dtype.
    """
    mask = np.ones(basic_image.shape, dtype=rasterio.bool_)
    mask[4:5, 4:5] = False
    truth = basic_image * np.invert(mask)

    sieved_image = sieve(basic_image, basic_image.sum(), mask=mask)
    assert sieved_image.sum() > 0

    assert np.array_equal(
        truth,
        sieved_image
    )

    assert np.array_equal(
        truth.astype(rasterio.uint8),
        sieved_image
    )


def test_sieve_blank_mask(basic_image):
    """A blank mask should have no effect."""
    mask = np.ones(basic_image.shape, dtype=rasterio.bool_)
    assert np.array_equal(
        basic_image,
        sieve(basic_image, basic_image.sum(), mask=mask)
    )


@pytest.mark.xfail(reason="Upstream bug in GDAL, see #3074.")
def test_sieve_all_masked(basic_image):
    """A blank mask should have no effect."""
    mask = np.full(basic_image.shape, False)
    # mask[-1, -1] = True
    sieved = sieve(basic_image, basic_image.sum(), mask=mask)
    assert np.array_equal(
        basic_image,
        sieved
    )


def test_sieve_invalid_mask_shape(basic_image):
    """A mask that is the wrong shape should fail."""
    with pytest.raises(ValueError):
        sieve(
            basic_image, basic_image.sum(),
            mask=np.ones(
                (basic_image.shape[0] + 10, basic_image.shape[1] + 10),
                dtype=rasterio.bool_
            )
        )


def test_sieve_invalid_mask_dtype(basic_image):
    """A mask that is the wrong dtype should fail."""
    for dtype in ('int8', 'int16', 'int32'):
        with pytest.raises(ValueError):
            sieve(
                basic_image, basic_image.sum(),
                mask=np.ones(basic_image.shape, dtype=dtype)
            )


def test_sieve_supported_dtypes(basic_image):
    """Supported data types should return valid results."""
    supported_types = (
        ('int16', -32768),
        ('int32', -2147483648),
        ('uint8', 255),
        ('uint16', 65535)
    )

    for dtype, test_value in supported_types:
        truth = (basic_image).astype(dtype) * test_value
        sieved_image = sieve(truth, basic_image.sum())
        assert np.array_equal(truth, sieved_image)
        assert np.dtype(sieved_image.dtype) == np.dtype(dtype)


def test_sieve_unsupported_dtypes(basic_image):
    """Unsupported data types should raise exceptions."""
    unsupported_types = (
        ('int8', -127),
        ('uint32', 4294967295),
        ('int64', 20439845334323),
        ('float16', -9343.232),
        ('float32', 1.434532),
        ('float64', -98332.133422114)
    )

    for dtype, test_value in unsupported_types:
        with pytest.raises(ValueError):
            sieve(
                (basic_image).astype(dtype) * test_value,
                basic_image.sum()
            )


def test_sieve_band(pixelated_image, pixelated_image_file):
    """Sieving a band from a raster file should match sieve of array."""

    truth = sieve(pixelated_image, 9)

    with rasterio.open(pixelated_image_file) as src:
        band = rasterio.band(src, 1)
        assert np.array_equal(truth, sieve(band, 9))

        # Mask band should also work but will be a no-op
        assert np.array_equal(
            pixelated_image,
            sieve(band, 9, mask=band)
        )


def test_sieve_internal_driver_manager(capfd, basic_image, pixelated_image):
    """Sieve should work without explicitly calling driver manager."""
    assert np.array_equal(
        basic_image,
        sieve(pixelated_image, basic_image.sum())
    )


def test_zz_no_dataset_leaks(capfd):
    with rasterio.Env() as env:
        env._dump_open_datasets()
        captured = capfd.readouterr()
        assert not captured.err


def test_sieve_bands(pixelated_image, pixelated_image_file):
    """Verify fix for gh-2782."""
    truth = sieve(pixelated_image, 9)

    with rasterio.open(pixelated_image_file) as src:
        assert np.array_equal(truth, sieve(rasterio.band(src, [1]), 9))

        # Mask band should also work but will be a no-op
        assert np.array_equal(
            pixelated_image,
            sieve(rasterio.band(src, [1]), 9, mask=rasterio.band(src, 1))
        )


def test_sieve_dataset(pixelated_image, pixelated_image_file):
    """Verify fix for gh-2782."""
    truth = sieve(pixelated_image, 9)

    with rasterio.open(pixelated_image_file) as src:
        assert np.array_equal(truth, sieve(src, 9))

        # Mask band should also work but will be a no-op
        assert np.array_equal(
            pixelated_image,
            sieve(src, 9, mask=rasterio.band(src, 1))
        )
