#!/usr/bin/env python
try:
    from StringIO import StringIO
except ImportError:
    from io import StringIO
from nose.tools import *
import networkx as nx
import networkx.readwrite.graph6 as g6

class TestGraph6Utils(object):

    def test_n_data_n_conversion(self):
        for i in [0, 1, 42, 62, 63, 64, 258047, 258048, 7744773, 68719476735]:
            assert_equal(g6.data_to_n(g6.n_to_data(i))[0], i)
            assert_equal(g6.data_to_n(g6.n_to_data(i))[1], [])
            assert_equal(g6.data_to_n(g6.n_to_data(i) + [42, 43])[1],
                         [42, 43])

    def test_data_sparse6_data_conversion(self):
        for data in [[], [0], [63], [63, 63], [0]*42,
                     [0, 1, 62, 42, 3, 11, 0, 11]]:
            assert_equal(g6.graph6_to_data(g6.data_to_graph6(data)), data)
            assert_equal(len(g6.data_to_graph6(data)), len(data))


class TestGraph6(object):

    def test_parse_graph6(self):
        data="""DF{"""
        G=nx.parse_graph6(data)
        assert_equal(sorted(G.nodes()),[0, 1, 2, 3, 4])
        assert_equal([e for e in sorted(G.edges())],
                     [(0, 3), (0, 4), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)])

    def test_read_graph6(self):
        data="""DF{"""
        G=nx.parse_graph6(data)
        fh = StringIO(data)
        Gin=nx.read_graph6(fh)
        assert_equal(sorted(G.nodes()),sorted(Gin.nodes()))
        assert_equal(sorted(G.edges()),sorted(Gin.edges()))

    def test_read_many_graph6(self):
        # Read many graphs into list
        data="""DF{\nD`{\nDqK\nD~{\n"""
        fh = StringIO(data)
        glist=nx.read_graph6(fh)
        assert_equal(len(glist),4)
        for G in glist:
            assert_equal(sorted(G.nodes()),[0, 1, 2, 3, 4])

    def test_generate_graph6(self):
        assert_equal(nx.generate_graph6(nx.empty_graph(0)), '>>graph6<<?')
        assert_equal(nx.generate_graph6(nx.empty_graph(1)), '>>graph6<<@')

        G1 = nx.complete_graph(4)
        assert_equal(nx.generate_graph6(G1, header=True), '>>graph6<<C~')
        assert_equal(nx.generate_graph6(G1, header=False), 'C~')

        G2 = nx.complete_bipartite_graph(6,9)
        assert_equal(nx.generate_graph6(G2, header=False),
                     'N??F~z{~Fw^_~?~?^_?') # verified by Sage

        G3 = nx.complete_graph(67)
        assert_equal(nx.generate_graph6(G3, header=False),
                     '~?@B' + '~' * 368 + 'w')

    def test_write_graph6(self):
        fh = StringIO()
        nx.write_graph6(nx.complete_bipartite_graph(6,9), fh)
        fh.seek(0)
        assert_equal(fh.read(), '>>graph6<<N??F~z{~Fw^_~?~?^_?\n')

    def test_generate_and_parse_graph6(self):
        for i in list(range(13)) + [31, 47, 62, 63, 64, 72]:
            g = nx.random_graphs.gnm_random_graph(i, i * i // 4, seed=i)
            gstr = nx.generate_graph6(g, header=False)
            assert_equal(len(gstr),
                         ((i-1) * i // 2 + 5) // 6 + (1 if i < 63 else 4))
            g2 = nx.parse_graph6(gstr)
            assert_equal(g2.order(), g.order())
            assert_equal(sorted(g2.edges()), sorted(g.edges()))

    @raises(nx.NetworkXError)
    def directed_error(self):
        nx.generate_graph6(nx.DiGraph())
