import networkx as nx
import pylab as plt
import pygraphviz as pgv
from networkx.drawing.nx_agraph import graphviz_layout, to_agraph

def de_bruijn_graph(st, k):
    edges = []
    nodes = set()
    for i in range(len(st) - k + 1):
        edges.append((st[i : i + k - 1], st[i + 1 : i + k]))
        nodes.add(st[i : i + k - 1])
        nodes.add(st[i + 1 : i + k])
    return nodes, edges

def draw(nodes, edges, fname):
    G = nx.DiGraph()
    G.add_nodes_from(nodes)
    G.add_edges_from(edges)
    G.graph['graph'] = {'rankdir': 'TD'}
    G.graph['node'] = {'shape': 'circle'}
    G.graph['edges'] = {'arrowsize': '4.0'}

    A = to_agraph(G)
    A.layout('dot')
    A.draw(fname)

def main():
    nodes1, edges1 = de_bruijn_graph('0001110100', 3)
    nodes2, edges2 = de_bruijn_graph('0000111101100101000', 4)
    draw(nodes1, edges1, 'deb_graph1.png')
    draw(nodes2, edges2, 'deb_graph2.png')

if __name__ == '__main__':
    main()