#!/usr/bin/env python
# coding: utf-8

# In[1]:


import ipywidgets as widgets

import geopandas as gpd

from lets_plot.geo_data import *
from lets_plot import *
LetsPlot.setup_html()


# In[2]:


class InteractiveGeocoder(object):
    LEVELS = ['country', 'state', 'county', 'city']
    RES = {'country': 3, 'state': 6, 'county': 9}

    def __init__(self, plot_width, plot_height):
        self.plot_width = plot_width
        self.plot_height = plot_height
        self.emulated = False
        self.select = {'type': None, 'value': None}
        self._init_level_widgets()
        self.wr = widgets.IntSlider(value=1, min=1, max=15, step=1, description='Resolution:')
        self.wo = widgets.Output(layout=widgets.Layout(height='{0}px'.format(self.plot_height + 20)))
        self._observe_widgets()

    def _init_level_widgets(self):
        self.wl = {}
        for level in self.LEVELS:
            self.wl[level] = widgets.Dropdown(options=[], description='{0}:'.format(level.title()))
        self.wl['country'].options = [''] + \
                                     geocode_countries().get_geocodes()\
                                                        .sort_values('found name')['found name']\
                                                        .to_list()

    def _observe_widgets(self):
        self.wl['country'].observe(self._on_change_select)
        self.wl['state'].observe(self._on_change_select)
        self.wl['county'].observe(self._on_change_select)
        self.wr.observe(self._on_change_slider)

    def _update_output(self, res_value=None):
        self.wo.outputs = ()
        if not self.select['type'] or not self.select['value']:
            return
        for i, level in enumerate(self.LEVELS[:-1]):
            if self.select['type'] != level or not self.wl[level].value:
                continue
            self._clear_options_for_lower_widgets(i)
            res_value = res_value or self.RES[level]
            p = self._get_plot(*(self._geocode(i, level, self._get_scope(i), res_value)))
            if p:
                self.wo.append_display_data(p)
            break
        self.emulated = True
        self.wr.value = res_value

    def _get_scope(self, level_id):
        scope = None
        for i in range(0, level_id):
            scope = geocode(level=self.LEVELS[i], names=self.wl[self.LEVELS[i]].value, scope=scope).ignore_all_errors()
        return scope

    def _clear_options_for_lower_widgets(self, level_id):
        for i in range(level_id+1, len(self.LEVELS)-1):
            self.emulated = True
            self.wl[self.LEVELS[i]].options = []

    def _geocode(self, i, level, scope, res_value):
        if self.select['value'] == '':
            return geocode(level=level, scope=scope).ignore_all_errors().get_boundaries(res_value), \
                   gpd.GeoDataFrame()
        b_gdf = geocode(level=self.LEVELS[i+1], scope=self.wl[level].value).ignore_all_errors().get_boundaries(res_value)
        p_gdf = gpd.GeoDataFrame()
        if self.LEVELS[i+1] in self.wl.keys():
            self.emulated = True
            self.wl[self.LEVELS[i+1]].options = [''] + b_gdf.sort_values('found name')['found name'].to_list()
        if b_gdf.empty or level == 'county':
            geocoded_level = geocode(level=level, names=self.wl[level].value, scope=scope).allow_ambiguous().ignore_all_errors()
            p_gdf = geocode(level='city', scope=geocoded_level).ignore_all_errors().get_centroids()
            if b_gdf.empty:
                b_gdf = geocoded_level.get_boundaries(res_value)
        return b_gdf, p_gdf

    def _get_plot(self, b_gdf, p_gdf):
        if b_gdf.empty:
            return None
        p = ggplot() + ggsize(self.plot_width, self.plot_height) + \
            theme_void()
        if p_gdf.empty:
            p += geom_map(data=b_gdf, fill='black', color='white', tooltips=layer_tooltips().line('@{found name}'))
        else:
            p += geom_map(data=b_gdf, fill='black', color='white') + \
                 geom_point(data=p_gdf, shape=1, color='white', tooltips=layer_tooltips().line('@{found name}'))
        return p

    def _on_change_select(self, change):
        if change['type'] != 'change' or change['name'] != 'value':
            return
        if self.emulated:
            self.emulated = False
            return
        self.select['type'] = change.owner.description[:-1].lower()
        self.select['value'] = change['new']
        self._update_output()

    def _on_change_slider(self, change):
        if change['type'] != 'change' or change['name'] != 'value':
            return
        if self.emulated:
            self.emulated = False
            return
        self._update_output(change['new'])

    def display(self):
        display(self.wl['country'], self.wl['state'], self.wl['county'], self.wr, self.wo)


# In[3]:


geocoder = InteractiveGeocoder(400, 300)
geocoder.display()