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

# In[1]:


import os
import folium

print(folium.__version__)


# # GeoJSON and choropleth
# 
# **A few examples of how to do that with `folium`.**
# 
# 
# ## Using `GeoJson`
# 
# 
# Let us load a GeoJSON file representing the US states.

# In[2]:


import json

us_states = os.path.join('data', 'us-states.json')

geo_json_data = json.load(open(us_states))


# It is a classical GeoJSON `FeatureCollection` (see https://en.wikipedia.org/wiki/GeoJSON) of the form :
# 
#     {
#         "type": "FeatureCollection",
#         "features": [
#             {
#                 "properties": {"name": "Alabama"},
#                 "id": "AL",
#                 "type": "Feature",
#                 "geometry": {
#                     "type": "Polygon",
#                     "coordinates": [[[-87.359296, 35.00118], ...]]
#                     }
#                 },
#             {
#                 "properties": {"name": "Alaska"},
#                 "id": "AK",
#                 "type": "Feature",
#                 "geometry": {
#                     "type": "MultiPolygon",
#                     "coordinates": [[[[-131.602021, 55.117982], ... ]]]
#                     }
#                 },
#             ...
#             ]
#         }
# 

# A first way of drawing it on a map, is simply to use `folium.GeoJson` :

# In[3]:


m = folium.Map([43, -100], zoom_start=4)

folium.GeoJson(geo_json_data).add_to(m)

m.save(os.path.join('results', 'GeoJSON_and_choropleth_0.html'))

m


# Note that you can avoid loading the file on yourself ; in simply providing a file path.

# In[4]:


m = folium.Map([43, -100], zoom_start=4)

folium.GeoJson(us_states).add_to(m)

m.save(os.path.join('results', 'GeoJSON_and_choropleth_1.html'))

m


# You can pass a geopandas object.

# In[5]:


import geopandas

gdf = geopandas.read_file(us_states)

m = folium.Map([43, -100], zoom_start=4)

folium.GeoJson(
    gdf,
).add_to(m)

m.save(os.path.join('results', 'GeoJSON_and_choropleth_3.html'))

m


# Now this is cool and simple, but we may be willing to choose the style of the data.
# 
# You can provide a function of the form `lambda feature: {}` that sets the style of each feature.
# 
# For possible options, see:
# 
# * For `Point` and `MultiPoint`, see http://leafletjs.com/reference.html#marker
# * For other features, see http://leafletjs.com/reference.html#path-options and http://leafletjs.com/reference.html#polyline-options
# 

# In[6]:


m = folium.Map([43, -100], zoom_start=4)

folium.GeoJson(
    geo_json_data,
    style_function=lambda feature: {
        'fillColor': '#ffff00',
        'color': 'black',
        'weight': 2,
        'dashArray': '5, 5'
    }
).add_to(m)

m.save(os.path.join('results', 'GeoJSON_and_choropleth_3.html'))

m


# What's cool in providing a function, is that you can specify a style depending on the feature. For example, if you want to visualize in green all states whose name contains the letter 'E', just do:

# In[7]:


m = folium.Map([43, -100], zoom_start=4)

folium.GeoJson(
    geo_json_data,
    style_function=lambda feature: {
        'fillColor': 'green' if 'e' in feature['properties']['name'].lower() else '#ffff00',
        'color': 'black',
        'weight': 2,
        'dashArray': '5, 5'
    }
).add_to(m)

m.save(os.path.join('results', 'GeoJSON_and_choropleth_4.html'))

m


# Wow, this looks almost like a choropleth. To do one, we just need to compute a color for each state.
# 
# Let's imagine we want to draw a choropleth of unemployment in the US.
# 
# First, we may load the data:

# In[8]:


import pandas as pd

US_Unemployment_Oct2012 = os.path.join('data', 'US_Unemployment_Oct2012.csv')
unemployment = pd.read_csv(US_Unemployment_Oct2012)

unemployment.head(5)


# Now we need to create a function that maps one value to a RGB color (of the form `#RRGGBB`).
# For this, we'll use colormap tools from `folium.colormap`.

# In[9]:


from branca.colormap import linear

