#!/usr/bin/env python3
"""What the surface is, and how deep the sea is: land, ice, lakes and sea depth on a half-degree raster.

Reads the land raster already used for the coastlines (land_now.json, from Natural Earth's 1:110m land),
Natural Earth's 1:110m glaciated areas and lakes, and its 1:10m depth contours of the sea, and writes
ground_now.json with three rasters of yes and no (land, ice, lake) and one of depth: for each cell the
number of depth contours it lies inside (0 = shallower than 200 m, 1 = 200 to 1000 m, 2 = 1000 to 2000 m,
and so on by the thousand). Rows run from 90 N southward, columns from 180 W eastward.

Usage:  python3 build_ground.py [--fetch] NE_DIR land_now.json ground_now.json
        --fetch downloads the outlines into NE_DIR first (about 45 MB, nearly all of it the depth contours).
"""
import base64, json, os, sys, urllib.request, zlib
import numpy as np

NE = 'https://raw.githubusercontent.com/nvkelso/natural-earth-vector/master/geojson/'
FILES = {'ice': 'ne_110m_glaciated_areas.geojson', 'lake': 'ne_110m_lakes.geojson'}
DEPTHS = [('K', 200), ('J', 1000), ('I', 2000), ('H', 3000), ('G', 4000), ('F', 5000), ('E', 6000), ('D', 7000), ('C', 8000), ('B', 9000), ('A', 10000)]
DEPTH_FILE = 'ne_10m_bathymetry_{}_{}.geojson'


def rings_of(path):
    out = []
    for f in json.load(open(path))['features']:
        g = f['geometry']
        polys = g['coordinates'] if g['type'] == 'MultiPolygon' else [g['coordinates']]
        out.extend(polys)                      # each polygon: [outer, hole, hole, ...]
    return out


def raster(polys, res, nx, ny):
    lon = -180 + (np.arange(nx) + 0.5) * res
    lat = 90 - (np.arange(ny) + 0.5) * res
    grid = np.zeros((ny, nx), bool)
    for poly in polys:
        inside = np.zeros((ny, nx), bool)
        for ring in poly:                       # even-odd over outer ring and holes
            r = np.asarray(ring, float)
            x0, y0, x1, y1 = r[:-1, 0], r[:-1, 1], r[1:, 0], r[1:, 1]
            j0, j1 = np.searchsorted(-lat, -r[:, 1].max()), np.searchsorted(-lat, -r[:, 1].min())
            for j in range(max(0, j0 - 1), min(ny, j1 + 1)):
                y = lat[j]
                cross = (y0 > y) != (y1 > y)
                if not cross.any():
                    continue
                xs = x0[cross] + (y - y0[cross]) * (x1[cross] - x0[cross]) / (y1[cross] - y0[cross])
                inside[j] ^= (np.sum(xs[:, None] > lon[None, :], 0) % 2).astype(bool)
        grid |= inside
    return grid


def main():
    args = [a for a in sys.argv[1:] if a != '--fetch']
    ne, land_file, out = args
    if '--fetch' in sys.argv:
        os.makedirs(ne, exist_ok=True)
        for f in list(FILES.values()) + [DEPTH_FILE.format(*d) for d in DEPTHS]:
            if not os.path.exists(os.path.join(ne, f)):
                urllib.request.urlretrieve(NE + f, os.path.join(ne, f))
    d = json.load(open(land_file))
    res, nx, ny = d['res'], d['nx'], d['ny']
    land = np.unpackbits(np.frombuffer(base64.b64decode(d['bits']), dtype=np.uint8))[:ny * nx].reshape(ny, nx).astype(bool)
    ice = raster(rings_of(os.path.join(ne, FILES['ice'])), res, nx, ny)
    south = (90 - (np.arange(ny) + 0.5) * res) < -86   # the outlines stop just short of the South Pole, which stands on the ice sheet
    land[south, :] = True
    ice[south, :] = True
    lake = raster(rings_of(os.path.join(ne, FILES['lake'])), res, nx, ny)
    ice &= land
    lake &= land
    depth = np.zeros((ny, nx), np.uint8)
    for d in DEPTHS:
        depth += raster(rings_of(os.path.join(ne, DEPTH_FILE.format(*d))), res, nx, ny)
    depth[land] = 0
    w = np.cos(np.radians(90 - (np.arange(ny) + 0.5) * res))[:, None] * np.ones((1, nx))
    share = lambda m: round(float((w * m).sum() / w.sum()) * 100, 2)
    pack = lambda m: base64.b64encode(np.packbits(m.astype(np.uint8).ravel()).tobytes()).decode()
    sea = ~land
    classes = [round(float((w * (sea & (depth == k))).sum() / (w * sea).sum()) * 100, 2) for k in range(len(DEPTHS) + 1)]
    json.dump({'res': res, 'nx': nx, 'ny': ny, 'land': pack(land), 'ice': pack(ice), 'lake': pack(lake),
               'depth': base64.b64encode(zlib.compress(depth.tobytes(), 9)).decode(), 'depthLevels': [d[1] for d in DEPTHS],
               'shares': {'land': share(land), 'ice': share(ice), 'lake': share(lake), 'seaByDepthClass': classes},
               'source': 'Natural Earth 1:110m land, glaciated areas and lakes; 1:10m bathymetry'}, open(out, 'w'))
    print('land', share(land), 'ice', share(ice), 'lake', share(lake), '% of the surface; the sea by depth class, %:', classes)


main()