# test zonal stats
import json
import os
import pytest
import simplejson
import sys

import numpy as np
import rasterio
from rasterstats import zonal_stats, raster_stats
from rasterstats.utils import VALID_STATS
from rasterstats.io import read_featurecollection, read_features
from shapely.geometry import Polygon
from affine import Affine

sys.path.append(os.path.dirname(os.path.abspath(__file__)))

DATA = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
raster = os.path.join(DATA, "slope.tif")


def test_main():
    polygons = os.path.join(DATA, "polygons.shp")
    stats = zonal_stats(polygons, raster)
    for key in ["count", "min", "max", "mean"]:
        assert key in stats[0]
    assert len(stats) == 2
    assert stats[0]["count"] == 75
    assert stats[1]["count"] == 50
    assert round(stats[0]["mean"], 2) == 14.66


# remove after band_num alias is removed
def test_band_alias():
    polygons = os.path.join(DATA, "polygons.shp")
    stats_a = zonal_stats(polygons, raster)
    stats_b = zonal_stats(polygons, raster, band=1)
    with pytest.deprecated_call():
        stats_c = zonal_stats(polygons, raster, band_num=1)
    assert stats_a[0]["count"] == stats_b[0]["count"] == stats_c[0]["count"]


def test_zonal_global_extent():
    polygons = os.path.join(DATA, "polygons.shp")
    stats = zonal_stats(polygons, raster)
    global_stats = zonal_stats(polygons, raster, global_src_extent=True)
    assert stats == global_stats


def test_zonal_nodata():
    polygons = os.path.join(DATA, "polygons.shp")
    stats = zonal_stats(polygons, raster, nodata=0)
    assert len(stats) == 2
    assert stats[0]["count"] == 75
    assert stats[1]["count"] == 50


def test_doesnt_exist():
    nonexistent = os.path.join(DATA, "DOESNOTEXIST.shp")
    with pytest.raises(ValueError):
        zonal_stats(nonexistent, raster)


def test_nonsense():
    polygons = os.path.join(DATA, "polygons.shp")
    with pytest.raises(ValueError):
        zonal_stats("blaghrlargh", raster)
    with pytest.raises(IOError):
        zonal_stats(polygons, "blercherlerch")
    with pytest.raises(ValueError):
        zonal_stats(
            [
                "blaghrlargh",
            ],
            raster,
        )


# Different geometry types
def test_points():
    points = os.path.join(DATA, "points.shp")
    stats = zonal_stats(points, raster)
    # three features
    assert len(stats) == 3
    # three pixels
    assert sum([x["count"] for x in stats]) == 3
    assert round(stats[0]["mean"], 3) == 11.386
    assert round(stats[1]["mean"], 3) == 35.547


def test_points_categorical():
    points = os.path.join(DATA, "points.shp")
    categorical_raster = os.path.join(DATA, "slope_classes.tif")
    stats = zonal_stats(points, categorical_raster, categorical=True)
    # three features
    assert len(stats) == 3
    assert "mean" not in stats[0]
    assert stats[0][1.0] == 1
    assert stats[1][2.0] == 1


def test_lines():
    lines = os.path.join(DATA, "lines.shp")
    stats = zonal_stats(lines, raster)
    assert len(stats) == 2
    assert stats[0]["count"] == 58
    assert stats[1]["count"] == 32


# Test multigeoms
def test_multipolygons():
    multipolygons = os.path.join(DATA, "multipolygons.shp")
    stats = zonal_stats(multipolygons, raster)
    assert len(stats) == 1
    assert stats[0]["count"] == 125


def test_multilines():
    multilines = os.path.join(DATA, "multilines.shp")
    stats = zonal_stats(multilines, raster)
    assert len(stats) == 1
    # can differ slightly based on platform/gdal version
    assert stats[0]["count"] in [89, 90]


def test_multipoints():
    multipoints = os.path.join(DATA, "multipoints.shp")
    stats = zonal_stats(multipoints, raster)
    assert len(stats) == 1
    assert stats[0]["count"] == 3


