"""test interpolation functions."""

import dask_geopandas
import geopandas
import numpy as np
import pytest
from geopandas.testing import assert_geodataframe_equal
from numpy.testing import assert_almost_equal
from scipy.sparse import csr_matrix
from shapely.geometry import box

from tobler.area_weighted import area_interpolate, area_interpolate_dask
from tobler.area_weighted.area_interpolate import _area_tables_binning


def test_area_interpolate_singlecore(datasets):
    sac1, sac2 = datasets
    with pytest.WARN_VAR_VALS_INF:
        area = area_interpolate(
            source_df=sac1,
            target_df=sac2,
            extensive_variables=["TOT_POP"],
            intensive_variables=["pct_poverty"],
            categorical_variables=["animal"],
            n_jobs=1,
        )
    assert_almost_equal(area.TOT_POP.sum(), 1796856, decimal=0)
    assert_almost_equal(area.pct_poverty.sum(), 2140, decimal=0)
    assert_almost_equal(area.animal_cat.sum(), 32, decimal=0)
    assert_almost_equal(area.animal_dog.sum(), 19, decimal=0)
    assert_almost_equal(area.animal_donkey.sum(), 22, decimal=0)
    assert_almost_equal(area.animal_wombat.sum(), 23, decimal=0)
    assert_almost_equal(area.animal_capybara.sum(), 20, decimal=0)


def test_area_interpolate_extensive(datasets):
    sac1, sac2 = datasets
    area = area_interpolate(
        source_df=sac1.to_crs(4326),  # trigger warning once
        target_df=sac2.to_crs(4326),
        extensive_variables=["TOT_POP"],
        n_jobs=1,
    )
    assert_almost_equal(area.TOT_POP.sum(), 1796856, decimal=0)


def test_area_interpolate_intensive(datasets):
    sac1, sac2 = datasets
    with pytest.WARN_VAR_VALS_INF:
        area = area_interpolate(
            source_df=sac1,
            target_df=sac2,
            intensive_variables=["pct_poverty"],
            n_jobs=1,
        )
    assert_almost_equal(area.pct_poverty.sum(), 2140, decimal=0)


def test_area_interpolate_categorical(datasets):
    sac1, sac2 = datasets
    with pytest.WARN_VAR_VALS_INF:
        area = area_interpolate(
            source_df=sac1,
            target_df=sac2,
            extensive_variables=["TOT_POP"],
            intensive_variables=["pct_poverty"],
            categorical_variables=["animal"],
            n_jobs=1,
        )
    assert_almost_equal(area.animal_cat.sum(), 32, decimal=0)
    assert_almost_equal(area.animal_dog.sum(), 19, decimal=0)
    assert_almost_equal(area.animal_donkey.sum(), 22, decimal=0)
    assert_almost_equal(area.animal_wombat.sum(), 23, decimal=0)
    assert_almost_equal(area.animal_capybara.sum(), 20, decimal=0)


@pytest.mark.xfail(reason="dask_geopandas is broken with dask-expr backend")
def test_area_interpolate_categorical_dask(datasets):
    sac1, sac2 = datasets
    sac1["animal"] = sac1["animal"].astype("category")
    dsac1 = dask_geopandas.from_geopandas(sac1, npartitions=2).spatial_shuffle(
        by="hilbert", shuffle="tasks"
    )
    dsac2 = dask_geopandas.from_geopandas(sac2, npartitions=2).spatial_shuffle(
        by="hilbert", shuffle="tasks"
    )
    with pytest.WARN_VAR_VALS_INF:
        area = area_interpolate_dask(
            source_dgdf=dsac1,
            target_dgdf=dsac2,
            id_col="ZIP",
            categorical_variables=["animal"],
        ).compute()
    assert_almost_equal(area.animal_cat.sum(), 32, decimal=0)
    assert_almost_equal(area.animal_dog.sum(), 19, decimal=0)
    assert_almost_equal(area.animal_donkey.sum(), 22, decimal=0)
    assert_almost_equal(area.animal_wombat.sum(), 23, decimal=0)
    assert_almost_equal(area.animal_capybara.sum(), 20, decimal=0)


