import numpy as np
from random import randint
import math
import matplotlib as mpl
import networkx as nx

from visJS2jupyter import visJS_module
reload(visJS_module)

# create a simple graph
G = nx.connected_watts_strogatz_graph(30,5,.2)
nodes = list(G.nodes()) # type cast to list in order to make compatible with networkx 1.11 and 2.0
edges = list(G.edges()) # for nx 2.0, returns an "EdgeView" object rather than an iterable

"""
Explanation: More complex network styling for visJS2jupyter

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

Import packages
End of explanation
"""


# add a node attributes to color-code by
cc = nx.clustering(G)
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 = 'clustering_coefficient', values = cc) # parameter order for name and values is switched 
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 edge attribute weight
weights=[randint(0,9) for p in range(len(edges))]
edge_to_weight = dict(zip(edges, weights))
nx.set_edge_attributes(G, name = 'weight', values = edge_to_weight)

# setting shape of each node using a dictionary
stars = ['star']*20 # some will be stars
circles = ['dot']*10 # some will be circles
shapes = stars + circles
node_to_shape = dict(zip(nodes, shapes))

"""
Explanation: Map node attributes to visual properties, and style the nodes and edges

To map node attributes to properties, simply add the property to the graph as a node-attribute, and use the return_node_to_color function
End of explanation
"""


# 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, vmin = 0, vmax = 0.08)

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

# per node attributes
nodes_dict = [{"id":n,
               "color":node_to_color[n],
               "degree":nx.degree(G,n),
               "node_shape": node_to_shape[n],
               "x":pos[n][0]*700,
               "y":pos[n][1]*700} for n in nodes
              ]

# map to indices for source/target in edges
node_map = dict(zip(nodes,range(len(nodes)))) 

# map colors to edges based on degree
edge_to_color = visJS_module.return_edge_to_color(G,field_to_map='weight',cmap=mpl.cm.jet,alpha = 1.0, color_vals_transform = None,ceil_val=10,
                        vmin=0,vmax=5)

# per edge attributes
edges_dict = [{"source":node_map[edges[i][0]], "target":node_map[edges[i][1]], 
              "color":edge_to_color[edges[i]]} for i in range(len(edges))]

# set some network-wide styles
visJS_module.visjs_network(nodes_dict,edges_dict,
                          node_size_multiplier=7,
                          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)



"""
Explanation: <a id='interactive_network'></a>
Interactive network
TIP: If you get errors running the following cell, try updating visJS2jupyter (pip install visJS2jupyter --upgrade). This cell is dependent on an update made on March 9th, 2018
End of explanation
"""


