#!/usr/bin/env python
import cProfile
import pstats
import string
import random
import networkx as nx
from bokeh.plotting import figure, output_file, show
from bokeh.models import (
        ColumnDataSource,
        WheelZoomTool,
        PanTool,
        ResetTool)

def generateName(size=10, chars=string.ascii_uppercase + string.ascii_lowercase):
    return ''.join(random.choice(chars) for _ in range(size))

def generateCoord(rmin=0, rmax=500):
    return random.randint(rmin, rmax)

def createGraph(size=500):
    graph = nx.Graph()

    # nodes
    for i in range(size):
        name = generateName()
        if not graph.has_node(name):
            x = generateCoord()
            y = generateCoord()
            graph.add_node(name, x=x, y=y)

    # edges
    nodes = graph.nodes()
    for node in nodes:
        while len(graph.neighbors(node)) < 5:
            other = random.choice(nodes)
            graph.add_edge(node, other)

    return graph

def createPlot(fn, title, graph, edges):
    origFA = 0.4
    origLA = 0.9

    output_file(fn, mode='inline')
    tools = [WheelZoomTool(), PanTool(), ResetTool()]

    names = graph.nodes()
    xs = [graph.node[name]['x'] for name in names]
    ys = [graph.node[name]['y'] for name in names]
    fa = [origFA] * len(names)
    la = [origLA] * len(names)
    source = ColumnDataSource(data=dict(name=names, x=xs, y=ys, fa=fa, la=la))

    plot = figure(title=title,
            x_range = (min(xs)-10, max(xs)+10),
            y_range = (min(ys)-10, max(ys)+10),
            height = 800,
            width = 800,
            tools = tools)

    scatter = plot.scatter('x', 'y', fill_alpha='fa', line_alpha='la', source=source, size=7, color='blue')

    if edges:
        for edge in graph.edges():
            start = edge[0]
            end = edge[1]
            startInd = names.index(start)
            endInd = names.index(end)
            startx = xs[startInd]
            starty = ys[startInd]
            endx = xs[endInd]
            endy = ys[endInd]
            plot.line(x=[startx, endx], y=[starty, endy], color='red')

    show(plot)

g = createGraph()

# with edges
pr = cProfile.Profile()
pr.enable()
createPlot('edges.html', 'Minimal Example', g, True)
pr.disable()
pr.dump_stats('edges.txt')
stats = pstats.Stats('edges.txt')
stats.strip_dirs()
stats.sort_stats('tottime')
stats.print_stats(.05)

# without edges
pr2 = cProfile.Profile()
pr2.enable()
createPlot('noedges.html', 'Minimal Example', g, False)
pr2.disable()
pr2.dump_stats('noedges.txt')
stats = pstats.Stats('noedges.txt')
stats.strip_dirs()
stats.sort_stats('tottime')
stats.print_stats(.05)