import random
from copy import copy

import pytest

import networkx as nx
from networkx.utils import (
    PythonRandomInterface,
    PythonRandomViaNumpyBits,
    arbitrary_element,
    create_py_random_state,
    create_random_state,
    dict_to_numpy_array,
    discrete_sequence,
    edges_equal,
    flatten,
    groups,
    make_list_of_ints,
    pairwise,
    powerlaw_sequence,
)
from networkx.utils.misc import _dict_to_numpy_array1, _dict_to_numpy_array2

nested_depth = (
    1,
    2,
    (3, 4, ((5, 6, (7,), (8, (9, 10), 11), (12, 13, (14, 15)), 16), 17), 18, 19),
    20,
)

nested_set = {
    (1, 2, 3, 4),
    (5, 6, 7, 8, 9),
    (10, 11, (12, 13, 14), (15, 16, 17, 18)),
    19,
    20,
}

nested_mixed = [
    1,
    (2, 3, {4, (5, 6), 7}, [8, 9]),
    {10: "foo", 11: "bar", (12, 13): "baz"},
    {(14, 15): "qwe", 16: "asd"},
    (17, (18, "19"), 20),
]


@pytest.mark.parametrize("result", [None, [], ["existing"], ["existing1", "existing2"]])
@pytest.mark.parametrize("nested", [nested_depth, nested_mixed, nested_set])
def test_flatten(nested, result):
    if result is None:
        val = flatten(nested, result)
        assert len(val) == 20
    else:
        _result = copy(result)  # because pytest passes parameters as is
        nexisting = len(_result)
        val = flatten(nested, _result)
        assert len(val) == len(_result) == 20 + nexisting

    assert issubclass(type(val), tuple)


def test_make_list_of_ints():
    mylist = [1, 2, 3.0, 42, -2]
    assert make_list_of_ints(mylist) is mylist
    assert make_list_of_ints(mylist) == mylist
    assert isinstance(make_list_of_ints(mylist)[2], int)
    pytest.raises(nx.NetworkXError, make_list_of_ints, [1, 2, 3, "kermit"])
    pytest.raises(nx.NetworkXError, make_list_of_ints, [1, 2, 3.1])


def test_random_number_distribution():
    # smoke test only
    z = powerlaw_sequence(20, exponent=2.5)
    z = discrete_sequence(20, distribution=[0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 3])


class TestNumpyArray:
    @classmethod
    def setup_class(cls):
        global np
        np = pytest.importorskip("numpy")

    def test_numpy_to_list_of_ints(self):
        a = np.array([1, 2, 3], dtype=np.int64)
        b = np.array([1.0, 2, 3])
        c = np.array([1.1, 2, 3])
        assert isinstance(make_list_of_ints(a), list)
        assert make_list_of_ints(b) == list(b)
        B = make_list_of_ints(b)
        assert isinstance(B[0], int)
        pytest.raises(nx.NetworkXError, make_list_of_ints, c)

    def test__dict_to_numpy_array1(self):
        d = {"a": 1, "b": 2}
        a = _dict_to_numpy_array1(d, mapping={"a": 0, "b": 1})
        np.testing.assert_allclose(a, np.array([1, 2]))
        a = _dict_to_numpy_array1(d, mapping={"b": 0, "a": 1})
        np.testing.assert_allclose(a, np.array([2, 1]))

        a = _dict_to_numpy_array1(d)
        np.testing.assert_allclose(a.sum(), 3)

    def test__dict_to_numpy_array2(self):
        d = {"a": {"a": 1, "b": 2}, "b": {"a": 10, "b": 20}}

        mapping = {"a": 1, "b": 0}
        a = _dict_to_numpy_array2(d, mapping=mapping)
        np.testing.assert_allclose(a, np.array([[20, 10], [2, 1]]))

        a = _dict_to_numpy_array2(d)
        np.testing.assert_allclose(a.sum(), 33)

    def test_dict_to_numpy_array_a(self):
        d = {"a": {"a": 1, "b": 2}, "b": {"a": 10, "b": 20}}

        mapping = {"a": 0, "b": 1}
        a = dict_to_numpy_array(d, mapping=mapping)
        np.testing.assert_allclose(a, np.array([[1, 2], [10, 20]]))

        mapping = {"a": 1, "b": 0}
        a = dict_to_numpy_array(d, mapping=mapping)
        np.testing.assert_allclose(a, np.array([[20, 10], [2, 1]]))

        a = _dict_to_numpy_array2(d)
        np.testing.assert_allclose(a.sum(), 33)

    def test_dict_to_numpy_array_b(self):
        d = {"a": 1, "b": 2}

        mapping = {"a": 0, "b": 1}
        a = dict_to_numpy_array(d, mapping=mapping)
        np.testing.assert_allclose(a, np.array([1, 2]))

        a = _dict_to_numpy_array1(d)
        np.testing.assert_allclose(a.sum(), 3)