def test_categorical():
    polygons = os.path.join(DATA, "polygons.shp")
    categorical_raster = os.path.join(DATA, "slope_classes.tif")
    stats = zonal_stats(polygons, categorical_raster, categorical=True)
    assert len(stats) == 2
    assert stats[0][1.0] == 75
    assert 5.0 in stats[1]


def test_categorical_map():
    polygons = os.path.join(DATA, "polygons.shp")
    categorical_raster = os.path.join(DATA, "slope_classes.tif")
    catmap = {5.0: "cat5"}
    stats = zonal_stats(
        polygons, categorical_raster, categorical=True, category_map=catmap
    )
    assert len(stats) == 2
    assert stats[0][1.0] == 75
    assert 5.0 not in stats[1]
    assert "cat5" in stats[1]


def test_specify_stats_list():
    polygons = os.path.join(DATA, "polygons.shp")
    stats = zonal_stats(polygons, raster, stats=["min", "max"])
    assert sorted(stats[0].keys()) == sorted(["min", "max"])
    assert "count" not in list(stats[0].keys())


def test_specify_all_stats():
    polygons = os.path.join(DATA, "polygons.shp")
    stats = zonal_stats(polygons, raster, stats="ALL")
    assert sorted(stats[0].keys()) == sorted(VALID_STATS)
    stats = zonal_stats(polygons, raster, stats="*")
    assert sorted(stats[0].keys()) == sorted(VALID_STATS)


def test_specify_stats_string():
    polygons = os.path.join(DATA, "polygons.shp")
    stats = zonal_stats(polygons, raster, stats="min max")
    assert sorted(stats[0].keys()) == sorted(["min", "max"])
    assert "count" not in list(stats[0].keys())


def test_specify_stats_invalid():
    polygons = os.path.join(DATA, "polygons.shp")
    with pytest.raises(ValueError):
        zonal_stats(polygons, raster, stats="foo max")


def test_optional_stats():
    polygons = os.path.join(DATA, "polygons.shp")
    stats = zonal_stats(polygons, raster, stats="min max sum majority median std")
    assert stats[0]["min"] <= stats[0]["median"] <= stats[0]["max"]


def test_range():
    polygons = os.path.join(DATA, "polygons.shp")
    stats = zonal_stats(polygons, raster, stats="range min max")
    for stat in stats:
        assert stat["range"] == stat["max"] - stat["min"]
    ranges = [x["range"] for x in stats]
    # without min/max specified
    stats = zonal_stats(polygons, raster, stats="range")
    assert "min" not in stats[0]
    assert ranges == [x["range"] for x in stats]


def test_nodata():
    polygons = os.path.join(DATA, "polygons.shp")
    categorical_raster = os.path.join(DATA, "slope_classes.tif")
    stats = zonal_stats(
        polygons, categorical_raster, stats="*", categorical=True, nodata=1.0
    )
    assert stats[0]["majority"] is None
    assert stats[0]["count"] == 0  # no pixels; they're all null
    assert stats[1]["minority"] == 2.0
    assert stats[1]["count"] == 49  # used to be 50 if we allowed 1.0
    assert "1.0" not in stats[0]


def test_dataset_mask():
    polygons = os.path.join(DATA, "polygons.shp")
    raster = os.path.join(DATA, "dataset_mask.tif")
    stats = zonal_stats(polygons, raster, stats="*")
    assert stats[0]["count"] == 75
    assert stats[1]["count"] == 0


def test_partial_overlap():
    polygons = os.path.join(DATA, "polygons_partial_overlap.shp")
    stats = zonal_stats(polygons, raster, stats="count")
    for res in stats:
        # each polygon should have at least a few pixels overlap
        assert res["count"] > 0


def test_no_overlap():
    polygons = os.path.join(DATA, "polygons_no_overlap.shp")
    stats = zonal_stats(polygons, raster, stats="count")
    for res in stats:
        # no polygon should have any overlap
        assert res["count"] == 0


