"""
For completeness, we need to test a shuffled dataframe
(i.e. always send unsorted data) with:
- numeric ids
- string ids
- point dataframe
- coordinates
- check two kernel functions
- check two tree types
- scikit/no scikit
"""

import geodatasets
import geopandas
import numpy as np
import pandas as pd
import pytest
from shapely import Point

from libpysal.graph import Graph
from libpysal.graph._kernel import (
    HAS_SKLEARN,
    _distance_band,
    _kernel,
    _kernel_functions,
)
from libpysal.graph._utils import CoplanarError


def _grocs():
    grocs = geopandas.read_file(geodatasets.get_path("geoda groceries"))[
        ["OBJECTID", "geometry"]
    ].explode(ignore_index=True)
    grocs["strID"] = grocs.OBJECTID.astype(str)
    grocs["intID"] = grocs.OBJECTID.values
    return grocs


@pytest.fixture(scope="session")
def grocs():
    return _grocs()


kernel_functions = list(_kernel_functions.keys())


def my_kernel(distances, bandwidth):
    output = np.cos(distances / distances.max())
    output[distances < bandwidth] = 0
    return output


kernel_functions.append(my_kernel)

np.random.seed(6301)
# create a 2-d laplace distribution as a "degenerate"
# over-concentrated distribution
# and rescale to match the lenght-scale in groceries
lap_coords = np.random.laplace(size=(200, 2)) / 50

# create a 2-d cauchy as a "degenerate"
# spatial outlier-y distribution
cau_coords = np.random.standard_cauchy(size=(200, 2))


@pytest.fixture(
    scope="session",
    params=[pytest.param("grocs", marks=pytest.mark.network), "lap", "cau"],
)
def data(request):
    if request.param == "grocs":
        return _grocs()
    elif request.param == "lap":
        return lap_coords
    elif request.param == "cau":
        return cau_coords


parametrize_ids = pytest.mark.parametrize("ids", [None, "strID", "intID"])
parametrize_kernelfunctions = pytest.mark.parametrize("kernel", kernel_functions)

# how do we parameterize conditional on sklearn in env?


@pytest.mark.network
@parametrize_ids
def test_neighbors(ids, grocs):
    data = grocs.set_index(ids) if ids else grocs
    head, tail, weight = _kernel(data, bandwidth=5000, kernel="boxcar")
    assert head.shape[0] == 437
    assert tail.shape == head.shape
    assert weight.shape == head.shape
    np.testing.assert_array_equal(pd.unique(head), data.index)
    known = np.linspace(9, 436, 10, dtype=int)
    head_exp = [2, 16, 28, 41, 55, 72, 92, 111, 135, 147]
    if ids:
        head_exp = data.index.values[head_exp]
    np.testing.assert_array_equal(head.values[known], head_exp)
    tail_exp = [13, 92, 66, 31, 33, 73, 16, 103, 9, 147]
    if ids:
        tail_exp = data.index.values[tail_exp]
    np.testing.assert_array_equal(tail.values[known], tail_exp)
    np.testing.assert_array_equal(
        weight[known],
        np.array([1, 1, 1, 1, 1, 1, 1, 1, 1, 0], dtype="int64"),
    )


def test_no_taper(data):
    head, tail, weight = _kernel(data, taper=False)
    assert head.shape[0] == len(data) * (len(data) - 1)
    assert tail.shape == head.shape
    assert weight.shape == head.shape
    if hasattr(data, "index"):
        np.testing.assert_array_equal(np.unique(head), data.index)
    else:
        np.testing.assert_array_equal(np.unique(head), np.arange(len(data)))


@pytest.mark.network
@parametrize_ids
def test_ids(ids, grocs):
    data = grocs.set_index(ids) if ids else grocs
    head, tail, _ = _kernel(data)
    np.testing.assert_array_equal(pd.unique(head), data.index)
    assert np.isin(tail, data.index).all()