def test_pairwise():
    nodes = range(4)
    node_pairs = [(0, 1), (1, 2), (2, 3)]
    node_pairs_cycle = node_pairs + [(3, 0)]
    assert list(pairwise(nodes)) == node_pairs
    assert list(pairwise(iter(nodes))) == node_pairs
    assert list(pairwise(nodes, cyclic=True)) == node_pairs_cycle
    empty_iter = iter(())
    assert list(pairwise(empty_iter)) == []
    empty_iter = iter(())
    assert list(pairwise(empty_iter, cyclic=True)) == []


def test_groups():
    many_to_one = dict(zip("abcde", [0, 0, 1, 1, 2]))
    actual = groups(many_to_one)
    expected = {0: {"a", "b"}, 1: {"c", "d"}, 2: {"e"}}
    assert actual == expected
    assert {} == groups({})


def test_create_random_state():
    np = pytest.importorskip("numpy")
    rs = np.random.RandomState

    assert isinstance(create_random_state(1), rs)
    assert isinstance(create_random_state(None), rs)
    assert isinstance(create_random_state(np.random), rs)
    assert isinstance(create_random_state(rs(1)), rs)
    # Support for numpy.random.Generator
    rng = np.random.default_rng()
    assert isinstance(create_random_state(rng), np.random.Generator)
    pytest.raises(ValueError, create_random_state, "a")

    assert np.all(rs(1).rand(10) == create_random_state(1).rand(10))


def test_create_py_random_state():
    pyrs = random.Random

    assert isinstance(create_py_random_state(1), pyrs)
    assert isinstance(create_py_random_state(None), pyrs)
    assert isinstance(create_py_random_state(pyrs(1)), pyrs)
    pytest.raises(ValueError, create_py_random_state, "a")

    np = pytest.importorskip("numpy")

    rs = np.random.RandomState
    rng = np.random.default_rng(1000)
    rng_explicit = np.random.Generator(np.random.SFC64())
    old_nprs = PythonRandomInterface
    nprs = PythonRandomViaNumpyBits
    assert isinstance(create_py_random_state(np.random), nprs)
    assert isinstance(create_py_random_state(rs(1)), old_nprs)
    assert isinstance(create_py_random_state(rng), nprs)
    assert isinstance(create_py_random_state(rng_explicit), nprs)
    # test default rng input
    old_nprs_instance = old_nprs()
    nprs_instance = nprs()
    assert isinstance(old_nprs_instance, old_nprs)
    assert isinstance(nprs_instance, nprs)
    assert create_py_random_state(old_nprs_instance) == old_nprs_instance
    assert create_py_random_state(nprs_instance) == nprs_instance

    # VeryLargeIntegers Smoke test (they raise error for np.random)
    int64max = 9223372036854775807  # from np.iinfo(np.int64).max
    for r in (rng, rs(1)):
        prs = create_py_random_state(r)
        prs.randrange(3, int64max + 5)
        prs.randint(3, int64max + 5)


def test_PythonRandomInterface_RandomState():
    np = pytest.importorskip("numpy")

    seed = 42
    rs = np.random.RandomState
    rng = PythonRandomInterface(rs(seed))
    rs42 = rs(seed)

    # make sure these functions are same as expected outcome
    assert rng.randrange(3, 5) == rs42.randint(3, 5)
    assert rng.randrange(2) == rs42.randint(0, 2)
    assert rng.uniform(1, 10) == rs42.uniform(1, 10)
    assert rng.choice([1, 2, 3]) == rs42.choice([1, 2, 3])
    assert rng.gauss(0, 1) == rs42.normal(0, 1)
    assert rng.expovariate(1.5) == rs42.exponential(1 / 1.5)
    assert rng.paretovariate(2) == rs42.pareto(2)
    assert np.all(rng.shuffle([1, 2, 3]) == rs42.shuffle([1, 2, 3]))
    assert np.all(
        rng.sample([1, 2, 3], 2) == rs42.choice([1, 2, 3], (2,), replace=False)
    )
    assert np.all(
        [rng.randint(3, 5) for _ in range(100)]
        == [rs42.randint(3, 6) for _ in range(100)]
    )
    assert rng.random() == rs42.random_sample()


def test_PythonRandomInterface_Generator():
    np = pytest.importorskip("numpy")

    seed = 42
    rng = np.random.default_rng(seed)
    pri = PythonRandomInterface(np.random.default_rng(seed))

    # make sure these functions are same as expected outcome
    assert pri.randrange(3, 5) == rng.integers(3, 5)
    assert pri.randrange(2) == rng.integers(0, 2)
    assert pri.uniform(1, 10) == rng.uniform(1, 10)
    assert pri.choice([1, 2, 3]) == rng.choice([1, 2, 3])
    assert pri.gauss(0, 1) == rng.normal(0, 1)
    assert pri.expovariate(1.5) == rng.exponential(1 / 1.5)
    assert pri.paretovariate(2) == rng.pareto(2)
    assert np.all(pri.shuffle([1, 2, 3]) == rng.shuffle([1, 2, 3]))
    assert np.all(
        pri.sample([1, 2, 3], 2) == rng.choice([1, 2, 3], (2,), replace=False)
    )
    assert np.all(
        [pri.randint(3, 5) for _ in range(100)]
        == [rng.integers(3, 6) for _ in range(100)]
    )
    assert pri.random() == rng.random()


@pytest.mark.parametrize(
    ("iterable_type", "expected"), ((list, 1), (tuple, 1), (str, "["), (set, 1))
)
def test_arbitrary_element(iterable_type, expected):
    iterable = iterable_type([1, 2, 3])
    assert arbitrary_element(iterable) == expected


@pytest.mark.parametrize(
    "iterator",
    ((i for i in range(3)), iter([1, 2, 3])),  # generator
)
def test_arbitrary_element_raises(iterator):
    """Value error is raised when input is an iterator."""
    with pytest.raises(ValueError, match="from an iterator"):
        arbitrary_element(iterator)


@pytest.mark.parametrize("n", [5, 10, 20])
@pytest.mark.parametrize("gen", [nx.complete_graph, nx.path_graph, nx.cycle_graph])
@pytest.mark.parametrize("create_using", [nx.Graph, nx.DiGraph])
def test_edges_equal(n, gen, create_using):
    """Test whether edges_equal properly compares edges without attribute data."""
    G = gen(n, create_using=create_using)
    H = gen(n, create_using=create_using)
    assert edges_equal(G.edges(), H.edges(), directed=G.is_directed())
    assert edges_equal(H.edges(), G.edges(), directed=H.is_directed())

    H.remove_edge(0, 1)
    assert edges_equal(H.edges(), H.edges(), directed=H.is_directed())
    assert not edges_equal(G.edges(), H.edges(), directed=G.is_directed())
    assert not edges_equal(H.edges(), G.edges(), directed=H.is_directed())


@pytest.mark.parametrize("n", [5, 10, 20])
@pytest.mark.parametrize("gen", [nx.complete_graph, nx.path_graph, nx.cycle_graph])
@pytest.mark.parametrize("create_using", [nx.MultiGraph, nx.MultiDiGraph])
def test_edges_equal_multiedge(n, gen, create_using):
    """Test whether ``edges_equal`` properly compares edges in multigraphs."""
    G = gen(n, create_using=create_using)
    H = gen(n, create_using=create_using)

    G_edges = list(G.edges())
    G.add_edges_from(G_edges)
    H.add_edges_from(G_edges)
    assert edges_equal(G.edges(), H.edges(), directed=G.is_directed())

    H.remove_edge(0, 1)
    assert edges_equal(H.edges(), H.edges(), directed=H.is_directed())
    assert not edges_equal(G.edges(), H.edges(), directed=G.is_directed())


@pytest.mark.parametrize("n", [5, 10, 20])
@pytest.mark.parametrize("gen", [nx.complete_graph, nx.path_graph, nx.cycle_graph])
@pytest.mark.parametrize("weight", [1, 2, 3])
def test_edges_equal_weighted(n, gen, weight):
    """Test whether ``edges_equal`` properly compares edges with weight data."""
    G = gen(n)
    H = gen(n)

    G_edges = list(G.edges())
    G.add_weighted_edges_from((*e, weight) for e in G_edges)
    assert edges_equal(G.edges(), G.edges())

    H.add_weighted_edges_from((*e, weight + 1) for e in G_edges)
    assert edges_equal(H.edges(), H.edges())
    assert not edges_equal(G.edges(data=True), H.edges(data=True))


def test_edges_equal_data():
    """Test whether ``edges_equal`` properly compares edges with attribute dictionaries."""
    G = nx.path_graph(3)
    H = nx.path_graph(3)
    I = nx.path_graph(3, create_using=nx.MultiGraph)

    attrs = {(0, 1): {"attr1": 20, "attr2": "nothing"}, (1, 2): {"attr2": 3}}
    nx.set_edge_attributes(G, attrs)
    assert edges_equal(G.edges(data=True), G.edges(data=True))
    assert not edges_equal(G.edges(data=True), G.edges())

    nx.set_edge_attributes(H, attrs)
    assert edges_equal(G.edges(), H.edges())
    assert edges_equal(G.edges(data=True), H.edges(data=True))

    H[0][1]["attr2"] = "something"
    assert edges_equal(G.edges(), H.edges())
    assert not edges_equal(G.edges(data=True), H.edges(data=True))


def test_edges_equal_multigraph_data():
    """Test whether ``edges_equal`` properly compares edges with attribute dictionaries in ``MultiGraphs``."""
    G = nx.path_graph(3, create_using=nx.MultiGraph)
    I = nx.path_graph(3, create_using=nx.MultiGraph)

    G.add_edge(0, 1, 0, attr1="blue")
    G.add_edge(1, 2, 1, attr2="green")
    I.add_edge(0, 1, 0, attr1="blue")
    I.add_edge(0, 1, 1, attr2="green")
    assert edges_equal(G.edges(data=True), G.edges(data=True))
    assert not edges_equal(G.edges(), I.edges())
    assert not edges_equal(G.edges(data=True), I.edges(data=True))
    assert not edges_equal(G.edges(keys=True), I.edges(keys=True))
    assert not edges_equal(G.edges(keys=True, data=True), I.edges(keys=True, data=True))


def test_edges_equal_directed():
    """Test whether ``edges_equal`` properly compares directed edges."""
    G = nx.DiGraph([(0, 1)])
    I = nx.DiGraph([(1, 0)])

    assert edges_equal(G.edges(), I.edges(), directed=False)
    assert not edges_equal(G.edges(), I.edges(), directed=True)


def test_edges_equal_directed_data():
    """Test whether ``edges_equal`` properly compares directed edges with attribute dictionaries."""
    G = nx.DiGraph()
    I = nx.DiGraph()

    G.add_edge(0, 1, attr1="blue")
    I.add_edge(0, 1, attr1="blue")
    assert edges_equal(G.edges(data=True), G.edges(data=True), directed=True)
    I.add_edge(1, 2, attr2="green")
    assert not edges_equal(G.edges(data=True), I.edges(data=True), directed=True)
    G.add_edge(1, 2, attr2="green")
    assert edges_equal(G.edges(data=True), I.edges(data=True), directed=True)
    G.remove_edge(1, 2)
    G.add_edge(2, 1, attr2="green")
    assert edges_equal(G.edges(data=True), I.edges(data=True), directed=False)
    assert not edges_equal(G.edges(data=True), I.edges(data=True), directed=True)
