# Copyright (c) 2021, NVIDIA CORPORATION.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path

import pytest
import numpy as np
import pandas as pd
import cupy as cp
from scipy.sparse import coo_matrix, csr_matrix

from . import utils

_test_data = {
    "graph1":  # asymmetric
    {
        "input": [
            [0, 1, 1, 0, 0],
            [0, 0, 1, 0, 0],
            [0, 0, 0, 0, 0],
            [0, 0, 0, 0, 1],
            [0, 0, 0, 0, 0],
        ],
        "scc_comp_vertices": [
            [0],
            [1],
            [2],
            [3],
            [4],
        ],
        "wcc_comp_vertices": [
            [0, 1, 2],
            [3, 4],
        ],
    },

    "graph2":  # symmetric
    {
        "input": [
            [0, 1, 1, 0, 0],
            [1, 0, 1, 0, 0],
            [1, 1, 0, 0, 0],
            [0, 0, 0, 0, 1],
            [0, 0, 0, 1, 0],
        ],
        "scc_comp_vertices": [
            [0, 1, 2],
            [3, 4],
        ],
        "wcc_comp_vertices": [
            [0, 1, 2],
            [3, 4],
        ],
    },

    "karate-disjoint-sequential":
    {
        "input":
        utils.RAPIDS_DATASET_ROOT_DIR_PATH/"karate-disjoint-sequential.csv",
        "scc_comp_vertices": [
            [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
             17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31,
             32, 33],
            [34],
            [35],
            [36],
        ],
        "wcc_comp_vertices": [
            [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
             17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31,
             32, 33],
            [34, 35, 36],
        ],
    },

    "dolphins":  # dolphins contains only one component
    {
        "input": utils.RAPIDS_DATASET_ROOT_DIR_PATH/"dolphins.csv",
        "scc_comp_vertices": [
            list(range(62)),
        ],
        "wcc_comp_vertices": [
            list(range(62)),
        ],
    },
}


@pytest.fixture(scope="module",
                params=[pytest.param(value, id=key)
                        for (key, value) in _test_data.items()])
def input_and_expected_output(request):
    """
    This fixture takes the above test data and converts it into everything
    needed to run a pylibcugraph CC algo for a specific input (either a
    adjacency matrix or a CSV edgelist file), and returns it along with the
    expected WCC and SCC result for each.
    """
    d = request.param.copy()
    input = d.pop("input")
    expected_output_dict = d

    if isinstance(input, Path):
        pdf = pd.read_csv(
            input,
            delimiter=" ",
            header=None,
            names=["0", "1", "weight"],
            dtype={"0": "int32", "1": "int32", "weight": "float32"},
        )
        num_verts = len(set(pdf["0"].tolist() + pdf["1"].tolist()))
        num_edges = len(pdf)
        weights = np.ones(num_edges)
        coo = coo_matrix((weights, (pdf["0"], pdf["1"])),
                         shape=(num_verts, num_verts),
                         dtype=np.float32)
        csr = coo.tocsr()
    else:
        csr = csr_matrix(input)
        num_verts = csr.get_shape()[0]
        num_edges = csr.nnz

    offsets = cp.asarray(csr.indptr, dtype=np.int32)
    indices = cp.asarray(csr.indices, dtype=np.int32)
    labels_to_populate = cp.asarray(np.zeros(num_verts, dtype=np.int32))

    return ((offsets, indices, labels_to_populate, num_verts, num_edges),
            expected_output_dict)


def _check_labels(vertex_ordered_labels, expected_vertex_comps):
    """
    vertex_ordered_labels is a list of labels, ordered by the position of the
    vertex ID value, as returned by pylibcugraph.CC algos. For example:
    [9, 9, 7]
    means vertex 0 is labelled 9, vertex 1 is labelled 9, and vertex 2 is
    labelled 7.

    expected_vertex_comps is a list of components, where each component is a
    list of vertex IDs the component contains. Each component corresponds to
    some label For example:
    [[0, 1], [2]]
    is two components, the first
    containing vertices 0, 1, and the other 2.  [0, 1] has the label 9 and [2]
    has the label 7.

    This asserts if the vertex_ordered_labels do not correspond to the
    expected_vertex_comps.
    """
    # Group the vertex_ordered_labels list into components based on labels by
    # creating a dictionary of labels to lists of vertices with that label.
    d = {}
    for (vertex, label) in enumerate(vertex_ordered_labels):
        d.setdefault(label, []).append(vertex)

    assert len(d.keys()) == len(expected_vertex_comps), \
        "number of different labels does not match expected"

    # Compare the actual components (created from the dictionary above) to
    # expected.
    actual_vertex_comps = sorted(d.values())
    assert actual_vertex_comps == sorted(expected_vertex_comps)


