from itertools import chain

import networkx as nx
from networkx.utils import edges_equal, nodes_equal


def _check_custom_label_attribute(input_trees, res_tree, label_attribute):
    res_attr_dict = nx.get_node_attributes(res_tree, label_attribute)
    res_attr_set = set(res_attr_dict.values())
    input_label = (tree for tree, root in input_trees)
    input_label_set = set(chain.from_iterable(input_label))
    return res_attr_set == input_label_set


def test_empty_sequence():
    """Joining the empty sequence results in the tree with one node."""
    T = nx.join_trees([])
    assert len(T) == 1
    assert T.number_of_edges() == 0


def test_single():
    """Joining just one tree yields a tree with one more node."""
    T = nx.empty_graph(1)
    trees = [(T, 0)]
    actual_with_label = nx.join_trees(trees, label_attribute="custom_label")
    expected = nx.path_graph(2)
    assert nodes_equal(list(expected), list(actual_with_label))
    assert edges_equal(list(expected.edges()), list(actual_with_label.edges()))


def test_basic():
    """Joining multiple subtrees at a root node."""
    trees = [(nx.full_rary_tree(2, 2**2 - 1), 0) for i in range(2)]
    expected = nx.full_rary_tree(2, 2**3 - 1)
    actual = nx.join_trees(trees, label_attribute="old_labels")
    assert nx.is_isomorphic(actual, expected)
    assert _check_custom_label_attribute(trees, actual, "old_labels")

    actual_without_label = nx.join_trees(trees)
    assert nx.is_isomorphic(actual_without_label, expected)
    # check that no labels were stored
    assert all(not data for _, data in actual_without_label.nodes(data=True))


def test_first_label():
    """Test the functionality of the first_label argument."""
    T1 = nx.path_graph(3)
    T2 = nx.path_graph(2)
    actual = nx.join_trees([(T1, 0), (T2, 0)], first_label=10)
    expected_nodes = set(range(10, 16))
    assert set(actual.nodes()) == expected_nodes
    assert set(actual.neighbors(10)) == {11, 14}