import matplotlib as mpl
import networkx as nx
import pandas as pd
import random
import numpy as np

from visJS2jupyter import visJS_module
reload(visJS_module)

"""
Explanation: Multigraph Network Styling for visJS2jupyter

Authors: Brin Rosenthal (sbrosenthal@ucsd.edu), Mikayla Webster (m1webste@ucsd.edu)

Import packages
End of explanation
"""


G = nx.connected_watts_strogatz_graph(30,5,.2)
G = nx.MultiGraph(G)
edges = G.edges(keys = True) # for multigraphs every edge has to be represented by a three-tuple (source, target, key)

"""
Explanation: We start by creating a randomized, single-edged graph, and convert that to a multigraph
End of explanation
"""


sources = list(zip(*edges)[0])
targets = list(zip(*edges)[1])
backward_edges = list(zip(targets, sources)) # demonstarting adding backwards edges

G.add_edges_from(backward_edges)
edges = list(G.edges(data = True))

nodes = list(G.nodes()) # type cast to list in order to make compatible with networkx 1.11 and 2.0
edges = list(G.edges(keys = True)) # for nx 2.0, returns an "EdgeView" object rather than an iterable

"""
Explanation: We duplicate every edge in the graph to make it a true multigraph. 
Note: NetworkX does not support duplicate edges with opposite directions. NetworkX will flip any backwards edges you try to add to your graph. For example, if your graph currently contains the edges [(0,1), (1,2)] and you add the edge (1,0) to your graph, your graph will now contain edges [(0,1), (0,1), (1,2)]
End of explanation
"""


# add some node attributes to color-code by
degree = dict(G.degree()) # nx 2.0 returns a "DegreeView" object. Cast to dict to maintain compatibility with nx 1.11
bc = nx.betweenness_centrality(G)
nx.set_node_attributes(G, name = 'degree', values = degree)             # between networkx 1.11 and 2.0, therefore we must
nx.set_node_attributes(G, name = 'betweenness_centrality', values = bc) # explicitly pass our arguments 
                                                                        # (not implicitly through position)

# add the edge attribute 'weight' to color-code by
weights = []
for i in range(len(edges)):
    weights.append(float(random.randint(1,5))) 
    
w_dict = dict(zip(edges, weights))
nx.set_edge_attributes(G, name = 'weight', values = w_dict)


# map the betweenness centrality to the node color, using matplotlib spring_r colormap
node_to_color = visJS_module.return_node_to_color(G,field_to_map='betweenness_centrality',cmap=mpl.cm.spring_r,alpha = 1,
                                                 color_max_frac = .9,color_min_frac = .1)

# map weight to edge color, using default settings
edge_to_color = visJS_module.return_edge_to_color(G,field_to_map='weight')

"""
Explanation: Multigraph Node and Edge Styling
There is no difference between multigraph and single-edged-graph styling. Just map the node and edge attributes to some visual properties, and style the nodes and edges according to these properties (like usual!)
End of explanation
"""


# set node initial positions using networkx's spring_layout function
pos = nx.spring_layout(G)

nodes_dict = [{"id":n,"color":node_to_color[n],
               "degree":nx.degree(G,n),
              "x":pos[n][0]*1000,
              "y":pos[n][1]*1000} for n in nodes
              ]
node_map = dict(zip(nodes,range(len(nodes))))  # map to indices for source/target in edges
edges_dict = [{"source":node_map[edges[i][0]], "target":node_map[edges[i][1]], 
              "color":edge_to_color[(edges[i][0],edges[i][1],edges[i][2])],"title":'test'} # remember (source, target, key)
              for i in range(len(edges))]

# set some network-wide styles
visJS_module.visjs_network(nodes_dict,edges_dict,
                          node_size_multiplier=3,
                          node_size_transform = '',
                          node_color_highlight_border='red',
                          node_color_highlight_background='#D3918B',
                          node_color_hover_border='blue',
                          node_color_hover_background='#8BADD3',
                          node_font_size=25,
                          edge_arrow_to=True,
                          physics_enabled=True,
                          edge_color_highlight='#8A324E',
                          edge_color_hover='#8BADD3',
                          edge_width=3,
                          max_velocity=15,
                          min_velocity=1,
                          edge_smooth_enabled = True)

"""
Explanation: Interactive network
Note that this example is simply the multigraph version of our "Complex Parameters" notebook.
End of explanation
"""


