#!/usr/bin/env python
# coding: utf-8

# In[1]:


get_ipython().run_line_magic('pylab', 'inline')
import networkx as nx
import pandas as pd
from skbio import DNA


# In[2]:


# create a graph with skbio.DNASequence objects as nodes
G=nx.Graph()
d1 = DNA('AAAAAAAAAA', 'd1')
d2 = DNA('CAAAAAAAAA', 'd2')
d3 = DNA('TTTTTTTTTT', 'd3')
d4 = DNA('TTTTTTTTAA', 'd4')
d5 = DNA('TAAAAAAATT', 'd5')
G.add_edge(d1, d2)
G.add_edge(d3, d4)
G.add_edge(d3, d5)
G.add_edge(d4, d5)
pos=nx.spring_layout(G)


# In[3]:


def get_edge_widths(G, num_bins=3, min_width=1):
    result = pd.qcut([u.distance(v) for (u,v) in G.edges()], 
                     num_bins,
                     labels=range(min_width,min_width+num_bins))
    return result


# In[4]:


nx.draw_networkx_nodes(G, pos, node_size=700)
## Question: is passing width a list of floats supported functionality?
## Answer: sounds like it - https://github.com/networkx/networkx/issues/1441
nx.draw_networkx_edges(G, pos, width=get_edge_widths(G))
nx.draw_networkx_labels(G, pos, labels={d:d.id for d in G.nodes()},
                        font_size=20, font_family='sans-serif')
_ = plt.axis('off')


# In[ ]: