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

# # From NumPy to Leaflet

# This notebook shows how to display some raster geographic data in IPyLeaflet. The data is a NumPy array, which means that you have all the power of the Python scientific stack at your disposal to process it.
# 
# The following libraries are needed:
# * `requests`
# * `tqdm`
# * `rasterio`
# * `numpy`
# * `scipy`
# * `pillow`
# * `matplotlib`
# * `ipyleaflet`
# 
# The recommended way is to try to `conda install` them first, and if they are not found then `pip install`.

# In[ ]:


import requests
import os
from tqdm import tqdm
import zipfile
import rasterio
from affine import Affine
import numpy as np
import scipy.ndimage
from rasterio.warp import reproject, Resampling
import PIL
import matplotlib.pyplot as plt
from base64 import b64encode

try:
    from StringIO import StringIO

    py3 = False
except ImportError:
    from io import StringIO, BytesIO

    py3 = True
from ipyleaflet import Map, ImageOverlay, basemap_to_tiles, basemaps


# Download a raster file representing the flow accumulation for South America. This gives an idea of the river network.

# In[ ]:


url = "https://edcintl.cr.usgs.gov/downloads/sciweb1/shared/hydrosheds/sa_30s_zip_grid/sa_acc_30s_grid.zip"
filename = os.path.basename(url)
name = filename[: filename.find("_grid")]
adffile = name + "/" + name + "/w001001.adf"

if not os.path.exists(adffile):
    r = requests.get(url, stream=True)
    with open(filename, "wb") as f:
        total_length = int(r.headers.get("content-length"))
        for chunk in tqdm(
            r.iter_content(chunk_size=1024), total=(total_length / 1024) + 1
        ):
            if chunk:
                f.write(chunk)
                f.flush()
    zip = zipfile.ZipFile(filename)
    zip.extractall(".")


# We transform the data a bit so that rivers appear thicker.

# In[ ]:


dataset = rasterio.open(adffile)
acc_orig = dataset.read()[0]
acc = np.where(acc_orig < 0, 0, acc_orig)

shrink = 1  # if you are out of RAM try increasing this number (should be a power of 2)
radius = 5  # you can play with this number to change the width of the rivers
circle = np.zeros((2 * radius + 1, 2 * radius + 1)).astype("uint8")
y, x = np.ogrid[-radius : radius + 1, -radius : radius + 1]
index = x ** 2 + y ** 2 <= radius ** 2
circle[index] = 1
acc = np.sqrt(acc)
acc = scipy.ndimage.maximum_filter(acc, footprint=circle)
acc[acc_orig < 0] = np.nan
acc = acc[::shrink, ::shrink]


# The original data is in the WGS 84 projection, but Leaflet uses Web Mercator, so we need to reproject.

# In[ ]:


# At this point if GDAL complains about not being able to open EPSG support file gcs.csv, try in the terminal:
# export GDAL_DATA=`gdal-config --datadir`

with rasterio.Env():
    rows, cols = acc.shape
    src_transform = list(dataset.transform)
    src_transform[0] *= shrink
    src_transform[4] *= shrink
    src_transform = Affine(*src_transform[:6])
    src_crs = {"init": "EPSG:4326"}
    source = acc

    dst_crs = {"init": "EPSG:3857"}
    dst_transform, width, height = rasterio.warp.calculate_default_transform(
        src_crs, dst_crs, cols, rows, *dataset.bounds
    )
    dst_shape = height, width

    destination = np.zeros(dst_shape)

    reproject(
        source,
        destination,
        src_transform=src_transform,
        src_crs=src_crs,
        dst_transform=dst_transform,
        dst_crs=dst_crs,
        resampling=Resampling.nearest,
    )

acc_web = destination


# Let's convert our NumPy array to an image. For that we must specify a colormap (here `plt.cm.jet`).

# In[ ]:


acc_norm = acc_web - np.nanmin(acc_web)
acc_norm = acc_norm / np.nanmax(acc_norm)
acc_norm = np.where(np.isfinite(acc_web), acc_norm, 0)
acc_im = PIL.Image.fromarray(np.uint8(plt.cm.jet(acc_norm) * 255))
acc_mask = np.where(np.isfinite(acc_web), 255, 0)
mask = PIL.Image.fromarray(np.uint8(acc_mask), mode="L")
im = PIL.Image.new("RGBA", acc_norm.shape[::-1], color=None)
im.paste(acc_im, mask=mask)


# The image is embedded in the URL as a PNG file, so that it can be sent to the browser.

# In[ ]:


if py3:
    f = BytesIO()
else:
    f = StringIO()
im.save(f, "png")
data = b64encode(f.getvalue())
if py3:
    data = data.decode("ascii")
imgurl = "data:image/png;base64," + data


# Finally we can overlay our image and if everything went fine it should be exactly over South America.

# In[ ]:


b = dataset.bounds
bounds = [(b.bottom, b.left), (b.top, b.right)]
io = ImageOverlay(url=imgurl, bounds=bounds)


# In[ ]:


center = [-10, -60]
zoom = 2
m = Map(center=center, zoom=zoom, interpolation="nearest")
m


# In[ ]:


tile = basemap_to_tiles(basemaps.Esri.WorldStreetMap)
m.add_layer(tile)


# You can play with the opacity slider and check that rivers from our data file match the rivers on OpenStreetMap.

# In[ ]:


m.add_layer(io)
io.interact(opacity=(0.0, 1.0, 0.01))


# In[ ]: