#!/usr/bin/env python3
"""Build an island-preserving GSHHG mask and repack existing exact counts."""
import argparse
import gzip
import hashlib
import json
import os
from pathlib import Path
import shutil
import zipfile
import fiona
import numpy as np
from rasterio.features import rasterize
from rasterio.transform import from_origin
from land_mask import WIDTH, HEIGHT, load_mask, mask_metadata

ROOT = Path(__file__).resolve().parents[1]


def polar_geometry(feature):
    geometry = feature['geometry'].__geo_interface__
    if feature['properties']['id'] == '4-E' and feature['properties']['level'] == 5:
        # GSHHG 2.3.7's eastern Antarctic shapefile repeats (0,-90)
        # instead of closing along the pole. Its diagonal closure drops
        # most of East Antarctica when rasterized. Preserve the coastline
        # and explicitly close the eastern half through the South Pole.
        ring = geometry['coordinates'][0]
        coast = [p for p in ring if p[1] > -90]
        if coast[-1] == coast[0]: coast.pop()
        while len(coast) > 1 and coast[0] == coast[1]: coast.pop(0)
        assert coast[0][0] == 0 and coast[-1][0] == 180
        geometry = {'type': 'Polygon', 'coordinates': [coast + [(180, -90), (0, -90), coast[0]]]}
    return geometry


def build_mask(archive, output):
    # Extra dateline column preserves cells touching either +/-180 shoreline.
    shape = (HEIGHT, WIDTH + 1)
    transform = from_origin(-180.125, 90.125, .25, .25)
    land = np.zeros(shape, dtype=np.uint8)
    with zipfile.ZipFile(archive) as z:
        for level in [1, 5, 2, 3, 4]:
            suffix = f'/f/GSHHS_f_L{level}.shp'
            name = next(n for n in z.namelist() if n.endswith(suffix))
            with fiona.open(f'zip://{archive.resolve()}!{name}') as source:
                geometries = [polar_geometry(f) for f in source]
            if level in [1, 5, 3]:
                rasterize(geometries, out=land, transform=transform, all_touched=True, default_value=1)
            else:
                # Remove lake interiors only. Cells crossed by a shoreline still
                # contain land, so protect them, even if their centre is water.
                water = rasterize(geometries, out_shape=shape, transform=transform, dtype='uint8')
                boundaries = []
                for g in geometries:
                    rings = g['coordinates'] if g['type'] == 'Polygon' else [r for p in g['coordinates'] for r in p]
                    boundaries.append({'type': 'MultiLineString', 'coordinates': rings})
                edge = rasterize(boundaries, out_shape=shape, transform=transform, all_touched=True, dtype='uint8')
                land[(water == 1) & (edge == 0)] = 0
            print(f'MASK level {level}: {len(geometries):,} polygons', flush=True)
    land[:, 0] |= land[:, -1]
    # Entire deep polar cap must survive, across both hemispheres and the
    # dateline. A check only at longitude 0 misses malformed polar closures.
    assert land[int((90 + 85) * 4):, :WIDTH].all(), 'Incomplete Antarctic polar cap'
    packed = np.packbits(land[:, :WIDTH].reshape(-1), bitorder='little').tobytes()
    output.write_bytes(gzip.compress(packed, compresslevel=9, mtime=0))
    return load_mask(output)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--shorelines', type=Path, default=ROOT / '.local/gshhg-shp-2.3.7.zip')
    parser.add_argument('--counts', type=Path, default=ROOT / '.local/full-grid-cutoffs')
    parser.add_argument('--output', type=Path, default=ROOT / 'dist/data/land')
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    mask_path = args.output / 'land-mask.bin.gz'
    mask = build_mask(args.shorelines, mask_path)
    manifest = json.loads((args.counts / 'manifest.json').read_text())
    if manifest['version'] != 2:
        raise ValueError('Input must contain the original complete grids')
    original_bytes = manifest['totalBytes']
    manifest.update(version=3, encoding='gzip little-endian uint16; retained land cells only, in mask order',
                    mask=mask_metadata(mask_path))
    for key, layer in manifest['layers'].items():
        full_raw = gzip.decompress((args.counts / layer['file']).read_bytes())
        assert hashlib.sha256(full_raw).hexdigest() == layer['sha256']
        values = np.frombuffer(full_raw, dtype='<u2')[mask]
        raw = values.tobytes()
        data = gzip.compress(raw, compresslevel=9, mtime=0)
        path = args.output / layer['file']
        path.parent.mkdir(parents=True, exist_ok=True)
        temporary = path.with_suffix('.tmp')
        temporary.write_bytes(data)
        np.testing.assert_array_equal(np.frombuffer(gzip.decompress(temporary.read_bytes()), dtype='<u2'), values)
        os.replace(temporary, path)
        layer.update(bytes=len(data), uncompressedBytes=len(raw), sha256=hashlib.sha256(raw).hexdigest(),
                     min=float(values.min() / layer['divisor']), max=float(values.max() / layer['divisor']))
    manifest['totalBytes'] = sum(l['bytes'] for l in manifest['layers'].values()) + manifest['mask']['bytes']
    (args.output / 'manifest.json').write_text(json.dumps(manifest, indent=2) + '\n')
    with zipfile.ZipFile(args.shorelines) as z:
        for name in z.namelist():
            if Path(name).name in ['LICENSE.TXT', 'COPYING.LESSERv3', 'COPYINGv3', 'README.TXT']:
                (args.output / ('GSHHG-' + Path(name).name)).write_bytes(z.read(name))
    # Supply the mask's transformation source alongside the derived bitset.
    for name in ['pack_land_only.py', 'land_mask.py']:
        shutil.copyfile(Path(__file__).with_name(name), args.output / name)
    print(json.dumps({'landCells': int(mask.sum()), 'totalCells': len(mask), 'maskBytes': mask_path.stat().st_size,
                      'beforeBytes': original_bytes, 'afterBytes': manifest['totalBytes'],
                      'reductionPercent': 100 * (1 - manifest['totalBytes'] / original_bytes)}, indent=2), flush=True)


if __name__ == '__main__':
    main()