"""A heightmap PNG of any mountain, from the free Copernicus 30 m elevation model on AWS.

    python heightmap.py LAT LON WIDTH_KM HEIGHT_KM OUT.png

    python heightmap.py 35.875 76.52 18 13.5 k2.png        # K2, the map in the blog post

Needs rasterio, numpy and Pillow (pip install rasterio numpy pillow). The map has to lie inside one
1-degree tile of the model (the one containing LAT/LON). Prints the elevations that black and white stand for.

The model is Copernicus DEM GLO-30 (Copernicus WorldDEM-30). If you share the PNG, or a picture of it, the
licence asks for this notice (see NOTICE.txt next to this file):
"produced using Copernicus WorldDEM-30 © DLR e.V. 2010-2014 and © Airbus Defence and Space GmbH 2014-2018
provided under COPERNICUS by the European Union and ESA; all rights reserved", and
"The organisations in charge of the Copernicus programme by law or by delegation do not incur any liability for any
use of the Copernicus WorldDEM-30".
"""
import sys

import numpy as np
import rasterio
from PIL import Image
from rasterio.transform import Affine
from rasterio.warp import Resampling, reproject, transform
from rasterio.windows import Window, from_bounds

TILE = ("https://copernicus-dem-30m.s3.amazonaws.com/Copernicus_DSM_COG_10_N{lat:02d}_00_E{lon:03d}_00_DEM/"
        "Copernicus_DSM_COG_10_N{lat:02d}_00_E{lon:03d}_00_DEM.tif")
PIXEL = 30                                  # metres per pixel: the model's own resolution


def utm(lat, lon):
    """EPSG code of the UTM zone (northern hemisphere) a point lies in: a flat map in metres, no stretching."""
    return f"EPSG:{32600 + int((lon + 180) // 6) + 1}"


def to_xy(crs, lat, lon):
    """Latitude/longitude -> metres east and north in `crs`."""
    (x,), (y,) = transform("EPSG:4326", crs, [lon], [lat])
    return x, y


def frame(lat, lon, width_km, height_km):
    """The map's rectangle: its CRS, the transform from its pixels to metres, and its size in pixels."""
    crs = utm(lat, lon)
    cx, cy = to_xy(crs, lat, lon)
    w, h = round(width_km * 1000 / PIXEL), round(height_km * 1000 / PIXEL)
    return crs, Affine(PIXEL, 0, cx - w * PIXEL / 2, 0, -PIXEL, cy + h * PIXEL / 2), (w, h)


def elevation(lat, lon, width_km, height_km):
    """Metres above sea level for every pixel of the map, north up."""
    crs, to_metres, (w, h) = frame(lat, lon, width_km, height_km)
    d = max(width_km, height_km) / 111 + 0.02                   # degrees of margin around the rectangle
    out = np.zeros((h, w), "float32")
    with rasterio.open(TILE.format(lat=int(lat), lon=int(lon))) as src:
        win = from_bounds(lon - d, lat - d, lon + d, lat + d, src.transform).round_offsets().round_lengths()
        win = win.intersection(Window(0, 0, src.width, src.height))
        reproject(src.read(1, window=win), out, src_transform=src.window_transform(win), src_crs=src.crs,
                  dst_transform=to_metres, dst_crs=crs, resampling=Resampling.bilinear)
    return out


if __name__ == "__main__":
    lat, lon, w_km, h_km = map(float, sys.argv[1:5])
    z = elevation(lat, lon, w_km, h_km)
    lo, hi = float(z.min()), float(z.max())
    Image.fromarray(np.round((z - lo) / (hi - lo) * 65535).astype(np.uint16)).save(sys.argv[5])   # 16 bit
    print(f"{z.shape[1]} x {z.shape[0]} px, black = {lo:.0f} m, white = {hi:.0f} m")