%load_ext lab_black
%matplotlib inline

from IPython.display import HTML
import requests
import zipfile
import io
from datetime import timedelta, date
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import pandas as pd
import geopandas as gpd
import contextily as ctx
from PIL import Image

url_dep = "http://osm13.openstreetmap.fr/~cquest/openfla/export/departements-20140306-5m-shp.zip"
covid_url = (
    "https://www.data.gouv.fr/fr/datasets/r/63352e38-d353-4b54-bfd1-f1b3ee1cabd7"
)
filter_dep = ["971", "972", "973", "974", "976"]  # only metropolitan France
figsize = (15, 15)
tile_zoom = 7
frame_duration = 500  # in ms

local_path = "tmp/"
r = requests.get(url_dep)
z = zipfile.ZipFile(io.BytesIO(r.content))
z.extractall(path=local_path)
filenames = [
    y
    for y in sorted(z.namelist())
    for ending in ["dbf", "prj", "shp", "shx"]
    if y.endswith(ending)
]
dbf, prj, shp, shx = [filename for filename in filenames]
fr = gpd.read_file(local_path + shp)  #  + encoding="utf-8" if needed
fr.crs = "epsg:4326"  # {'init': 'epsg:4326'}
met = fr.query("code_insee not in @filter_dep")
met.set_index("code_insee", inplace=True)
met = met["geometry"]

w, s, e, n = met.total_bounds
bck, ext = ctx.bounds2img(w, s, e, n, zoom=tile_zoom, ll=True)

def save_img(df, title, img_name, vmin, vmax):
    gdf = gpd.GeoDataFrame(df, crs={"init": "epsg:4326"})
    gdf_3857 = gdf.to_crs(epsg=3857)  # web mercator
    f, ax = plt.subplots(figsize=figsize)
    ax.imshow(
        bck, extent=ext, interpolation="sinc", aspect="equal"
    )  # load background map
    divider = make_axes_locatable(ax)
    cax = divider.append_axes(
        "right", size="5%", pad=0.1
    )  # GeoPandas trick to adjust the legend bar
    gdf_3857.plot(
        column="hosp",  # Number of people currently hospitalized
        ax=ax,
        cax=cax,
        alpha=0.75,
        edgecolor="k",
        legend=True,
        cmap=matplotlib.cm.get_cmap("magma_r"),
        vmin=vmin,
        vmax=vmax,
    )

    ax.set_axis_off()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)
    ax.set_title(title, fontsize=25)
    plt.savefig(img_name, bbox_inches="tight")  # pad_inches=-0.1 to remove border
    plt.close(f)

cov = pd.read_csv(covid_url, sep=";", index_col=2, parse_dates=True,)
cov = cov.query("sexe == 0")  # sum of male/female
cov = cov.query("dep not in @filter_dep")
cov.dropna(inplace=True)
cov.head()

cov["geometry"] = cov["dep"].map(met)

def daterange(date1, date2):
    for n in range(int((date2 - date1).days) + 1):
        yield date1 + timedelta(n)

vmax = cov.hosp.max()
for i, dt in enumerate(daterange(cov.index.min(), cov.index.max())):
    title = dt.strftime("%d-%b-%Y")
    df = cov.query("jour == @dt")
    df = df.drop_duplicates(subset=["dep"], keep="first")
    img_name = "img/" + str(i) + ".png"
    save_img(df, title, img_name, 0, vmax)

frames = []
for i, dt in enumerate(daterange(cov.index.min(), cov.index.max())):
    name = "img/" + str(i) + ".png"
    frames.append(Image.open(name))

frames[0].save(
    "covid.gif",
    format="GIF",
    append_images=frames[1:],
    save_all=True,
    duration=frame_duration,
    loop=0,
)

from IPython.display import HTML

HTML("<img src='covid.gif'>")