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

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

# In[1]:


import json

import requests

url = (
    "https://raw.githubusercontent.com/python-visualization/folium/master/examples/data"
)
us_states = f"{url}/us-states.json"

geo_json_data = json.loads(requests.get(us_states).text)


# 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[2]:


import folium


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

folium.GeoJson(geo_json_data).add_to(m)

m


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

# In[3]:


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

folium.GeoJson(us_states).add_to(m)

m


# You can pass a geopandas object.

# In[4]:


import geopandas

gdf = geopandas.read_file(us_states)

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

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

m


# ### Click on zoom
# 
# You can enable an option that if you click on a part of the geometry the map will zoom in to that.
# 
# Try it on the map below:

# In[5]:


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

folium.GeoJson(geo_json_data, zoom_on_click=True).add_to(m)

m


# ### Styling
# 
# 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


# 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


# 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 = f"{url}/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_09.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


# Of course, if you can create and/or use a dictionary 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)


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

# In[14]:


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

m


# # Using `Choropleth` class
# 
# Now if you want to get faster, you can use the `Choropleth` class. Have a look at it's docstring, it has several styling options.
# 
# Just like the `GeoJson` class you can provide it a filename, a dict, or a geopandas object.

# In[15]:


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

folium.Choropleth(
    geo_data=us_states,
    fill_opacity=0.3,
    line_weight=2,
).add_to(m)

m


# Then, in playing with keyword arguments, you can get a choropleth in a few lines:

# In[16]:


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

folium.Choropleth(
    geo_data=us_states,
    data=unemployment,
    columns=["State", "Unemployment"],
    key_on="feature.id",
).add_to(m)

m


# You can force the color scale to a given number of bins (or directly list the bins you would like), by providing the `bins` argument.

# In[17]:


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

folium.Choropleth(
    geo_data=us_states,
    data=unemployment,
    columns=["State", "Unemployment"],
    key_on="feature.id",
    fill_color="YlGn",
    bins=[3, 4, 9, 11],
).add_to(m)

m


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

# In[18]:


state_data = pd.read_csv(US_Unemployment_Oct2012)

m = folium.Map(location=[48, -102], zoom_start=3)
folium.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,
).add_to(m)

m


# You can customize the way missing and `nan` values are displayed on your map using the two parameters `nan_fill_color` and `nan_fill_opacity`.

# In[19]:


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

messed_up_data = unemployment.drop(0)
messed_up_data.loc[4, "Unemployment"] = float("nan")

folium.Choropleth(
    geo_data=us_states,
    data=messed_up_data,
    columns=["State", "Unemployment"],
    nan_fill_color="purple",
    nan_fill_opacity=0.4,
    key_on="feature.id",
    fill_color="YlGn",
).add_to(m)

m


# Internally Choropleth uses the `GeoJson` or `TopoJson` class, depending on your settings, and the `StepColormap` class. Both objects are attributes of your `Choropleth` object called `geojson` and `color_scale`. You can make changes to them, but for regular things you won't have to. For example setting a name for in the layer controls or disabling showing the layer on opening the map is possible in `Choropleth` itself.

# In[20]:


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

choropleth = folium.Choropleth(
    geo_data=us_states,
    data=state_data,
    columns=["State", "Unemployment"],
    key_on="feature.id",
    fill_color="YlGn",
    name="Unenployment",
    show=False,
).add_to(m)

# The underlying GeoJson and StepColormap objects are reachable
print(type(choropleth.geojson))
print(type(choropleth.color_scale))

folium.LayerControl(collapsed=False).add_to(m)

m