#!/usr/bin/env python

from nose.tools import *
from networkx import *
from networkx.algorithms.bipartite.generators import *

"""Generators - Bipartite
----------------------
"""


class TestGeneratorsBipartite():
    def test_complete_bipartite_graph(self):
        G = complete_bipartite_graph(0, 0)
        assert_true(is_isomorphic(G, null_graph()))

        for i in [1, 5]:
            G = complete_bipartite_graph(i, 0)
            assert_true(is_isomorphic(G, empty_graph(i)))
            G = complete_bipartite_graph(0, i)
            assert_true(is_isomorphic(G, empty_graph(i)))

        G = complete_bipartite_graph(2, 2)
        assert_true(is_isomorphic(G, cycle_graph(4)))

        G = complete_bipartite_graph(1, 5)
        assert_true(is_isomorphic(G, star_graph(5)))

        G = complete_bipartite_graph(5, 1)
        assert_true(is_isomorphic(G, star_graph(5)))

        # complete_bipartite_graph(m1,m2) is a connected graph with
        # m1+m2 nodes and  m1*m2 edges
        for m1, m2 in [(5, 11), (7, 3)]:
            G = complete_bipartite_graph(m1, m2)
            assert_equal(number_of_nodes(G), m1 + m2)
            assert_equal(number_of_edges(G), m1 * m2)

        assert_raises(networkx.exception.NetworkXError,
                      complete_bipartite_graph, 7, 3, create_using=DiGraph())

        mG = complete_bipartite_graph(7, 3, create_using=MultiGraph())
        assert_equal(sorted(mG.edges()), sorted(G.edges()))

        # specify nodes rather than number of nodes
        G = complete_bipartite_graph([1, 2], ['a', 'b'])
        has_edges = G.has_edge(1, 'a') & G.has_edge(1, 'b') &\
            G.has_edge(2, 'a') & G.has_edge(2, 'b')
        assert_true(has_edges)
        assert_equal(G.size(), 4)

    def test_configuration_model(self):
        aseq = [3, 3, 3, 3]
        bseq = [2, 2, 2, 2, 2]
        assert_raises(networkx.exception.NetworkXError,
                      configuration_model, aseq, bseq)

        aseq = [3, 3, 3, 3]
        bseq = [2, 2, 2, 2, 2, 2]
        G = configuration_model(aseq, bseq)
        assert_equal(sorted(d for n, d in G.degree()),
                     [2, 2, 2, 2, 2, 2, 3, 3, 3, 3])

        aseq = [2, 2, 2, 2, 2, 2]
        bseq = [3, 3, 3, 3]
        G = configuration_model(aseq, bseq)
        assert_equal(sorted(d for n, d in G.degree()),
                     [2, 2, 2, 2, 2, 2, 3, 3, 3, 3])

        aseq = [2, 2, 2, 1, 1, 1]
        bseq = [3, 3, 3]
        G = configuration_model(aseq, bseq)
        assert_equal(sorted(d for n, d in G.degree()),
                     [1, 1, 1, 2, 2, 2, 3, 3, 3])

        GU = project(Graph(G), range(len(aseq)))
        assert_equal(GU.number_of_nodes(), 6)

        GD = project(Graph(G), range(len(aseq), len(aseq) + len(bseq)))
        assert_equal(GD.number_of_nodes(), 3)

        assert_raises(networkx.exception.NetworkXError,
                      configuration_model, aseq, bseq,
                      create_using=DiGraph())

    def test_havel_hakimi_graph(self):
        aseq = [3, 3, 3, 3]
        bseq = [2, 2, 2, 2, 2]
        assert_raises(networkx.exception.NetworkXError,
                      havel_hakimi_graph, aseq, bseq)

        bseq = [2, 2, 2, 2, 2, 2]
        G = havel_hakimi_graph(aseq, bseq)
        assert_equal(sorted(d for n, d in G.degree()),
                     [2, 2, 2, 2, 2, 2, 3, 3, 3, 3])

        aseq = [2, 2, 2, 2, 2, 2]
        bseq = [3, 3, 3, 3]
        G = havel_hakimi_graph(aseq, bseq)
        assert_equal(sorted(d for n, d in G.degree()),
                     [2, 2, 2, 2, 2, 2, 3, 3, 3, 3])

        GU = project(Graph(G), range(len(aseq)))
        assert_equal(GU.number_of_nodes(), 6)

        GD = project(Graph(G), range(len(aseq), len(aseq) + len(bseq)))
        assert_equal(GD.number_of_nodes(), 4)
        assert_raises(networkx.exception.NetworkXError,
                      havel_hakimi_graph, aseq, bseq,
                      create_using=DiGraph())

    def test_reverse_havel_hakimi_graph(self):
        aseq = [3, 3, 3, 3]
        bseq = [2, 2, 2, 2, 2]
        assert_raises(networkx.exception.NetworkXError,
                      reverse_havel_hakimi_graph, aseq, bseq)

        bseq = [2, 2, 2, 2, 2, 2]
        G = reverse_havel_hakimi_graph(aseq, bseq)
        assert_equal(sorted(d for n, d in G.degree()),
                     [2, 2, 2, 2, 2, 2, 3, 3, 3, 3])

        aseq = [2, 2, 2, 2, 2, 2]
        bseq = [3, 3, 3, 3]
        G = reverse_havel_hakimi_graph(aseq, bseq)
        assert_equal(sorted(d for n, d in G.degree()),
                     [2, 2, 2, 2, 2, 2, 3, 3, 3, 3])

        aseq = [2, 2, 2, 1, 1, 1]
        bseq = [3, 3, 3]
        G = reverse_havel_hakimi_graph(aseq, bseq)
        assert_equal(sorted(d for n, d in G.degree()),
                     [1, 1, 1, 2, 2, 2, 3, 3, 3])

        GU = project(Graph(G), range(len(aseq)))
        assert_equal(GU.number_of_nodes(), 6)

        GD = project(Graph(G), range(len(aseq), len(aseq) + len(bseq)))
        assert_equal(GD.number_of_nodes(), 3)
        assert_raises(networkx.exception.NetworkXError,
                      reverse_havel_hakimi_graph, aseq, bseq,
                      create_using=DiGraph())

    def test_alternating_havel_hakimi_graph(self):
        aseq = [3, 3, 3, 3]
        bseq = [2, 2, 2, 2, 2]
        assert_raises(networkx.exception.NetworkXError,
                      alternating_havel_hakimi_graph, aseq, bseq)

        bseq = [2, 2, 2, 2, 2, 2]
        G = alternating_havel_hakimi_graph(aseq, bseq)
        assert_equal(sorted(d for n, d in G.degree()),
                     [2, 2, 2, 2, 2, 2, 3, 3, 3, 3])

        aseq = [2, 2, 2, 2, 2, 2]
        bseq = [3, 3, 3, 3]
        G = alternating_havel_hakimi_graph(aseq, bseq)
        assert_equal(sorted(d for n, d in G.degree()),
                     [2, 2, 2, 2, 2, 2, 3, 3, 3, 3])

        aseq = [2, 2, 2, 1, 1, 1]
        bseq = [3, 3, 3]
        G = alternating_havel_hakimi_graph(aseq, bseq)
        assert_equal(sorted(d for n, d in G.degree()),
                     [1, 1, 1, 2, 2, 2, 3, 3, 3])

        GU = project(Graph(G), range(len(aseq)))
        assert_equal(GU.number_of_nodes(), 6)

        GD = project(Graph(G), range(len(aseq), len(aseq) + len(bseq)))
        assert_equal(GD.number_of_nodes(), 3)

        assert_raises(networkx.exception.NetworkXError,
                      alternating_havel_hakimi_graph, aseq, bseq,
                      create_using=DiGraph())

    def test_preferential_attachment(self):
        aseq = [3, 2, 1, 1]
        G = preferential_attachment_graph(aseq, 0.5)
        assert_raises(networkx.exception.NetworkXError,
                      preferential_attachment_graph, aseq, 0.5,
                      create_using=DiGraph())

    def test_random_graph(self):
        n = 10
        m = 20
        G = random_graph(n, m, 0.9)
        assert_equal(len(G), 30)
        assert_true(is_bipartite(G))
        X, Y = nx.algorithms.bipartite.sets(G)
        assert_equal(set(range(n)), X)
        assert_equal(set(range(n, n + m)), Y)

    def test_random_graph(self):
        n = 10
        m = 20
        G = random_graph(n, m, 0.9, directed=True)
        assert_equal(len(G), 30)
        assert_true(is_bipartite(G))
        X, Y = nx.algorithms.bipartite.sets(G)
        assert_equal(set(range(n)), X)
        assert_equal(set(range(n, n + m)), Y)

    def test_gnmk_random_graph(self):
        n = 10
        m = 20
        edges = 200
        G = gnmk_random_graph(n, m, edges)
        assert_equal(len(G), 30)
        assert_true(is_bipartite(G))
        X, Y = nx.algorithms.bipartite.sets(G)
        print(X)
        assert_equal(set(range(n)), X)
        assert_equal(set(range(n, n + m)), Y)
        assert_equal(edges, len(list(G.edges())))