import networkx as nx
import network2tikz as tikz


def draw_graph(G: nx.Graph, filename=None):
    """Draw graph marking stations as double circles."""
    layout = nx.planar_layout(G)

    tikz.plot(
        G,
        filename=filename,
        layout=layout,
        standalone=False,
        **_get_edge_styles(G),
        **_get_node_styles(G)
    )


def _get_edge_styles(G):
    styles = {
        'edges_label': {(u, v): '({}, {})'.format(
            G[u][v]['weight'],
            G[u][v]['c']) for u, v in G.edges}
    }

    return styles


def _get_node_styles(G):
    styles = {
        'vertex_color': ['white' for n in G.nodes],
        'vertex_size': 1.
    }

    nodes_label = {}
    nodes_style = {}
    for n in G:
        style = 'circle'
        label = G.nodes[n].get('label', str(n))

        c = G.nodes[n].get('c', None)
        if c is not None:
            style += ',style={draw,double}'
            label += ',{}'.format(c)

        nodes_label[n] = label
        nodes_style[n] = style

    styles['nodes_label'] = nodes_label
    styles['nodes_style'] = nodes_style

    return styles


if __name__ == '__main__':
    # Write Latex files of all test case graph configs
    import os
    import sys
    from tests import config

    try:
        plot_dir = sys.argv[1]
    except IndexError:
        plot_dir = 'plots'

    if config.config_list:
        try:
            os.mkdir(plot_dir)
        except FileExistsError:
            pass

    for conf_name in config.config_list:
        conf = getattr(config, conf_name)
        G = config.get_graph(conf)
        draw_graph(G, os.path.join(plot_dir, conf_name + '.tex'))