"""
Unit tests of Graph Element.
"""

import numpy as np
import pandas as pd
import pytest

import holoviews as hv
from holoviews.element.util import circular_layout, connect_edges, connect_edges_pd
from holoviews.testing import assert_data_equal, assert_element_equal

from ..utils import optional_dependencies

nx, nx_skip = optional_dependencies("networkx")
scipy, scipy_skip = optional_dependencies("scipy")


class GraphTests:
    def setup_method(self):
        N = 8
        self.nodes = circular_layout(np.arange(N))
        self.source = np.arange(N, dtype=np.int32)
        self.target = np.zeros(N, dtype=np.int32)
        self.edge_info = np.arange(N)
        self.graph = hv.Graph(((self.source, self.target),))

    def test_basic_constructor(self):
        graph = hv.Graph(((self.source, self.target),))
        nodes = hv.Nodes(self.nodes)
        assert_element_equal(graph.nodes, nodes)

    def test_constructor_with_nodes(self):
        graph = hv.Graph(((self.source, self.target), self.nodes))
        nodes = hv.Nodes(self.nodes)
        assert_element_equal(graph.nodes, nodes)

    def test_graph_edge_segments(self):
        segments = connect_edges(self.graph)
        paths = []
        nodes = np.column_stack(self.nodes)
        for start, end in zip(nodes[self.source], nodes[self.target], strict=True):
            paths.append(np.array([start[:2], end[:2]]))
        np.testing.assert_equal(segments, paths)

    def test_graph_node_info_no_index(self):
        node_info = hv.Dataset(np.arange(8), vdims=["Label"])
        graph = hv.Graph(((self.source, self.target), node_info))
        assert_data_equal(graph.nodes.dimension_values(3), node_info.dimension_values(0))

    def test_graph_node_info_no_index_mismatch(self):
        node_info = hv.Dataset(np.arange(6), vdims=["Label"])
        with pytest.raises(ValueError):  # noqa: PT011
            hv.Graph(((self.source, self.target), node_info))

    def test_graph_node_info_merge_on_index(self):
        node_info = hv.Dataset((np.arange(8), np.arange(1, 9)), "index", "label")
        graph = hv.Graph(((self.source, self.target), node_info))
        assert_data_equal(graph.nodes.dimension_values(3), node_info.dimension_values(1))

    def test_graph_node_info_merge_on_index_partial(self):
        node_info = hv.Dataset((np.arange(5), np.arange(1, 6)), "index", "label")
        graph = hv.Graph(((self.source, self.target), node_info))
        expected = np.array([1.0, 2.0, 3.0, 4.0, 5.0, np.nan, np.nan, np.nan])
        assert_data_equal(graph.nodes.dimension_values(3), expected)

    def test_graph_edge_segments_pd(self):
        segments = connect_edges_pd(self.graph)
        paths = []
        nodes = np.column_stack(self.nodes)
        for start, end in zip(nodes[self.source], nodes[self.target], strict=True):
            paths.append(np.array([start[:2], end[:2]]))
        np.testing.assert_equal(segments, paths)

    def test_constructor_with_nodes_and_paths(self):
        paths = hv.Graph(((self.source, self.target), self.nodes)).edgepaths
        graph = hv.Graph(((self.source, self.target), self.nodes, paths.data))
        assert_element_equal(graph._edgepaths, paths)

    def test_constructor_with_nodes_and_paths_dimension_mismatch(self):
        paths = hv.Graph(((self.source, self.target), self.nodes)).edgepaths
        exception = (
            "Ensure that the first two key dimensions on Nodes and EdgePaths match: x != x2"
        )
        with pytest.raises(ValueError, match=exception):
            hv.Graph(((self.source, self.target), self.nodes, paths.redim(x="x2")))

    def test_graph_clone_static_plot_id(self):
        assert self.graph.clone()._plot_id == self.graph._plot_id

    def test_select_by_node_in_edges_selection_mode(self):
        graph = hv.Graph(((self.source, self.target),))
        selection = hv.Graph(([(1, 0), (2, 0)], list(zip(*self.nodes, strict=True))[0:3]))
        assert_element_equal(graph.select(index=(1, 3)), selection)

    def test_select_by_node_in_nodes_selection_mode(self):
        graph = hv.Graph(((self.source, self.source + 1), self.nodes))
        selection = hv.Graph(([(1, 2)], list(zip(*self.nodes, strict=True))[1:3]))
        assert_element_equal(graph.select(index=(1, 3), selection_mode="nodes"), selection)

    def test_select_by_source(self):
        graph = hv.Graph(((self.source, self.target),))
        selection = hv.Graph(([(0, 0), (1, 0)], list(zip(*self.nodes, strict=True))[:2]))
        assert_element_equal(graph.select(start=(0, 2)), selection)

    def test_select_by_target(self):
        graph = hv.Graph(((self.source, self.target),))
        selection = hv.Graph(([(0, 0), (1, 0)], list(zip(*self.nodes, strict=True))[:2]))
        assert_element_equal(graph.select(start=(0, 2)), selection)

    def test_select_by_source_and_target(self):
        graph = hv.Graph(((self.source, self.source + 1), self.nodes))
        selection = hv.Graph(([(0, 1)], list(zip(*self.nodes, strict=True))[:2]))
        assert_element_equal(graph.select(start=(0, 3), end=1), selection)

    def test_select_by_edge_data(self):
        graph = hv.Graph(((self.target, self.source, self.edge_info),), vdims=["info"])
        selection = hv.Graph(
            ([(0, 0, 0), (0, 1, 1)], list(zip(*self.nodes, strict=True))[:2]), vdims=["info"]
        )
        assert_element_equal(graph.select(info=(0, 2)), selection)

    def test_graph_node_range(self):
        graph = hv.Graph(((self.target, self.source),))
        assert graph.range("x") == (-1, 1)
        assert graph.range("y") == (-1, 1)

    def test_graph_redim_nodes(self):
        graph = hv.Graph(((self.target, self.source),))
        redimmed = graph.redim(x="x2", y="y2")
        assert_element_equal(redimmed.nodes, graph.nodes.redim(x="x2", y="y2"))
        assert_element_equal(redimmed.edgepaths, graph.edgepaths.redim(x="x2", y="y2"))


@scipy_skip
class TestFromSparse:
    @pytest.mark.parametrize("sparse_func", ["coo_array", "coo_matrix"])
    @pytest.mark.parametrize(
        "graph_kwargs",
        [
            {},
            {"kdims": ["source", "target"]},
            {"vdims": ["weight"]},
        ],
    )
    def test_from_sparse(self, sparse_func, graph_kwargs):
        import scipy.sparse as sp

        row = np.array([0, 0, 1, 2])
        col = np.array([1, 2, 2, 0])
        data = np.array([1.0, 2.0, 3.0, 4.0])
        sparse_data = getattr(sp, sparse_func)((data, (row, col)), shape=(3, 3))
        nodes_data = {"x": [0.0, 1.0, 0.5], "y": [0.0, 0.0, 1.0], "index": [0, 1, 2]}
        graph = hv.Graph.from_sparse(sparse_data, nodes_data, **graph_kwargs)

        if kdims := graph_kwargs.get("kdims"):
            np.testing.assert_array_equal(graph.dimension_values(kdims[0]), row)
            np.testing.assert_array_equal(graph.dimension_values(kdims[1]), col)
        else:
            np.testing.assert_array_equal(graph.dimension_values("start"), row)
            np.testing.assert_array_equal(graph.dimension_values("end"), col)

        if vdims := graph_kwargs.get("vdims"):
            np.testing.assert_array_equal(graph.dimension_values(vdims[0]), data)
        else:
            np.testing.assert_array_equal(graph.dimension_values("data"), data)

    def test_from_sparse_invalid_input(self):
        regular_array = np.array([[1, 2], [3, 4]])
        nodes_data = {"x": [0.0, 1.0], "y": [0.0, 1.0], "index": [0, 1]}

        with pytest.raises(TypeError, match=r"edges expected to be a scipy.sparse array"):
            hv.Graph.from_sparse(regular_array, nodes_data)