@pytest.mark.network
def test_distance(grocs):
    _, _, weight = _kernel(grocs, kernel="identity", taper=False)
    known = np.linspace(9, weight.shape[0], 10, dtype=int, endpoint=False)
    np.testing.assert_array_almost_equal(
        weight[known],
        np.array(
            [
                39028.10991144,
                51086.85388002,
                21270.55278224,
                8999.11607504,
                91203.25966722,
                36548.75743352,
                58917.81440314,
                63359.35143896,
                24952.48387721,
                65860.55093353,
            ]
        ),
    )

    _, _, weight = _kernel(lap_coords, kernel="identity", taper=False)
    known = np.linspace(9, weight.shape[0], 10, dtype=int, endpoint=False)
    np.testing.assert_array_almost_equal(
        weight[known],
        np.array(
            [
                0.0112305,
                0.03158595,
                0.06027445,
                0.01274032,
                0.07559474,
                0.02240698,
                0.07024776,
                0.05554498,
                0.06012029,
                0.03836196,
            ]
        ),
    )

    _, _, weight = _kernel(cau_coords, kernel="identity", taper=False)
    known = np.linspace(9, weight.shape[0], 10, dtype=int, endpoint=False)
    np.testing.assert_array_almost_equal(
        weight[known],
        np.array(
            [
                5.16946327,
                2.29669606,
                1.60345402,
                1.23831204,
                19.2459334,
                2.99545291,
                2.33316209,
                2.09711302,
                2.12594958,
                6.19347592,
            ]
        ),
    )


def test_k(data):
    head, tail, weight = _kernel(data, k=3, kernel="identity")
    assert head.shape[0] == data.shape[0] * 3
    assert tail.shape == head.shape
    assert weight.shape == head.shape

    # TODO: what shall be the behaviour when k is set but bandwidth
    # TODO: is too small so weight is zero, eliminating the link with taper?
    # TODO: now the affected focal has less neighbors than K without warning
    # TODO: test it with the code above and default kernel


@pytest.mark.network
@parametrize_kernelfunctions
def test_kernels(kernel, grocs):
    _, _, weight = _kernel(grocs, kernel=kernel, taper=False)
    if kernel == "triangular":
        assert weight.mean() == pytest.approx(0.09416822598434019)
        assert weight.max() == pytest.approx(0.9874476868358023)
    elif kernel == "parabolic":
        assert weight.mean() == pytest.approx(0.10312196315841769)
        assert weight.max() == pytest.approx(0.749881829575671)
    elif kernel == "gaussian":
        assert weight.mean() == pytest.approx(0.13787969156713978)
        assert weight.max() == pytest.approx(0.39891085285421685)
    elif kernel == "bisquare":
        assert weight.mean() == pytest.approx(0.09084085210598618)
        assert weight.max() == pytest.approx(0.9372045972129259)
    elif kernel == "tricube":
        assert weight.mean() == pytest.approx(0.0925592973486846)
        assert weight.max() == pytest.approx(0.8641924033756005)
    elif kernel == "cosine":
        assert weight.mean() == pytest.approx(0.1008306468068958)
        assert weight.max() == pytest.approx(0.7852455006403666)
    elif kernel in ["boxcar", "discrete"]:
        assert weight.mean() == pytest.approx(0.2499540356683214)
        assert weight.max() == 1
    elif kernel in ["identity", None]:
        assert weight.mean() == pytest.approx(39758.007361814016)
        assert weight.max() == pytest.approx(127937.75271993055)
    elif kernel in ["exponential", None]:
        assert weight.mean() == pytest.approx(0.25104208195691335)
        assert weight.max() == pytest.approx(0.9875261386315732)
    else:  # function
        assert weight.mean() == pytest.approx(0.6880384553732511)
        assert weight.max() == pytest.approx(0.9855481738848647)


@pytest.mark.parametrize("bandwidth", [None, 0.05, 0.4])
def test_bandwidth(data, bandwidth):
    head, tail, weight = _kernel(data, bandwidth=bandwidth)
    assert tail.shape == head.shape
    assert weight.shape == head.shape
    if hasattr(data, "index"):
        np.testing.assert_array_equal(np.unique(head), data.index)
    else:
        np.testing.assert_array_equal(np.unique(head), np.arange(len(data)))


