"""Unit tests for matplotlib drawing functions."""
import os
import itertools
from nose import SkipTest
import networkx as nx


class TestPylab(object):
    @classmethod
    def setupClass(cls):
        global plt
        try:
            import matplotlib as mpl
            mpl.use('PS', warn=False)
            import matplotlib.pyplot as plt
            plt.rcParams['text.usetex'] = False
        except ImportError:
            raise SkipTest('matplotlib not available.')
        except RuntimeError:
            raise SkipTest('matplotlib not available.')

    def setUp(self):
        self.G = nx.barbell_graph(4, 6)

    def test_draw(self):
        try:
            functions = [nx.draw_circular,
                         nx.draw_kamada_kawai,
                         nx.draw_random,
                         nx.draw_spectral,
                         nx.draw_spring,
                         nx.draw_shell]
            options = [{
                           'node_color': 'black',
                           'node_size': 100,
                           'width': 3,
                       }]
            for function, option in itertools.product(functions, options):
                function(self.G, **option)
                plt.savefig('test.ps')

        finally:
            try:
                os.unlink('test.ps')
            except OSError:
                pass

    def test_edge_colormap(self):
        colors = range(self.G.number_of_edges())
        nx.draw_spring(self.G, edge_color=colors, width=4,
                       edge_cmap=plt.cm.Blues, with_labels=True)
        plt.show()

    def test_arrows(self):
        nx.draw_spring(self.G.to_directed())
        plt.show()

    def test_edge_colors_and_widths(self):
        nx.draw_random(self.G, edgelist=[(0, 1), (0, 2)], width=[1, 2], edge_colors=['r', 'b'])

    def test_labels_and_colors(self):
        G = nx.cubical_graph()
        pos = nx.spring_layout(G)  # positions for all nodes
        # nodes
        nx.draw_networkx_nodes(G, pos,
                               nodelist=[0, 1, 2, 3],
                               node_color='r',
                               node_size=500,
                               alpha=0.8)
        nx.draw_networkx_nodes(G, pos,
                               nodelist=[4, 5, 6, 7],
                               node_color='b',
                               node_size=500,
                               alpha=0.8)
        # edges
        nx.draw_networkx_edges(G, pos, width=1.0, alpha=0.5)
        nx.draw_networkx_edges(G, pos,
                               edgelist=[(0, 1), (1, 2), (2, 3), (3, 0)],
                               width=8, alpha=0.5, edge_color='r')
        nx.draw_networkx_edges(G, pos,
                               edgelist=[(4, 5), (5, 6), (6, 7), (7, 4)],
                               width=8, alpha=0.5, edge_color='b')
        # some math labels
        labels = {}
        labels[0] = r'$a$'
        labels[1] = r'$b$'
        labels[2] = r'$c$'
        labels[3] = r'$d$'
        labels[4] = r'$\alpha$'
        labels[5] = r'$\beta$'
        labels[6] = r'$\gamma$'
        labels[7] = r'$\delta$'
        nx.draw_networkx_labels(G, pos, labels, font_size=16)
        plt.show()

    def test_axes(self):
        fig, ax = plt.subplots()
        nx.draw(self.G, ax=ax)

    def test_empty_graph(self):
        G = nx.Graph()
        nx.draw(G)

    def test_alpha_iter(self):
        pos = nx.random_layout(self.G)
        # with fewer alpha elements than nodes
        plt.subplot(131)
        nx.draw_networkx_nodes(self.G, pos, alpha=[0.1, 0.2])
        # with equal alpha elements and nodes
        num_nodes = len(self.G.nodes)
        alpha = [x / num_nodes for x in range(num_nodes)]
        colors = range(num_nodes)
        plt.subplot(132)
        nx.draw_networkx_nodes(self.G, pos, node_color=colors, alpha=alpha)
        # with more alpha elements than nodes
        alpha.append(1)
        plt.subplot(133)
        nx.draw_networkx_nodes(self.G, pos, alpha=alpha)