def test_all_touched():
    polygons = os.path.join(DATA, "polygons.shp")
    stats = zonal_stats(polygons, raster, all_touched=True)
    assert stats[0]["count"] == 95  # 75 if ALL_TOUCHED=False
    assert stats[1]["count"] == 73  # 50 if ALL_TOUCHED=False


def test_ndarray_without_affine():
    with rasterio.open(raster) as src:
        polygons = os.path.join(DATA, "polygons.shp")
        with pytest.raises(ValueError):
            zonal_stats(polygons, src.read(1))  # needs affine kwarg


def _assert_dict_eq(a, b):
    """Assert that dicts a and b similar within floating point precision"""
    err = 1e-5
    for k in set(a.keys()).union(set(b.keys())):
        if a[k] == b[k]:
            continue
        try:
            if abs(a[k] - b[k]) > err:
                raise AssertionError("{}: {} != {}".format(k, a[k], b[k]))
        except TypeError:  # can't take abs, nan
            raise AssertionError("{} != {}".format(a[k], b[k]))


def test_ndarray():
    with rasterio.open(raster) as src:
        arr = src.read(1)
        affine = src.transform

    polygons = os.path.join(DATA, "polygons.shp")
    stats = zonal_stats(polygons, arr, affine=affine)
    stats2 = zonal_stats(polygons, raster)
    for s1, s2 in zip(stats, stats2):
        _assert_dict_eq(s1, s2)
    with pytest.raises(AssertionError):
        _assert_dict_eq(stats[0], stats[1])
    assert stats[0]["count"] == 75
    assert stats[1]["count"] == 50

    points = os.path.join(DATA, "points.shp")
    stats = zonal_stats(points, arr, affine=affine)
    assert stats == zonal_stats(points, raster)
    assert sum([x["count"] for x in stats]) == 3
    assert round(stats[0]["mean"], 3) == 11.386
    assert round(stats[1]["mean"], 3) == 35.547


def test_alias():
    polygons = os.path.join(DATA, "polygons.shp")
    stats = zonal_stats(polygons, raster)
    with pytest.deprecated_call():
        stats2 = raster_stats(polygons, raster)
    assert stats == stats2


def test_add_stats():
    polygons = os.path.join(DATA, "polygons.shp")

    def mymean(x):
        return np.ma.mean(x)

    stats = zonal_stats(polygons, raster, add_stats={"mymean": mymean})
    for i in range(len(stats)):
        assert stats[i]["mean"] == stats[i]["mymean"]


def test_add_stats_prop():
    polygons = os.path.join(DATA, "polygons.shp")

    def mymean_prop(x, prop):
        return np.ma.mean(x) * prop["id"]

    stats = zonal_stats(polygons, raster, add_stats={"mymean_prop": mymean_prop})
    for i in range(len(stats)):
        assert stats[i]["mymean_prop"] == stats[i]["mean"] * (i + 1)


def test_mini_raster():
    polygons = os.path.join(DATA, "polygons.shp")
    stats = zonal_stats(polygons, raster, raster_out=True)
    stats2 = zonal_stats(
        polygons,
        stats[0]["mini_raster_array"],
        raster_out=True,
        affine=stats[0]["mini_raster_affine"],
    )
    assert (
        stats[0]["mini_raster_array"] == stats2[0]["mini_raster_array"]
    ).sum() == stats[0]["count"]


def test_percentile_good():
    polygons = os.path.join(DATA, "polygons.shp")
    stats = zonal_stats(polygons, raster, stats="median percentile_50 percentile_90")
    assert "percentile_50" in stats[0].keys()
    assert "percentile_90" in stats[0].keys()
    assert stats[0]["percentile_50"] == stats[0]["median"]
    assert stats[0]["percentile_50"] <= stats[0]["percentile_90"]


def test_zone_func_has_return():
    def example_zone_func(zone_arr):
        return np.ma.masked_array(np.full(zone_arr.shape, 1))

    polygons = os.path.join(DATA, "polygons.shp")
    stats = zonal_stats(polygons, raster, zone_func=example_zone_func)
    assert stats[0]["max"] == 1
    assert stats[0]["min"] == 1
    assert stats[0]["mean"] == 1