@pytest.mark.network
@pytest.mark.parametrize(
    "metric",
    [
        "euclidean",
        "minkowski",
        "cityblock",
        "chebyshev",
        "haversine",
    ],
)
def test_metric(metric, grocs):
    pytest.importorskip("pyproj")
    data = grocs.to_crs(4326) if metric == "haversine" else grocs
    if not HAS_SKLEARN and metric in ["chebyshev", "haversine"]:
        pytest.skip("metric not supported by scipy")
    head, tail, weight = _kernel(
        data, metric=metric, kernel="identity", p=1.5, taper=False
    )
    assert head.shape[0] == len(data) * (len(data) - 1)
    assert tail.shape == head.shape
    assert weight.shape == head.shape
    np.testing.assert_array_equal(pd.unique(head), data.index)

    if metric == "euclidean":
        assert weight.mean() == pytest.approx(39758.007362)
        assert weight.max() == pytest.approx(127937.75272)
    elif metric == "minkowski":
        assert weight.mean() == pytest.approx(42288.642129)
        assert weight.max() == pytest.approx(140674.095752)
    elif metric == "cityblock":
        assert weight.mean() == pytest.approx(49424.576155)
        assert weight.max() == pytest.approx(173379.431622)
    elif metric == "chebyshev":
        assert weight.mean() == pytest.approx(36590.352895)
        assert weight.max() == pytest.approx(123955.14249)
    else:
        assert weight.mean() == pytest.approx(0.115835)
        assert weight.max() == pytest.approx(0.371465)


@pytest.mark.network
@pytest.mark.parametrize(
    "metric",
    [
        "euclidean",
        "minkowski",
        "cityblock",
        "chebyshev",
        "haversine",
    ],
)
def test_metric_k(metric, grocs):
    pytest.importorskip("pyproj")
    data = grocs.to_crs(4326) if metric == "haversine" else grocs
    if not HAS_SKLEARN and metric in ["chebyshev", "haversine"]:
        pytest.skip("metric not supported by scipy")
    head, tail, weight = _kernel(data, k=3, metric=metric, kernel="identity", p=1.5)
    assert head.shape[0] == len(data) * 3
    assert tail.shape == head.shape
    assert weight.shape == head.shape
    np.testing.assert_array_equal(pd.unique(head), data.index)

    if metric == "euclidean":
        assert weight.mean() == pytest.approx(4577.237441)
        assert weight.max() == pytest.approx(18791.085051)
    elif metric == "minkowski":
        assert weight.mean() == pytest.approx(4884.254721)
        assert weight.max() == pytest.approx(20681.125211)
    elif metric == "cityblock":
        assert weight.mean() == pytest.approx(5665.288523)
        assert weight.max() == pytest.approx(23980.893147)
    elif metric == "chebyshev":
        assert weight.mean() == pytest.approx(4032.283559)
        assert weight.max() == pytest.approx(16374.141739)
    else:
        assert weight.mean() == pytest.approx(0.00021882448)
        assert weight.max() == pytest.approx(0.000897441)


# def test_precomputed(data, ids):
#     raise NotImplementedError()