@nx_skip
class FromNetworkXTests:
    def test_from_networkx_with_node_attrs(self):
        G = nx.karate_club_graph()
        graph = hv.Graph.from_networkx(G, nx.circular_layout)
        clubs = np.array(
            [
                "Mr. Hi",
                "Mr. Hi",
                "Mr. Hi",
                "Mr. Hi",
                "Mr. Hi",
                "Mr. Hi",
                "Mr. Hi",
                "Mr. Hi",
                "Mr. Hi",
                "Officer",
                "Mr. Hi",
                "Mr. Hi",
                "Mr. Hi",
                "Mr. Hi",
                "Officer",
                "Officer",
                "Mr. Hi",
                "Mr. Hi",
                "Officer",
                "Mr. Hi",
                "Officer",
                "Mr. Hi",
                "Officer",
                "Officer",
                "Officer",
                "Officer",
                "Officer",
                "Officer",
                "Officer",
                "Officer",
                "Officer",
                "Officer",
                "Officer",
                "Officer",
            ]
        )
        assert_data_equal(graph.nodes.dimension_values("club"), clubs)

    def test_from_networkx_with_invalid_node_attrs(self):
        FG = nx.Graph()
        FG.add_node(1, test=[])
        FG.add_node(2, test=[])
        FG.add_edge(1, 2)
        graph = hv.Graph.from_networkx(FG, nx.circular_layout)
        assert graph.nodes.vdims == []
        assert_data_equal(graph.nodes.dimension_values(2), np.array([1, 2]))
        assert_data_equal(graph.array(), np.array([(1, 2)]))

    def test_from_networkx_with_edge_attrs(self):
        FG = nx.Graph()
        FG.add_weighted_edges_from([(1, 2, 0.125), (1, 3, 0.75), (2, 4, 1.2), (3, 4, 0.375)])
        graph = hv.Graph.from_networkx(FG, nx.circular_layout)
        assert_data_equal(graph.dimension_values("weight"), np.array([0.125, 0.75, 1.2, 0.375]))

    def test_from_networkx_with_invalid_edge_attrs(self):
        FG = nx.Graph()
        FG.add_weighted_edges_from([(1, 2, []), (1, 3, []), (2, 4, []), (3, 4, [])])
        graph = hv.Graph.from_networkx(FG, nx.circular_layout)
        assert graph.vdims == []

    def test_from_networkx_only_nodes(self):
        G = nx.Graph()
        G.add_nodes_from([1, 2, 3])
        graph = hv.Graph.from_networkx(G, nx.circular_layout)
        assert_data_equal(graph.nodes.dimension_values(2), np.array([1, 2, 3]))

    def test_from_networkx_custom_nodes(self):
        FG = nx.Graph()
        FG.add_weighted_edges_from([(1, 2, 0.125), (1, 3, 0.75), (2, 4, 1.2), (3, 4, 0.375)])
        nodes = hv.Dataset([(1, "A"), (2, "B"), (3, "A"), (4, "B")], "index", "some_attribute")
        graph = hv.Graph.from_networkx(FG, nx.circular_layout, nodes=nodes)
        assert_data_equal(
            graph.nodes.dimension_values("some_attribute"), np.array(["A", "B", "A", "B"])
        )

    def test_from_networkx_dictionary_positions(self):
        G = nx.Graph()
        G.add_nodes_from([1, 2, 3])
        positions = nx.circular_layout(G)
        graph = hv.Graph.from_networkx(G, positions)
        assert_data_equal(graph.nodes.dimension_values(2), np.array([1, 2, 3]))


class ChordTests:
    def setup_method(self):
        self.simplices = [(0, 1, 2), (1, 2, 3), (2, 3, 4)]

    def test_chord_constructor_no_vdims(self):
        chord = hv.Chord(self.simplices)
        nodes = np.array(
            [
                [0.8660254037844387, 0.49999999999999994, 0],
                [-0.4999999999999998, 0.8660254037844388, 1],
                [-0.5000000000000004, -0.8660254037844384, 2],
                [0.8660254037844379, -0.5000000000000012, 3],
            ]
        )
        assert_element_equal(chord.nodes, hv.Nodes(nodes))
        assert_data_equal(chord.array(), np.array([s[:2] for s in self.simplices]))

    def test_chord_constructor_with_vdims(self):
        chord = hv.Chord(self.simplices, vdims=["z"])
        nodes = np.array(
            [
                [0.9396926207859084, 0.3420201433256687, 0],
                [6.123233995736766e-17, 1.0, 1],
                [-0.8660254037844388, -0.4999999999999998, 2],
                [0.7660444431189779, -0.6427876096865396, 3],
            ]
        )
        assert_element_equal(chord.nodes, hv.Nodes(nodes))
        assert_data_equal(chord.array(), np.array(self.simplices))

    def test_chord_constructor_self_reference(self):
        chord = hv.Chord([("A", "B", 2), ("B", "A", 3), ("A", "A", 2)])
        nodes = pd.DataFrame(
            [[-0.5, 0.866025, "A"], [0.5, -0.866025, "B"]], columns=["x", "y", "index"]
        )
        assert_element_equal(chord.nodes, hv.Nodes(nodes))


