import numpy as np
import networkx as nx
import holoviews as hv
from holoviews import opts
from holoviews import dim

from holoviews.element.graphs import layout_nodes
from bokeh.sampledata.airport_routes import routes, airports

hv.extension('bokeh')

# Create dataset indexed by AirportID and with additional value dimension
airports = hv.Dataset(airports, ['AirportID'], ['Name', 'IATA', 'City'])
source_airports = list(airports.select(City='New York').data.AirportID)

# Add connections count to routes then aggregate and select just routes connecting to NYC
routes['connections'] = 1
nyc_graph = hv.Graph((routes, airports), ['SourceID', "DestinationID"], ['connections'], label='NYC Airport Connections')\
    .aggregate(function=np.count_nonzero).select(SourceID=source_airports)

# Lay out graph weighting and weight by the number of connections
np.random.seed(14)
graph = layout_nodes(nyc_graph, layout=nx.layout.fruchterman_reingold_layout, kwargs={'weight': 'connections'})
labels = hv.Labels(graph.nodes, ['x', 'y'], ['IATA', 'City'])

nyc_labels = labels.select(City='New York').opts(
    text_color='white', yoffset=0.05, text_font_size='16pt')

other_labels = labels[labels['City']!='New York'].opts(
    text_color='white', text_font_size='8pt')

cmap = {3697: 'red', 3797: 'blue'}

(graph * nyc_labels * other_labels).opts(
    opts.Graph(bgcolor='gray', width=800, height=800,
               edge_color=dim('SourceID').categorize(cmap, 'gray'),
               node_color=dim('index').categorize(cmap, 'gray'),
               title='NYC Airport Connections',
               xaxis=None, yaxis=None)
)