@pytest.mark.network
def test_coplanar(grocs):
    grocs_duplicated = pd.concat(
        [grocs, grocs.iloc[:10], grocs.iloc[:3]], ignore_index=True
    )
    # plain kernel
    head, tail, weight = _kernel(grocs_duplicated, taper=False)
    assert head.shape[0] == len(grocs_duplicated) * (len(grocs_duplicated) - 1)
    assert tail.shape == head.shape
    assert weight.shape == head.shape
    np.testing.assert_array_equal(pd.unique(head), grocs_duplicated.index)

    # k, raise
    with pytest.raises(CoplanarError, match="There are"):
        _kernel(grocs_duplicated, k=2)

    # k, jitter
    head, tail, weight = _kernel(grocs_duplicated, taper=False, k=2, coplanar="jitter")
    assert head.shape[0] == len(grocs_duplicated) * 2
    assert tail.shape == head.shape
    assert weight.shape == head.shape
    np.testing.assert_array_equal(pd.unique(head), grocs_duplicated.index)

    # k, clique
    head, tail, weight = _kernel(grocs_duplicated, k=2, coplanar="clique", taper=False)
    assert head.shape[0] >= len(grocs_duplicated) * 2
    assert tail.shape == head.shape
    assert weight.shape == head.shape
    np.testing.assert_array_equal(pd.unique(head), grocs_duplicated.index)


def test_coplanar_clique_duplicate_labels():
    gs = geopandas.GeoSeries(
        [
            Point(0, 0),
            Point(0, 0),
            Point(1, 0),
            Point(2, 0),
            Point(3, 0),
            Point(4, 0),
        ]
    )
    g = Graph.build_kernel(gs, k=2, coplanar="clique")
    assert g.n == 6

    assert 0 in g.adjacency.index.get_level_values("focal")
    assert 1 in g.adjacency.index.get_level_values("focal")

    neighbors_0 = set(g.adjacency.loc[0].index) - {1}
    neighbors_1 = set(g.adjacency.loc[1].index) - {0}
    assert neighbors_0 == neighbors_1


def test_shape_preservation():
    coordinates = np.vstack(
        [np.repeat(np.arange(10), 10), np.tile(np.arange(10), 10)]
    ).T
    head, tail, weight = _kernel(
        coordinates,
        k=3,
        metric="euclidean",
        p=2,
        kernel="boxcar",
        bandwidth=0.5,
        taper=False,
    )
    np.testing.assert_array_equal(head, np.repeat(np.arange(100), 3))
    assert tail.shape == head.shape, "shapes of head and tail do not match"
    np.testing.assert_array_equal(weight, np.zeros((300,), dtype=int))

    head, tail, weight = _kernel(
        coordinates,
        k=3,
        metric="euclidean",
        p=2,
        kernel="boxcar",
        bandwidth=0.5,
        taper=True,
    )
    np.testing.assert_array_equal(head, np.arange(100))
    assert tail.shape == head.shape, "shapes of head and tail do not match"
    np.testing.assert_array_equal(weight, np.zeros((100,), dtype=int))


@pytest.mark.network
def test_haversine_check(grocs):
    with pytest.raises(ValueError, match="'haversine'"):
        _kernel(grocs, k=2, metric="haversine")


def test_distance_band_colocated():
    coordinates = np.array([[0, 0], [1, 0], [1, 0], [2, 0], [3, 0]])
    dist = _distance_band(coordinates, 1)
    assert dist.shape == (5, 5)
    np.testing.assert_array_equal(
        dist.data,
        np.array(
            [
                0.0,
                1.0,
                1.0,
                1.0,
                0.0,
                0.0,
                1.0,
                1.0,
                0.0,
                0.0,
                1.0,
                1.0,
                1.0,
                0.0,
                1.0,
                1.0,
                0.0,
            ]
        ),
    )


def test__kernel_precomputed_exclude_self_weights_affects_ranking():
    coords = np.array(
        [
            [0.1, 1.0, 2.0],
            [1.0, 0.1, 3.0],
            [2.0, 3.0, 0.1],
        ],
        dtype=float,
    )
    n = coords.shape[0]
    k = 2

    out_including_self = _kernel(
        coords,
        k=k,
        metric="precomputed",
        exclude_self_weights=False,
    )
    focal, neighbor, weight = out_including_self

    assert focal.shape == (n * k,)
    assert neighbor.shape == (n * k,)
    assert weight.shape == (n * k,)

    for i in range(n):
        neigh_i = set(neighbor[focal == i].tolist())
        assert i in neigh_i

    out_excluding_self = _kernel(
        coords,
        k=k,
        metric="precomputed",
        exclude_self_weights=True,
    )

    focal2, neighbor2, weight2 = out_excluding_self

    assert focal2.shape == (n * k,)
    assert neighbor2.shape == (n * k,)
    assert weight2.shape == (n * k,)

    # No (i -> i) edges at all
    assert not np.any(focal2 == neighbor2)

    for i in range(n):
        neigh_i = neighbor2[focal2 == i]
        assert len(neigh_i) == k
        assert len(set(neigh_i.tolist())) == k


