# -*- coding: utf-8 -*-
from bokeh.io import curdoc
from bokeh.layouts import row, widgetbox
from bokeh.models import CustomJS, ColumnDataSource, Select, HoverTool, ColorBar, LinearColorMapper, BoxZoomTool, PanTool, ResetTool, SaveTool
from bokeh.models.widgets import Slider, TextInput
from bokeh.palettes import Viridis256
from bokeh.plotting import figure, show

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import world_countries as wc


def set_colors(dataframe, data_set):
	plot_data_name = data_sets.get(data_set)
	norm_data = dataframe[plot_data_name]/dataframe[plot_data_name].max() #normalize data to turn into viridis rgba

	colors = [
	    "#%02x%02x%02x" % (int(r), int(g), int(b)) for r, g, b, _ in 255*mpl.cm.viridis(mpl.colors.Normalize()(norm_data))
	]
	return colors


# set up callbacks
def update_plot(attrname, old, new):
	# set variables for new data
	data_set = data_select.value
	plot_data_name = data_sets.get(data_set)
	data = dataframe[plot_data_name]
	
	# create new data
	colors = set_colors(dataframe, data_set)
	dataframe['colors'] = colors
	new_source = ColumnDataSource(dataframe)

	# update plot
	plot.title.text = data_set
	source.data.update(new_source.data)

	#update color bar
	data_min = data.min()
	data_max = data.max()
	color_mapper.low = data_min
	color_mapper.high = data_max

def update_location(attrname, old, new):
	lat_value = lat_slider.value
	lon_value = lon_slider.value
	new_location_source = ColumnDataSource( pd.DataFrame([[lon_value,lat_value]], columns=['x', 'y']) )
	location_source.data.update(new_location_source.data)


data_set = 'Prediction'
data_sets = {
	'Prediction': 'prediction',
	'Uncertainty':'uncertainty',
}

data_path = 'path_to_data'
dataframe = pd.read_csv(data_path)

colors = set_colors(dataframe, data_set)
dataframe['colors'] = colors
source = ColumnDataSource(dataframe)

#initialize the specific location to center of map
lat_value = 0.
lon_value = 0.
location_source = ColumnDataSource( pd.DataFrame([[lon_value,lat_value]], columns=['x', 'y']) )

plot_data_name = data_sets.get(data_set)
data = dataframe[plot_data_name]
data_min = data.min()
data_max = data.max()

# get the world map outlines data
world_countries = wc.data.copy()
worldmap = pd.DataFrame.from_dict(world_countries, orient='index')

#setting source for use in plot functionality
hover_tool = HoverTool(tooltips = [
    ('lon', '@lons'),
    ('lat', '@lats'),
    ('prediction', '@prediction'),
    ('uncertainty ', '@uncertainty'),
])

tools = [hover_tool, BoxZoomTool(), ResetTool(), SaveTool()]

plot = figure( x_range=(-180, 180), y_range=(-90, 90), plot_width=800, plot_height=600 ,\
    tools=tools , toolbar_location='below', toolbar_sticky=False,
    )

# define the countries' outlines
plot.patches(xs=worldmap['lons'], ys=worldmap['lats'], fill_color="white", fill_alpha=0.,
    line_color="black", line_width=1.2)

#need scatter plot for hover to work
scatter_plot = plot.scatter(x='lons', y='lats',
    fill_color='colors',  source=source, fill_alpha=0.5,
    line_color=None,
)   

plot.tools[0].renderers.append(scatter_plot)

# set up the color bar
color_mapper = LinearColorMapper(palette=Viridis256, low=data_min, high=data_max)
color_bar = ColorBar(color_mapper=color_mapper, \
                 label_standoff=12, border_line_color=None, location=(0,0))
plot.add_layout(color_bar, 'right')

plot.circle(x='x', y='y', size=15, color="white", alpha=0.75, source=location_source)

plot.title.text = data_set
plot.xaxis.axis_label = 'Lon'
plot.yaxis.axis_label = 'Lat'


#initialize the data selection widgets
data_select = Select(value=data_set, title='Data Set', options=sorted(data_sets.keys()))
lat_slider = Slider(title="Latitude", value=0.0, start=-90.0, end=90.0, step=1.)
lon_slider = Slider(title="Longitude", value=0.0, start=-180.0, end=180.0, step=1.)
controls = widgetbox(data_select,lat_slider, lon_slider)


data_select.on_change('value', update_plot)
lat_slider.on_change('value', update_location)
lon_slider.on_change('value', update_location)
layout = row(controls, plot)

# add to document
curdoc().add_root( layout )