from io import BytesIO
import tempfile
from unittest import TestCase

from nose.tools import assert_equal
from nose.tools import assert_true

import networkx as nx
from networkx.testing.utils import assert_edges_equal
from networkx.testing.utils import assert_nodes_equal


class TestSparseGraph6(object):

    def test_from_sparse6_bytes(self):
        data = b':Q___eDcdFcDeFcE`GaJ`IaHbKNbLM'
        G = nx.from_sparse6_bytes(data)
        assert_nodes_equal(sorted(G.nodes()),
                           [0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
                            10, 11, 12, 13, 14, 15, 16, 17])
        assert_edges_equal(G.edges(),
                           [(0, 1), (0, 2), (0, 3), (1, 12), (1, 14), (2, 13),
                            (2, 15), (3, 16), (3, 17), (4, 7), (4, 9), (4, 11),
                            (5, 6), (5, 8), (5, 9), (6, 10), (6, 11), (7, 8),
                            (7, 10), (8, 12), (9, 15), (10, 14), (11, 13),
                            (12, 16), (13, 17), (14, 17), (15, 16)])

    def test_from_bytes_multigraph_graph(self):
        graph_data = b':An'
        G = nx.from_sparse6_bytes(graph_data)
        assert_true(type(G), nx.Graph)
        multigraph_data = b':Ab'
        M = nx.from_sparse6_bytes(multigraph_data)
        assert_true(type(M), nx.MultiGraph)

    def test_read_sparse6(self):
        data = b':Q___eDcdFcDeFcE`GaJ`IaHbKNbLM'
        G = nx.from_sparse6_bytes(data)
        fh = BytesIO(data)
        Gin = nx.read_sparse6(fh)
        assert_nodes_equal(G.nodes(), Gin.nodes())
        assert_edges_equal(G.edges(), Gin.edges())

    def test_read_many_graph6(self):
        # Read many graphs into list
        data = (b':Q___eDcdFcDeFcE`GaJ`IaHbKNbLM\n'
                b':Q___dCfDEdcEgcbEGbFIaJ`JaHN`IM')
        fh = BytesIO(data)
        glist = nx.read_sparse6(fh)
        assert_equal(len(glist), 2)
        for G in glist:
            assert_nodes_equal(G.nodes(),
                               [0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
                                10, 11, 12, 13, 14, 15, 16, 17])


class TestWriteSparse6(TestCase):
    """Unit tests for writing graphs in the sparse6 format.

    Most of the test cases were checked against the sparse6 encoder in Sage.

    """

    def test_null_graph(self):
        G = nx.null_graph()
        result = BytesIO()
        nx.write_sparse6(G, result)
        self.assertEqual(result.getvalue(), b'>>sparse6<<:?\n')

    def test_trivial_graph(self):
        G = nx.trivial_graph()
        result = BytesIO()
        nx.write_sparse6(G, result)
        self.assertEqual(result.getvalue(), b'>>sparse6<<:@\n')

    def test_empty_graph(self):
        G = nx.empty_graph(5)
        result = BytesIO()
        nx.write_sparse6(G, result)
        self.assertEqual(result.getvalue(), b'>>sparse6<<:D\n')

    def test_large_empty_graph(self):
        G = nx.empty_graph(68)
        result = BytesIO()
        nx.write_sparse6(G, result)
        self.assertEqual(result.getvalue(), b'>>sparse6<<:~?@C\n')

    def test_very_large_empty_graph(self):
        G = nx.empty_graph(258049)
        result = BytesIO()
        nx.write_sparse6(G, result)
        self.assertEqual(result.getvalue(), b'>>sparse6<<:~~???~?@\n')

    def test_complete_graph(self):
        G = nx.complete_graph(4)
        result = BytesIO()
        nx.write_sparse6(G, result)
        self.assertEqual(result.getvalue(), b'>>sparse6<<:CcKI\n')

    def test_no_header(self):
        G = nx.complete_graph(4)
        result = BytesIO()
        nx.write_sparse6(G, result, header=False)
        self.assertEqual(result.getvalue(), b':CcKI\n')

    def test_padding(self):
        codes = (b':Cdv', b':DaYn', b':EaYnN', b':FaYnL', b':GaYnLz')
        for n, code in enumerate(codes, start=4):
            G = nx.path_graph(n)
            result = BytesIO()
            nx.write_sparse6(G, result, header=False)
            self.assertEqual(result.getvalue(), code + b'\n')

    def test_complete_bipartite(self):
        G = nx.complete_bipartite_graph(6, 9)
        result = BytesIO()
        nx.write_sparse6(G, result)
        # Compared with sage
        expected = b'>>sparse6<<:Nk' + b'?G`cJ' * 9 + b'\n'
        assert_equal(result.getvalue(), expected)

    def test_read_write_inverse(self):
        for i in list(range(13)) + [31, 47, 62, 63, 64, 72]:
            m = min(2 * i, i * i // 2)
            g = nx.random_graphs.gnm_random_graph(i, m, seed=i)
            gstr = BytesIO()
            nx.write_sparse6(g, gstr, header=False)
            # Strip the trailing newline.
            gstr = gstr.getvalue().rstrip()
            g2 = nx.from_sparse6_bytes(gstr)
            assert_equal(g2.order(), g.order())
            assert_edges_equal(g2.edges(), g.edges())

    def no_directed_graphs(self):
        with self.assertRaises(nx.NetworkXNotImplemented):
            nx.write_sparse6(nx.DiGraph(), BytesIO())

    def test_write_path(self):
        # On Windows, we can't reopen a file that is open
        # So, for test we get a valid name from tempfile but close it.
        with tempfile.NamedTemporaryFile() as f:
            fullfilename = f.name
        # file should be closed now, so write_sparse6 can open it
        nx.write_sparse6(nx.null_graph(), fullfilename)
        fh = open(fullfilename, mode='rb')
        self.assertEqual(fh.read(), b'>>sparse6<<:?\n')
        fh.close()
        import os
        os.remove(fullfilename)