def test_tree_parameter_knn(grocs):
    """Test that passing a pre-built tree produces the same results."""
    from scipy import spatial

    coords = np.array([[pt.x, pt.y] for pt in grocs.geometry.values])

    # Build KNN without tree
    g1 = _kernel(coords, k=5, kernel="boxcar", bandwidth=np.inf)

    # Build a KDTree and pass it — coords + tree triggers a warning
    tree = spatial.KDTree(coords)
    with pytest.warns(UserWarning, match="coordinate"):
        g2 = _kernel(coords, k=5, kernel="boxcar", bandwidth=np.inf, tree=tree)

    np.testing.assert_array_equal(g1[0], g2[0])  # focal
    np.testing.assert_array_equal(g1[1], g2[1])  # neighbor
    np.testing.assert_array_equal(g1[2], g2[2])  # weight


@pytest.mark.skipif(not HAS_SKLEARN, reason="scikit-learn not installed")
def test_tree_parameter_sklearn(grocs):
    """Test that passing a pre-built sklearn tree works."""
    from sklearn.neighbors import KDTree

    coords = np.array([[pt.x, pt.y] for pt in grocs.geometry.values])

    # Build KNN with sklearn tree — coords + tree triggers a warning
    tree = KDTree(coords)
    with pytest.warns(UserWarning, match="coordinate"):
        g = _kernel(coords, k=5, kernel="boxcar", bandwidth=np.inf, tree=tree)

    # Should produce valid output
    assert len(g[0]) > 0
    assert len(g[1]) > 0
    assert len(g[2]) > 0


def test_tree_parameter_distance_band(grocs):
    """Test that passing a pre-built tree to distance_band works."""
    from scipy import spatial

    coords = np.array([[pt.x, pt.y] for pt in grocs.geometry.values])

    # Build distance band without tree
    sp1 = _distance_band(coords, threshold=500)

    # Build a KDTree and pass it
    tree = spatial.KDTree(coords)
    sp2 = _distance_band(coords, threshold=500, tree=tree)

    np.testing.assert_array_equal(sp1.toarray(), sp2.toarray())


def test_tree_parameter_full_distance_matrix():
    """Test _kernel with k=None and a pre-built tree (uses pdist path)."""
    from scipy import spatial

    coords = lap_coords[:20]

    # Without tree
    g1 = _kernel(coords, k=None, kernel="identity", taper=False)

    # With tree — coords + tree triggers a warning
    tree = spatial.KDTree(coords)
    with pytest.warns(UserWarning, match="coordinate"):
        g2 = _kernel(coords, k=None, kernel="identity", taper=False, tree=tree)

    np.testing.assert_array_almost_equal(np.sort(g1[2]), np.sort(g2[2]))


def test_tree_jitter_raises():
    """Tree + coplanar='jitter' should raise a ValueError."""
    from scipy import spatial

    coords = lap_coords[:20]
    tree = spatial.KDTree(coords)

    with pytest.raises(ValueError, match="Cannot use a pre-built tree"):
        _kernel(coords, k=3, coplanar="jitter", tree=tree)


@pytest.mark.network
def test_tree_haversine_raises(grocs):
    """Tree + metric='haversine' should raise a ValueError."""
    from scipy import spatial

    grocs_4326 = grocs.to_crs(4326)
    coords = np.array([[pt.x, pt.y] for pt in grocs_4326.geometry.values])
    tree = spatial.KDTree(coords)

    with pytest.raises(ValueError, match="Cannot use a pre-built tree"):
        _kernel(coords, k=3, metric="haversine", tree=tree)