colormap = linear.YlGn.scale(
    unemployment.Unemployment.min(),
    unemployment.Unemployment.max())

print(colormap(5.0))

colormap


# We need also to convert the table into a dictionnary, in order to map a feature to it's unemployment value.

# In[10]:


unemployment_dict = unemployment.set_index('State')['Unemployment']

unemployment_dict['AL']


# Now we can do the choropleth.

# In[11]:


m = folium.Map([43, -100], zoom_start=4)

folium.GeoJson(
    geo_json_data,
    name='unemployment',
    style_function=lambda feature: {
        'fillColor': colormap(unemployment_dict[feature['id']]),
        'color': 'black',
        'weight': 1,
        'dashArray': '5, 5',
        'fillOpacity': 0.9,
    }
).add_to(m)

folium.LayerControl().add_to(m)

m.save(os.path.join('results', 'GeoJSON_and_choropleth_5.html'))

m


# Of course, if you can create and/or use a dictionnary providing directly the good color. Thus, the finishing seems faster:

# In[12]:


color_dict = {key: colormap(unemployment_dict[key]) for key in unemployment_dict.keys()}


# In[13]:


m = folium.Map([43, -100], zoom_start=4)

folium.GeoJson(
    geo_json_data,
    style_function=lambda feature: {
        'fillColor': color_dict[feature['id']],
        'color': 'black',
        'weight': 1,
        'dashArray': '5, 5',
        'fillOpacity': 0.9,
    }
).add_to(m)

m.save(os.path.join('results', 'GeoJSON_and_choropleth_6.html'))

m


# Note that adding a color legend may be a good idea.

# In[14]:


colormap.caption = 'Unemployment color scale'
colormap.add_to(m)

m.save(os.path.join('results', 'GeoJSON_and_choropleth_7.html'))

m


# # Using `choropleth` method
# 
# Now if you want to get faster, you can use the `Map.choropleth` method. Have a look at it's docstring, it has several styling options.
# 
# You can use it in providing a file name (`geo_path`) :

# In[15]:


m = folium.Map([43, -100], zoom_start=4)

m.choropleth(
    geo_data=geopandas.read_file(us_states),
    fill_color='red',
    fill_opacity=0.3,
    line_weight=2,
)

m.save(os.path.join('results', 'GeoJSON_and_choropleth_7.html'))

m


# Or in providing a GeoJSON string (`geo_str`) :

# In[16]:


m = folium.Map([43, -100], zoom_start=4)

m.choropleth(geo_data=us_states)

m.save(os.path.join('results', 'GeoJSON_and_choropleth_8.html'))

m


# Then, in playing with keyword arguments, you can get a choropleth in (*almost*) one line :

# In[17]:


m = folium.Map([43, -100], zoom_start=4)

m.choropleth(
    geo_data=us_states,
    data=unemployment,
    columns=['State', 'Unemployment'],
    key_on='feature.id',
    fill_color='YlGn',
)


m.save(os.path.join('results', 'GeoJSON_and_choropleth_9.html'))

m


# A cool thing: you can force the color scale to a given number of bins, in providing a `threshold_scale` argument.

# In[18]:


m = folium.Map([43, -100], zoom_start=4)

m.choropleth(
    geo_data=us_states,
    data=unemployment,
    columns=['State', 'Unemployment'],
    key_on='feature.id',
    fill_color='YlGn',
    threshold_scale=[3, 4, 9, 10]
)


m.save(os.path.join('results', 'GeoJSON_and_choropleth_10.html'))

m


# You can also enable the highlight function, to enable highlight functionality when you hover over each area.

# In[19]:


state_data = pd.read_csv(US_Unemployment_Oct2012)

m = folium.Map(location=[48, -102], zoom_start=3)
m.choropleth(
    geo_data=us_states,
    data=state_data,
    columns=['State', 'Unemployment'],
    key_on='feature.id',
    fill_color='YlGn',
    fill_opacity=0.7,
    line_opacity=0.2,
    legend_name='Unemployment Rate (%)',
    highlight=True
)

m.save(os.path.join('results', 'GeoJSON_and_choropleth_11.html'))

m