# Copyright (c) 2019-2024, 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.

import gc

import random
import networkx as nx
import pytest

import cudf
import cugraph
from cugraph.testing import utils, UNDIRECTED_DATASETS
from cugraph.datasets import karate_asymmetric
from pylibcugraph.testing.utils import gen_fixture_params_product


# =============================================================================
# Pytest Setup / Teardown - called for each test function
# =============================================================================
def setup_function():
    gc.collect()


# =============================================================================
# Pytest fixtures
# =============================================================================
datasets = UNDIRECTED_DATASETS
fixture_params = gen_fixture_params_product(
    (datasets, "graph_file"),
    ([True, False], "edgevals"),
    ([True, False], "start_list"),
)


@pytest.fixture(scope="module", params=fixture_params)
def input_combo(request):
    """
    This fixture returns a dictionary containing all input params required to
    run a Triangle Count algo
    """
    parameters = dict(zip(("graph_file", "edgevals", "start_list"), request.param))

    graph_file = parameters["graph_file"]
    input_data_path = graph_file.get_path()
    edgevals = parameters["edgevals"]

    G = graph_file.get_graph(ignore_weights=not edgevals)

    Gnx = utils.generate_nx_graph_from_file(
        input_data_path, directed=False, edgevals=edgevals
    )

    parameters["G"] = G
    parameters["Gnx"] = Gnx

    return parameters


# =============================================================================
# Tests
# =============================================================================
@pytest.mark.sg
def test_triangles(input_combo):
    G = input_combo["G"]
    Gnx = input_combo["Gnx"]
    nx_triangle_results = cudf.DataFrame()

    if input_combo["start_list"]:
        # sample k nodes from the nx graph
        k = random.randint(1, 10)
        start_list = random.sample(list(Gnx.nodes()), k)
    else:
        start_list = None

    cugraph_triangle_results = cugraph.triangle_count(G, start_list)

    triangle_results = (
        cugraph_triangle_results.sort_values("vertex")
        .reset_index(drop=True)
        .rename(columns={"counts": "cugraph_counts"})
    )

    dic_results = nx.triangles(Gnx, start_list)
    nx_triangle_results["vertex"] = dic_results.keys()
    nx_triangle_results["counts"] = dic_results.values()
    nx_triangle_results = nx_triangle_results.sort_values("vertex").reset_index(
        drop=True
    )

    triangle_results["nx_counts"] = nx_triangle_results["counts"]
    counts_diff = triangle_results.query("nx_counts != cugraph_counts")
    assert len(counts_diff) == 0


@pytest.mark.sg
def test_triangles_int64(input_combo):
    Gnx = input_combo["Gnx"]
    count_int32 = cugraph.triangle_count(Gnx)["counts"].sum()

    graph_file = input_combo["graph_file"]
    G = graph_file.get_graph()
    G.edgelist.edgelist_df = G.edgelist.edgelist_df.astype(
        {"src": "int64", "dst": "int64"}
    )
    count_int64 = cugraph.triangle_count(G)["counts"].sum()

    assert G.edgelist.edgelist_df["src"].dtype == "int64"
    assert G.edgelist.edgelist_df["dst"].dtype == "int64"
    assert count_int32 == count_int64


@pytest.mark.sg
def test_triangles_no_weights(input_combo):
    G_weighted = input_combo["Gnx"]
    count_triangles_nx_graph = cugraph.triangle_count(G_weighted)["counts"].sum()

    graph_file = input_combo["graph_file"]
    G = graph_file.get_graph(ignore_weights=True)

    assert G.is_weighted() is False
    count_triangles = cugraph.triangle_count(G)["counts"].sum()

    assert count_triangles_nx_graph == count_triangles


@pytest.mark.sg
def test_triangles_directed_graph():
    input_data_path = karate_asymmetric.get_path()
    M = utils.read_csv_for_nx(input_data_path)
    G = cugraph.Graph(directed=True)
    cu_M = cudf.DataFrame()
    cu_M["src"] = cudf.Series(M["0"])
    cu_M["dst"] = cudf.Series(M["1"])

    cu_M["weights"] = cudf.Series(M["weight"])
    G.from_cudf_edgelist(cu_M, source="src", destination="dst", edge_attr="weights")

    with pytest.raises(ValueError):
        cugraph.triangle_count(G)