class TriMeshTests:
    def setup_method(self):
        self.nodes = [(0, 0, 0), (0.5, 1, 1), (1.0, 0, 2), (1.5, 1, 4)]
        self.simplices = [(0, 1, 2), (1, 2, 3)]

    def test_trimesh_constructor(self):
        nodes = [n[:2] for n in self.nodes]
        trimesh = hv.TriMesh((self.simplices, nodes))
        assert_data_equal(trimesh.array(), np.array(self.simplices))
        assert_data_equal(trimesh.nodes.array([0, 1]), np.array(nodes))
        assert_data_equal(trimesh.nodes.dimension_values(2), np.arange(4))

    def test_trimesh_empty(self):
        trimesh = hv.TriMesh([])
        assert_data_equal(trimesh.array(), np.empty((0, 3)))
        assert_data_equal(trimesh.nodes.array(), np.empty((0, 3)))

    def test_trimesh_empty_clone(self):
        trimesh = hv.TriMesh([]).clone()
        assert_data_equal(trimesh.array(), np.empty((0, 3)))
        assert_data_equal(trimesh.nodes.array(), np.empty((0, 3)))

    def test_trimesh_constructor_tuple_nodes(self):
        nodes = tuple(zip(*self.nodes, strict=True))[:2]
        trimesh = hv.TriMesh((self.simplices, nodes))
        assert_data_equal(trimesh.array(), np.array(self.simplices))
        assert_data_equal(trimesh.nodes.array([0, 1]), np.array(nodes).T)
        assert_data_equal(trimesh.nodes.dimension_values(2), np.arange(4))

    def test_trimesh_constructor_df_nodes(self):
        nodes_df = pd.DataFrame(self.nodes, columns=["x", "y", "z"])
        trimesh = hv.TriMesh((self.simplices, nodes_df))
        nodes = hv.Nodes([(0, 0, 0, 0), (0.5, 1, 1, 1), (1.0, 0, 2, 2), (1.5, 1, 3, 4)], vdims="z")
        assert_data_equal(trimesh.array(), np.array(self.simplices))
        assert_element_equal(trimesh.nodes, nodes)

    def test_trimesh_constructor_point_nodes(self):
        nodes = [n[:2] for n in self.nodes]
        trimesh = hv.TriMesh((self.simplices, hv.Points(self.nodes)))
        assert_data_equal(trimesh.array(), np.array(self.simplices))
        assert_data_equal(trimesh.nodes.array([0, 1]), np.array(nodes))
        assert_data_equal(trimesh.nodes.dimension_values(2), np.arange(4))

    def test_trimesh_edgepaths(self):
        trimesh = hv.TriMesh((self.simplices, self.nodes))
        paths = [
            np.array(
                [
                    (0, 0),
                    (0.5, 1),
                    (1, 0),
                    (0, 0),
                    (np.nan, np.nan),
                    (0.5, 1),
                    (1, 0),
                    (1.5, 1),
                    (0.5, 1),
                ]
            )
        ]
        for p1, p2 in zip(trimesh.edgepaths.split(datatype="array"), paths, strict=True):
            assert_data_equal(p1, p2)

    def test_trimesh_select(self):
        trimesh = hv.TriMesh((self.simplices, self.nodes)).select(x=(0.1, None))
        assert_data_equal(trimesh.array(), np.array(self.simplices[1:]))


class TestSankey:
    def test_single_edge_sankey(self):
        sankey = hv.Sankey([("A", "B", 1)])
        links = list(sankey._sankey["links"])
        assert len(links) == 1
        del links[0]["source"]["sourceLinks"]
        del links[0]["target"]["targetLinks"]
        link = {
            "index": 0,
            "source": {
                "index": "A",
                "values": (),
                "targetLinks": [],
                "value": 1,
                "depth": 0,
                "height": 1,
                "column": 0,
                "x0": 0,
                "x1": 15,
                "y0": 0.0,
                "y1": 500.0,
            },
            "target": {
                "index": "B",
                "values": (),
                "sourceLinks": [],
                "value": 1,
                "depth": 1,
                "height": 0,
                "column": 1,
                "x0": 985.0,
                "x1": 1000.0,
                "y0": 0.0,
                "y1": 500.0,
            },
            "value": 1,
            "width": 500.0,
            "y0": 250.0,
            "y1": 250.0,
        }
        assert links[0] == link