def test_area_interpolate_custom_index(datasets):
    sac1, sac2 = datasets
    sac1.index = sac1.index * 2
    sac2.index = sac2.index * 13
    with pytest.WARN_VAR_VALS_INF:
        area = area_interpolate(
            source_df=sac1,
            target_df=sac2,
            extensive_variables=["TOT_POP"],
            intensive_variables=["pct_poverty"],
            categorical_variables=["animal"],
        )
    assert_almost_equal(area.TOT_POP.sum(), 1796856, decimal=0)
    assert_almost_equal(area.pct_poverty.sum(), 2140, decimal=0)
    assert_almost_equal(area.animal_cat.sum(), 32, decimal=0)
    assert_almost_equal(area.animal_dog.sum(), 19, decimal=0)
    assert_almost_equal(area.animal_donkey.sum(), 22, decimal=0)
    assert_almost_equal(area.animal_wombat.sum(), 23, decimal=0)
    assert_almost_equal(area.animal_capybara.sum(), 20, decimal=0)
    assert not area.isna().any().any()


def test_area_interpolate_sindex_options(datasets):
    sac1, sac2 = datasets
    with pytest.WARN_VAR_VALS_INF:
        auto = area_interpolate(
            source_df=sac1,
            target_df=sac2,
            extensive_variables=["TOT_POP"],
            intensive_variables=["pct_poverty"],
        )
    with pytest.WARN_VAR_VALS_INF:
        source = area_interpolate(
            source_df=sac1,
            target_df=sac2,
            extensive_variables=["TOT_POP"],
            intensive_variables=["pct_poverty"],
            spatial_index="source",
        )
    with pytest.WARN_VAR_VALS_INF:
        target = area_interpolate(
            source_df=sac1,
            target_df=sac2,
            extensive_variables=["TOT_POP"],
            intensive_variables=["pct_poverty"],
            spatial_index="target",
        )

    assert_geodataframe_equal(auto, source)
    assert_geodataframe_equal(auto, target)

    with pytest.raises(ValueError):
        area_interpolate(
            source_df=sac1,
            target_df=sac2,
            extensive_variables=["TOT_POP"],
            intensive_variables=["pct_poverty"],
            spatial_index="non-existent",
        )


def test_area_interpolate_parallel(datasets):
    sac1, sac2 = datasets
    with pytest.WARN_VAR_VALS_INF:
        area = area_interpolate(
            source_df=sac1,
            target_df=sac2,
            extensive_variables=["TOT_POP"],
            intensive_variables=["pct_poverty"],
            n_jobs=-1,
        )
    assert_almost_equal(area.TOT_POP.sum(), 1796856, decimal=0)
    assert_almost_equal(area.pct_poverty.sum(), 2140, decimal=0)


def test_area_tables_binning(datasets):
    sac1, sac2 = datasets
    sac1 = sac1.to_crs(4326)
    sac2 = sac2.to_crs(4326)

    auto = _area_tables_binning(source_df=sac1, target_df=sac2, spatial_index="auto")
    source = _area_tables_binning(
        source_df=sac1, target_df=sac2, spatial_index="source"
    )
    target = _area_tables_binning(
        source_df=sac1, target_df=sac2, spatial_index="target"
    )

    assert (auto != source).sum() == 0
    assert (auto != target).sum() == 0

    assert auto.sum() == pytest.approx(1.3879647)
    assert auto.mean() == pytest.approx(2.7552649e-05)

    assert (auto[5][0].toarray() > 0).sum() == 7


def test_passed_table(datasets):
    sac1, sac2 = datasets
    csr = _area_tables_binning(source_df=sac1, target_df=sac2, spatial_index="auto")

    with pytest.WARN_VAR_VALS_INF:
        area = area_interpolate(
            source_df=sac1,
            target_df=sac2,
            extensive_variables=["TOT_POP"],
            intensive_variables=["pct_poverty"],
            table=csr,
        )
    assert_almost_equal(area.TOT_POP.sum(), 1796856, decimal=0)
    assert_almost_equal(area.pct_poverty.sum(), 2140, decimal=0)

    dok = csr.todok()

    with pytest.WARN_VAR_VALS_INF:
        area = area_interpolate(
            source_df=sac1,
            target_df=sac2,
            extensive_variables=["TOT_POP"],
            intensive_variables=["pct_poverty"],
            table=dok,
        )
    assert_almost_equal(area.TOT_POP.sum(), 1796856, decimal=0)
    assert_almost_equal(area.pct_poverty.sum(), 2140, decimal=0)


