#!/usr/bin/env python3
"""Fetch real elevation for a level's OSM area from the AWS terrain tiles.

Usage: fetch_dem.py <area.json> <out_prefix>

Reads the area's bounding box straight from its OSM geometry (the original
fetch bboxes were never recorded -- same trap as the junction specs), grabs
the covering terrarium tiles at z15 (~4.2 m/px at Brisbane's latitude), and
writes:
    <out_prefix>_dem.f32   row-major float32 heights in metres (little-endian)
    <out_prefix>_dem.json  {"w", "h", "min_lat", "min_lon", "max_lat", "max_lon"}

Terrarium encoding: height = R*256 + G + B/256 - 32768.
Tiles: https://s3.amazonaws.com/elevation-tiles-prod/terrarium/{z}/{x}/{y}.png
(open data, no key). Data (c) Mapzen/AWS Terrain Tiles contributors.
"""
import io
import json
import math
import struct
import sys
import urllib.request

from PIL import Image

Z = 15
URL = "https://s3.amazonaws.com/elevation-tiles-prod/terrarium/%d/%d/%d.png"


def tile_of(lat, lon, z):
    n = 2 ** z
    xt = (lon + 180.0) / 360.0 * n
    lat_r = math.radians(lat)
    yt = (1.0 - math.asinh(math.tan(lat_r)) / math.pi) / 2.0 * n
    return xt, yt


def tile_bounds(xt, yt, z):
    n = 2 ** z
    lon0 = xt / n * 360.0 - 180.0
    lat0 = math.degrees(math.atan(math.sinh(math.pi * (1 - 2 * yt / n))))
    return lat0, lon0


def main():
    src, prefix = sys.argv[1], sys.argv[2]
    d = json.load(open(src))
    lats, lons = [], []
    for e in d["elements"]:
        for g in e.get("geometry") or []:
            lats.append(g["lat"])
            lons.append(g["lon"])
        if e.get("type") == "node" and "lat" in e:
            lats.append(e["lat"])
            lons.append(e["lon"])
    lo_lat, hi_lat = min(lats), max(lats)
    lo_lon, hi_lon = min(lons), max(lons)

    x0, y1 = tile_of(lo_lat, lo_lon, Z)   # south-west -> larger y tile index
    x1, y0 = tile_of(hi_lat, hi_lon, Z)
    tx0, tx1 = int(x0), int(x1)
    ty0, ty1 = int(y0), int(y1)
    cols = tx1 - tx0 + 1
    rows = ty1 - ty0 + 1
    print("bbox lat %.4f..%.4f lon %.4f..%.4f -> %dx%d tiles @z%d"
          % (lo_lat, hi_lat, lo_lon, hi_lon, cols, rows, Z))

    mosaic = Image.new("RGB", (cols * 256, rows * 256))
    for ty in range(ty0, ty1 + 1):
        for tx in range(tx0, tx1 + 1):
            req = urllib.request.Request(URL % (Z, tx, ty),
                                         headers={"User-Agent": "BurnoutShitbox/1.0"})
            png = urllib.request.urlopen(req, timeout=60).read()
            mosaic.paste(Image.open(io.BytesIO(png)), ((tx - tx0) * 256, (ty - ty0) * 256))

    # mosaic geographic bounds (tile edges, not the request bbox)
    m_hi_lat, m_lo_lon = tile_bounds(tx0, ty0, Z)
    m_lo_lat, m_hi_lon = tile_bounds(tx1 + 1, ty1 + 1, Z)

    w, h = mosaic.size
    px = mosaic.load()
    out = open(prefix + "_dem.f32", "wb")
    lo, hi = 1e9, -1e9
    for yy in range(h):
        row = bytearray()
        for xx in range(w):
            r, g, b = px[xx, yy]
            hm = r * 256 + g + b / 256.0 - 32768.0
            hm = max(hm, 0.0)   # clamp bathymetry: the river reads as 0, not -8
            lo, hi = min(lo, hm), max(hi, hm)
            row += struct.pack("<f", hm)
        out.write(row)
    out.close()
    json.dump({"w": w, "h": h, "min_lat": m_lo_lat, "max_lat": m_hi_lat,
               "min_lon": m_lo_lon, "max_lon": m_hi_lon},
              open(prefix + "_dem.json", "w"))
    print("wrote %s_dem.f32 (%dx%d, %.1f..%.1f m)" % (prefix, w, h, lo, hi))


if __name__ == "__main__":
    main()