import pytest
import altair.vegalite.v5 as alt


def geom_obj(geom):
    class Geom:
        pass

    geom_obj = Geom()
    geom_obj.__geo_interface__ = geom
    return geom_obj


# correct translation of Polygon geometry to Feature type
def test_geo_interface_polygon_feature():
    geom = {
        "coordinates": [[(0, 0), (0, 2), (2, 2), (2, 0), (0, 0)]],
        "type": "Polygon",
    }
    feat = geom_obj(geom)

    with alt.data_transformers.enable(consolidate_datasets=False):
        spec = alt.Chart(feat).mark_geoshape().to_dict()
    assert spec["data"]["values"]["type"] == "Feature"


# merge geometry with empty properties dictionary
def test_geo_interface_removal_empty_properties():
    geom = {
        "geometry": {
            "coordinates": [
                [[6.90, 53.48], [5.98, 51.85], [6.07, 53.51], [6.90, 53.48]]
            ],
            "type": "Polygon",
        },
        "id": None,
        "properties": {},
        "type": "Feature",
    }
    feat = geom_obj(geom)

    with alt.data_transformers.enable(consolidate_datasets=False):
        spec = alt.Chart(feat).mark_geoshape().to_dict()
    assert spec["data"]["values"]["type"] == "Feature"


# only register metadata in the properties member
def test_geo_interface_register_foreign_member():
    geom = {
        "geometry": {
            "coordinates": [
                [[6.90, 53.48], [5.98, 51.85], [6.07, 53.51], [6.90, 53.48]]
            ],
            "type": "Polygon",
        },
        "id": 2,
        "properties": {"foo": "bah"},
        "type": "Feature",
    }
    feat = geom_obj(geom)

    with alt.data_transformers.enable(consolidate_datasets=False):
        spec = alt.Chart(feat).mark_geoshape().to_dict()
    with pytest.raises(KeyError):
        spec["data"]["values"]["id"]
    assert spec["data"]["values"]["foo"] == "bah"


# correct serializing of arrays and nested tuples
def test_geo_interface_serializing_arrays_tuples():
    import array as arr

    geom = {
        "bbox": arr.array("d", [1, 2, 3, 4]),
        "geometry": {
            "coordinates": [
                (
                    (6.90, 53.48),
                    (5.98, 51.85),
                    (6.07, 53.51),
                    (6.90, 53.48),
                )
            ],
            "type": "Polygon",
        },
        "id": 27,
        "properties": {},
        "type": "Feature",
    }
    feat = geom_obj(geom)

    with alt.data_transformers.enable(consolidate_datasets=False):
        spec = alt.Chart(feat).mark_geoshape().to_dict()
    assert spec["data"]["values"]["geometry"]["coordinates"][0][0] == [6.9, 53.48]


# overwrite existing 'type' value in properties with `Feature`
def test_geo_interface_reserved_members():
    geom = {
        "geometry": {
            "coordinates": [
                [[6.90, 53.48], [5.98, 51.85], [6.07, 53.51], [6.90, 53.48]]
            ],
            "type": "Polygon",
        },
        "id": 27,
        "properties": {"type": "foo"},
        "type": "Feature",
    }
    feat = geom_obj(geom)

    with alt.data_transformers.enable(consolidate_datasets=False):
        spec = alt.Chart(feat).mark_geoshape().to_dict()
    assert spec["data"]["values"]["type"] == "Feature"


# an empty FeatureCollection is valid
def test_geo_interface_empty_feature_collection():
    geom = {"type": "FeatureCollection", "features": []}
    feat = geom_obj(geom)

    with alt.data_transformers.enable(consolidate_datasets=False):
        spec = alt.Chart(feat).mark_geoshape().to_dict()
    assert spec["data"]["values"] == []


# Features in a FeatureCollection only keep properties and geometry
def test_geo_interface_feature_collection():
    geom = {
        "type": "FeatureCollection",
        "features": [
            {
                "geometry": {
                    "coordinates": [
                        [[6.90, 53.48], [5.98, 51.85], [6.07, 53.51], [6.90, 53.48]]
                    ],
                    "type": "Polygon",
                },
                "id": 27,
                "properties": {"type": "foo", "id": 1, "geometry": 1},
                "type": "Feature",
            },
            {
                "geometry": {
                    "coordinates": [
                        [[8.90, 53.48], [7.98, 51.85], [8.07, 53.51], [8.90, 53.48]]
                    ],
                    "type": "Polygon",
                },
                "id": 28,
                "properties": {"type": "foo", "id": 2, "geometry": 1},
                "type": "Feature",
            },
        ],
    }
    feat = geom_obj(geom)

    with alt.data_transformers.enable(consolidate_datasets=False):
        spec = alt.Chart(feat).mark_geoshape().to_dict()
    assert spec["data"]["values"][0]["id"] == 1
    assert spec["data"]["values"][1]["id"] == 2
    assert "coordinates" in spec["data"]["values"][0]["geometry"]
    assert "coordinates" in spec["data"]["values"][1]["geometry"]
    assert spec["data"]["values"][0]["type"] == "Feature"
    assert spec["data"]["values"][1]["type"] == "Feature"


# typical output of a __geo_interface__ from geopandas GeoDataFrame
# notic that the index value is registerd as a commonly used identifier
# with the name "id" (in this case 49). Similar to serialization of a
# pandas DataFrame is the index not included in the output
def test_geo_interface_feature_collection_gdf():
    geom = {
        "bbox": (19.89, -26.82, 29.43, -17.66),
        "features": [
            {
                "bbox": (19.89, -26.82, 29.43, -17.66),
                "geometry": {
                    "coordinates": [
                        [[6.90, 53.48], [5.98, 51.85], [6.07, 53.51], [6.90, 53.48]]
                    ],
                    "type": "Polygon",
                },
                "id": "49",
                "properties": {
                    "continent": "Africa",
                    "gdp_md_est": 35900.0,
                    "id": "BWA",
                    "iso_a3": "BWA",
                    "name": "Botswana",
                    "pop_est": 2214858,
                },
                "type": "Feature",
            }
        ],
        "type": "FeatureCollection",
    }
    feat = geom_obj(geom)

    with alt.data_transformers.enable(consolidate_datasets=False):
        spec = alt.Chart(feat).mark_geoshape().to_dict()
    assert spec["data"]["values"][0]["id"] == "BWA"