###############################################################################
# Tests
def test_import():
    """
    Ensure pylibcugraph is importable.
    """
    # suppress F401 (imported but never used) in flake8
    import pylibcugraph  # noqa: F401


def test_scc(package_under_test, input_and_expected_output):
    """
    Tests strongly_connected_components()
    """
    pylibcugraph = package_under_test
    ((cupy_offsets, cupy_indices, cupy_labels_to_populate,
      num_verts, num_edges),
     expected_output_dict) = input_and_expected_output

    pylibcugraph.strongly_connected_components(
        cupy_offsets,
        cupy_indices,
        None,
        num_verts,
        num_edges,
        cupy_labels_to_populate
    )

    _check_labels(cupy_labels_to_populate.tolist(),
                  expected_output_dict["scc_comp_vertices"])


def test_wcc(package_under_test, input_and_expected_output):
    """
    Tests weakly_connected_components()
    """
    pylibcugraph = package_under_test
    ((cupy_offsets, cupy_indices, cupy_labels_to_populate,
      num_verts, num_edges),
     expected_output_dict) = input_and_expected_output

    pylibcugraph.weakly_connected_components(
        cupy_offsets,
        cupy_indices,
        None,
        num_verts,
        num_edges,
        cupy_labels_to_populate
    )

    _check_labels(cupy_labels_to_populate.tolist(),
                  expected_output_dict["wcc_comp_vertices"])


@pytest.mark.parametrize("api_name", ["strongly_connected_components",
                                      "weakly_connected_components"])
def test_non_CAI_input(package_under_test, api_name):
    """
    Ensures that the *_connected_components() APIs only accepts instances of
    objects that have a __cuda_array_interface__
    """
    pylibcugraph = package_under_test
    cupy_array = cp.ndarray(range(8))
    python_list = list(range(8))
    api = getattr(pylibcugraph, api_name)

    with pytest.raises(TypeError):
        api(src=cupy_array,
            dst=cupy_array,
            weights=cupy_array,  # should raise, weights must be None
            num_verts=2,
            num_edges=8,
            labels=cupy_array)
    with pytest.raises(TypeError):
        api(src=cupy_array,
            dst=python_list,  # should raise, no __cuda_array_interface__
            weights=None,
            num_verts=2,
            num_edges=8,
            labels=cupy_array)
    with pytest.raises(TypeError):
        api(src=python_list,  # should raise, no __cuda_array_interface__
            dst=cupy_array,
            weights=None,
            num_verts=2,
            num_edges=8,
            labels=cupy_array)
    with pytest.raises(TypeError):
        api(src=cupy_array,
            dst=cupy_array,
            weights=None,
            num_verts=2,
            num_edges=8,
            labels=python_list)  # should raise, no __cuda_array_interface__


@pytest.mark.parametrize("api_name", ["strongly_connected_components",
                                      "weakly_connected_components"])
def test_bad_dtypes(package_under_test, api_name):
    """
    Ensures that only supported dtypes are accepted.
    """
    graph = [
        [0, 1, 1, 0, 0],
        [0, 0, 1, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 1],
        [0, 0, 0, 0, 0],
    ]
    scipy_csr = csr_matrix(graph)
    num_verts = scipy_csr.get_shape()[0]
    num_edges = scipy_csr.nnz

    pylibcugraph = package_under_test
    api = getattr(pylibcugraph, api_name)

    cp_offsets = cp.asarray(scipy_csr.indptr)
    cp_indices = cp.asarray(scipy_csr.indices)
    cp_labels = cp.asarray(np.zeros(num_verts,
                                    dtype=np.int64))  # unsupported
    with pytest.raises(TypeError):
        api(offsets=cp_offsets,
            indices=cp_indices,
            weights=None,
            num_verts=num_verts,
            num_edges=num_edges,
            labels=cp_labels)

    cp_offsets = cp.asarray(scipy_csr.indptr,
                            dtype=np.int64)  # unsupported
    cp_indices = cp.asarray(scipy_csr.indices)
    cp_labels = cp.asarray(np.zeros(num_verts),
                           dtype=np.int32)
    with pytest.raises(TypeError):
        api(offsets=cp_offsets,
            indices=cp_indices,
            weights=None,
            num_verts=num_verts,
            num_edges=num_edges,
            labels=cp_labels)

    cp_offsets = cp.asarray(scipy_csr.indptr)
    cp_indices = cp.asarray(scipy_csr.indices,
                            dtype=np.float32)  # unsupported
    cp_labels = cp.asarray(np.zeros(num_verts),
                           dtype=np.int32)
    with pytest.raises(TypeError):
        api(offsets=cp_offsets,
            indices=cp_indices,
            weights=None,
            num_verts=num_verts,
            num_edges=num_edges,
            labels=cp_labels)