@pytest.mark.parametrize(
    ("variable_kind", "fill_nan", "warning_fill", "expected"),
    [
        ("extensive", 0.0, 0.0, np.array([1.0, 0.0])),
        ("extensive", 5.0, 5.0, np.array([1.0, 5.0])),
        ("extensive", "mean", 1.0, np.array([1.0, 1.0])),
        ("extensive", "median", 1.0, np.array([1.0, 1.0])),
        ("extensive", "max", 1.0, np.array([1.0, 1.0])),
        ("extensive", "min", 1.0, np.array([1.0, 1.0])),
        ("intensive", 0.0, 0.0, np.array([1.0, 0.0])),
        ("intensive", 5.0, 5.0, np.array([1.0, 5.0])),
        ("intensive", "mean", 1.0, np.array([1.0, 1.0])),
        ("intensive", "median", 1.0, np.array([1.0, 1.0])),
        ("intensive", "max", 1.0, np.array([1.0, 1.0])),
        ("intensive", "min", 1.0, np.array([1.0, 1.0])),
    ],
)
def test_area_interpolate_fill_nan_strategies(
    variable_kind, fill_nan, warning_fill, expected
):
    source = geopandas.GeoDataFrame(
        {"value": [1.0, np.nan]},
        geometry=[box(0, 0, 1, 1), box(1, 0, 2, 1)],
    )
    target = geopandas.GeoDataFrame(
        geometry=[box(0, 0, 1, 1), box(1, 0, 2, 1)],
    )
    table = csr_matrix(np.eye(2))

    kwargs = {f"{variable_kind}_variables": ["value"]}
    with pytest.warns(
        UserWarning,
        match=f"nan values in variable: value, replacing with {warning_fill}",
    ):
        result = area_interpolate(
            source_df=source,
            target_df=target,
            table=table,
            fill_nan=fill_nan,
            **kwargs,
        )

    np.testing.assert_allclose(result["value"].to_numpy(), expected)


@pytest.mark.parametrize("variable_kind", ["extensive", "intensive"])
def test_area_interpolate_fill_nan_none_propagates(variable_kind):
    source = geopandas.GeoDataFrame(
        {"value": [1.0, np.nan]},
        geometry=[box(0, 0, 1, 1), box(1, 0, 2, 1)],
    )
    target = geopandas.GeoDataFrame(
        geometry=[box(0, 0, 1, 1), box(1, 0, 2, 1)],
    )
    table = csr_matrix(np.eye(2))

    kwargs = {f"{variable_kind}_variables": ["value"]}
    with pytest.warns(
        UserWarning, match="nan values in variable: value, replacing with None"
    ):
        result = area_interpolate(
            source_df=source,
            target_df=target,
            table=table,
            fill_nan=None,
            **kwargs,
        )

    assert result["value"].isna().tolist() == [False, True]


@pytest.mark.parametrize("variable_kind", ["extensive", "intensive"])
def test_area_interpolate_fill_nan_invalid_strategy(variable_kind):
    source = geopandas.GeoDataFrame(
        {"value": [1.0, np.nan]},
        geometry=[box(0, 0, 1, 1), box(1, 0, 2, 1)],
    )
    target = geopandas.GeoDataFrame(
        geometry=[box(0, 0, 1, 1), box(1, 0, 2, 1)],
    )
    table = csr_matrix(np.eye(2))

    kwargs = {f"{variable_kind}_variables": ["value"]}
    with pytest.raises(
        ValueError, match="fill_value should be either None, a numeric value"
    ):
        area_interpolate(
            source_df=source,
            target_df=target,
            table=table,
            fill_nan="mode",
            **kwargs,
        )