def test_adaptive_bandwidth_requires_k():
    """bandwidth='adaptive' without k should raise a ValueError."""
    with pytest.raises(ValueError, match="bandwidth='adaptive'"):
        _kernel(lap_coords, bandwidth="adaptive", k=None)


def test_adaptive_bandwidth_basic():
    """bandwidth='adaptive' with k should produce per-observation bandwidths."""
    k = 5
    head_f, tail_f, weight_f = _kernel(lap_coords, k=k)
    head_a, tail_a, weight_a = _kernel(lap_coords, k=k, bandwidth="adaptive")

    # Both should have the same neighbors
    assert head_a.shape == head_f.shape
    assert tail_a.shape == tail_f.shape

    # each focal should have exactly k neighbors
    unique, counts = np.unique(head_a, return_counts=True)
    assert (counts == k).all()

    assert weight_a.mean() == pytest.approx(0.2950555038)
    assert weight_a.max() == pytest.approx(0.3981074891)

    assert not np.allclose(weight_a, weight_f)


@pytest.mark.parametrize(
    "kernel",
    ["bisquare", "boxcar", "triangular", "tricube", "cosine", "parabolic"],
)
def test_compact_support_sparse_path_matches_dense(kernel):
    """Compact-support kernels with a fixed bandwidth should take the sparse path
    and return results identical to the dense path."""
    coords = lap_coords[:100]
    bw = 0.05

    sparse_heads, sparse_tails, sparse_weights = _kernel(
        coords, bandwidth=bw, kernel=kernel, metric="euclidean"
    )
    # Force dense path: wrap bw so isinstance(bandwidth, (int, float)) is False
    dense_heads, dense_tails, dense_weights = _kernel(
        coords, bandwidth=complex(bw), kernel=kernel, metric="euclidean"
    )

    np.testing.assert_array_equal(
        np.sort(sparse_heads), np.sort(dense_heads), err_msg=f"{kernel}: heads differ"
    )
    np.testing.assert_array_equal(
        np.sort(sparse_tails), np.sort(dense_tails), err_msg=f"{kernel}: tails differ"
    )
    np.testing.assert_array_almost_equal(
        np.sort(sparse_weights),
        np.sort(dense_weights),
        err_msg=f"{kernel}: weights differ",
    )


def test_gaussian_not_sparse_path():
    """Gaussian without taper must not take the sparse path."""

    coords = lap_coords[:50]
    bw = 0.05

    heads, tails, weights = _kernel(
        coords, bandwidth=bw, kernel="gaussian", taper=False
    )
    # Gaussian with a tight bandwidth should still return weights for all pairs
    # (non-zero gaussian never reaches exactly zero), so the edge count should
    # exceed what a compact-support kernel would return at the same bandwidth.
    compact_heads, _, _ = _kernel(coords, bandwidth=bw, kernel="bisquare")
    assert len(heads) >= len(compact_heads), (
        "Gaussian should have at least as many edges as a compact kernel"
        " at the same bandwidth"
    )


def test_gaussian_taper_sparse_path_matches_dense():
    """Gaussian with taper=True should take the sparse path and match the dense path."""
    coords = lap_coords[:100]
    bw = 0.05

    sparse_heads, sparse_tails, sparse_weights = _kernel(
        coords, bandwidth=bw, kernel="gaussian", taper=True, metric="euclidean"
    )
    # Force dense path by wrapping bw so isinstance(bandwidth, (int, float)) is False
    dense_heads, dense_tails, dense_weights = _kernel(
        coords, bandwidth=complex(bw), kernel="gaussian", taper=True, metric="euclidean"
    )

    np.testing.assert_array_equal(np.sort(sparse_heads), np.sort(dense_heads))
    np.testing.assert_array_equal(np.sort(sparse_tails), np.sort(dense_tails))
    np.testing.assert_array_almost_equal(
        np.sort(sparse_weights), np.sort(dense_weights)
    )