def test_zone_func_good():
    def example_zone_func(zone_arr):
        zone_arr[:] = 0

    polygons = os.path.join(DATA, "polygons.shp")
    stats = zonal_stats(polygons, raster, zone_func=example_zone_func)
    assert stats[0]["max"] == 0
    assert stats[0]["min"] == 0
    assert stats[0]["mean"] == 0


def test_zone_func_bad():
    not_a_func = "jar jar binks"
    polygons = os.path.join(DATA, "polygons.shp")
    with pytest.raises(TypeError):
        zonal_stats(polygons, raster, zone_func=not_a_func)


def test_percentile_nodata():
    polygons = os.path.join(DATA, "polygons.shp")
    categorical_raster = os.path.join(DATA, "slope_classes.tif")
    # By setting nodata to 1, one of our polygons is within the raster extent
    # but has an empty masked array
    stats = zonal_stats(polygons, categorical_raster, stats=["percentile_90"], nodata=1)
    assert "percentile_90" in stats[0].keys()
    assert [None, 5.0] == [x["percentile_90"] for x in stats]


def test_percentile_bad():
    polygons = os.path.join(DATA, "polygons.shp")
    with pytest.raises(ValueError):
        zonal_stats(polygons, raster, stats="percentile_101")


def test_json_serializable():
    polygons = os.path.join(DATA, "polygons.shp")
    stats = zonal_stats(
        polygons, raster, stats=VALID_STATS + ["percentile_90"], categorical=True
    )
    try:
        json.dumps(stats)
        simplejson.dumps(stats)
    except TypeError:
        pytest.fail("zonal_stats returned a list that wasn't JSON-serializable")


def test_direct_features_collections():
    polygons = os.path.join(DATA, "polygons.shp")
    features = read_features(polygons)
    collection = read_featurecollection(polygons)

    stats_direct = zonal_stats(polygons, raster)
    stats_features = zonal_stats(features, raster)
    stats_collection = zonal_stats(collection, raster)

    assert stats_direct == stats_features == stats_collection


def test_all_nodata():
    polygons = os.path.join(DATA, "polygons.shp")
    raster = os.path.join(DATA, "all_nodata.tif")
    stats = zonal_stats(polygons, raster, stats=["nodata", "count"])
    assert stats[0]["nodata"] == 75
    assert stats[0]["count"] == 0
    assert stats[1]["nodata"] == 50
    assert stats[1]["count"] == 0


def test_some_nodata():
    polygons = os.path.join(DATA, "polygons.shp")
    raster = os.path.join(DATA, "slope_nodata.tif")
    stats = zonal_stats(polygons, raster, stats=["nodata", "count"])
    assert stats[0]["nodata"] == 36
    assert stats[0]["count"] == 39
    assert stats[1]["nodata"] == 19
    assert stats[1]["count"] == 31


# update this if nan end up being incorporated into nodata
def test_nan_nodata():
    polygon = Polygon([[0, 0], [2, 0], [2, 2], [0, 2]])
    arr = np.array([[np.nan, 12.25], [-999, 12.75]])
    affine = Affine(1, 0, 0, 0, -1, 2)

    stats = zonal_stats(
        polygon, arr, affine=affine, nodata=-999, stats="nodata count sum mean min max"
    )

    assert stats[0]["nodata"] == 1
    assert stats[0]["count"] == 2
    assert stats[0]["mean"] == 12.5
    assert stats[0]["min"] == 12.25
    assert stats[0]["max"] == 12.75


def test_some_nodata_ndarray():
    polygons = os.path.join(DATA, "polygons.shp")
    raster = os.path.join(DATA, "slope_nodata.tif")
    with rasterio.open(raster) as src:
        arr = src.read(1)
        affine = src.transform

    # without nodata
    stats = zonal_stats(polygons, arr, affine=affine, stats=["nodata", "count", "min"])
    assert stats[0]["min"] == -9999.0
    assert stats[0]["nodata"] == 0
    assert stats[0]["count"] == 75

    # with nodata
    stats = zonal_stats(
        polygons, arr, affine=affine, nodata=-9999.0, stats=["nodata", "count", "min"]
    )
    assert stats[0]["min"] >= 0.0
    assert stats[0]["nodata"] == 36
    assert stats[0]["count"] == 39


