import piplite
await piplite.install(['bqplot', 'ipyleaflet'])

import os
import io
import json
from urllib.request import urlopen
from datetime import datetime

import numpy as np
import pandas as pd

from ipywidgets import Dropdown

from bqplot import Lines, Figure, LinearScale, DateScale, Axis

from ipyleaflet import Map, GeoJSON, WidgetControl

with open('data/nations.json', 'r') as fobj:
    data = pd.read_json(fobj.read())

def clean_data(data):
    for column in ['income', 'lifeExpectancy', 'population']:
        data = data.drop(data[data[column].apply(len) <= 4].index)
    return data

def extrap_interp(data):
    data = np.array(data)
    x_range = np.arange(1800, 2009, 1.)
    y_range = np.interp(x_range, data[:, 0], data[:, 1])
    return y_range

def extrap_data(data):
    for column in ['income', 'lifeExpectancy', 'population']:
        data[column] = data[column].apply(extrap_interp)
    return data

data = clean_data(data)
data = extrap_data(data)

data

date_start = datetime(1800, 12, 31)
date_end = datetime(2009, 12, 31)

date_scale = DateScale(min=date_start, max=date_end)

date_data = pd.date_range(start=date_start, end=date_end, freq='A', normalize=True)

country_name = 'Angola'
data_name = 'income'

x_data = data[data.name == country_name][data_name].values[0]

x_scale = LinearScale()

lines = Lines(x=date_data, y=x_data, scales={'x': date_scale, 'y': x_scale})

ax_x = Axis(label='Year', scale=date_scale, num_ticks=10, tick_format='%Y')
ax_y = Axis(label=data_name.capitalize(), scale=x_scale, orientation='vertical', side='left')

figure = Figure(axes=[ax_x, ax_y], title=country_name, marks=[lines], animation_duration=500,
                layout={'max_height': '250px', 'max_width': '400px'})

def update_figure(country_name, data_name):
    try:
        lines.y = data[data.name == country_name][data_name].values[0]
        ax_y.label = data_name.capitalize()
        figure.title = country_name
    except IndexError:
        pass

with open('data/countries.geo.json', 'r') as fobj:
    countries = json.loads(fobj.read())

m = Map(zoom=3)

geo = GeoJSON(data=countries, style={'fillColor': 'white', 'weight': 0.5}, hover_style={'fillColor': '#1f77b4'}, name='Countries')
m.add_layer(geo)

widget_control1 = WidgetControl(widget=figure, position='bottomright')

m.add_control(widget_control1)

def on_hover(event, feature, **kwargs):
    global country_name

    country_name = feature['properties']['name']
    update_figure(country_name, data_name)

geo.on_hover(on_hover)

dropdown = Dropdown(
    options=['income', 'population', 'lifeExpectancy'],
    value=data_name,
    description='Plotting:'
)

def on_click(change):
    global data_name

    data_name = change['new']
    update_figure(country_name, data_name)

dropdown.observe(on_click, 'value')

widget_control2 = WidgetControl(widget=dropdown, position='bottomleft')

m.add_control(widget_control2)

m