def test_transform():
    with rasterio.open(raster) as src:
        arr = src.read(1)
        affine = src.transform
    polygons = os.path.join(DATA, "polygons.shp")

    stats = zonal_stats(polygons, arr, affine=affine)
    with pytest.deprecated_call():
        stats2 = zonal_stats(polygons, arr, transform=affine.to_gdal())
    assert stats == stats2


def test_prefix():
    polygons = os.path.join(DATA, "polygons.shp")
    stats = zonal_stats(polygons, raster, prefix="TEST")
    for key in ["count", "min", "max", "mean"]:
        assert key not in stats[0]
    for key in ["TESTcount", "TESTmin", "TESTmax", "TESTmean"]:
        assert key in stats[0]


def test_geojson_out():
    polygons = os.path.join(DATA, "polygons.shp")
    features = zonal_stats(polygons, raster, geojson_out=True)
    for feature in features:
        assert feature["type"] == "Feature"
        assert "id" in feature["properties"]  # from orig
        assert "count" in feature["properties"]  # from zonal stats


# do not think this is actually testing the line i wanted it to
# since the read_features func for this data type is generating
# the properties field
def test_geojson_out_with_no_properties():
    polygon = Polygon([[0, 0], [0, 0.5], [1, 1.5], [1.5, 2], [2, 2], [2, 0]])
    arr = np.array([[100, 1], [100, 1]])
    affine = Affine(1, 0, 0, 0, -1, 2)

    stats = zonal_stats(polygon, arr, affine=affine, geojson_out=True)
    assert "properties" in stats[0]
    for key in ["count", "min", "max", "mean"]:
        assert key in stats[0]["properties"]

    assert stats[0]["properties"]["mean"] == 34


# remove when copy_properties alias is removed
def test_copy_properties_warn():
    polygons = os.path.join(DATA, "polygons.shp")
    # run once to trigger any other unrelated deprecation warnings
    # so the test does not catch them instead
    stats_a = zonal_stats(polygons, raster)
    with pytest.deprecated_call():
        stats_b = zonal_stats(polygons, raster, copy_properties=True)
    assert stats_a == stats_b


def test_nan_counts():
    from affine import Affine

    transform = Affine(1, 0, 1, 0, -1, 3)

    data = np.array([[np.nan, np.nan, np.nan], [0, 0, 0], [1, 4, 5]])

    # geom extends an additional row to left
    geom = "POLYGON ((1 0, 4 0, 4 3, 1 3, 1 0))"

    # nan stat is requested
    stats = zonal_stats(geom, data, affine=transform, nodata=0.0, stats="*")

    for res in stats:
        assert res["count"] == 3  # 3 pixels of valid data
        assert res["nodata"] == 3  # 3 pixels of nodata
        assert res["nan"] == 3  # 3 pixels of nans

    # nan are ignored if nan stat is not requested
    stats = zonal_stats(geom, data, affine=transform, nodata=0.0, stats="count nodata")

    for res in stats:
        assert res["count"] == 3  # 3 pixels of valid data
        assert res["nodata"] == 3  # 3 pixels of nodata
        assert "nan" not in res


# Optional tests
def test_geodataframe_zonal():
    gpd = pytest.importorskip("geopandas")

    polygons = os.path.join(DATA, "polygons.shp")
    df = gpd.read_file(polygons)
    if not hasattr(df, "__geo_interface__"):
        pytest.skip("This version of geopandas doesn't support df.__geo_interface__")

    expected = zonal_stats(polygons, raster)
    assert zonal_stats(df, raster) == expected


# TODO #  gen_zonal_stats(<features_without_props>)
# TODO #  gen_zonal_stats(stats=nodata)
# TODO #  gen_zonal_stats(<raster with non-integer dtype>)
# TODO #  gen_zonal_stats(transform AND affine>)
