"""Extract Hajodae/Gisamun DSM and OSM coast using .tools/terrain dependencies.

Run: python scripts/prepare-hajodae.py [hajodae|gisamun]
The 10m grid is interpolated from 30m DSM, not higher-accuracy survey data.
"""
import hashlib
import json
from pathlib import Path
import subprocess
import sys
from contextlib import ExitStack

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / '.tools/terrain'))
import numpy as np
import rasterio
from rasterio.features import rasterize
from rasterio.merge import merge
from rasterio.transform import from_origin
from rasterio.warp import reproject, Resampling
from pyproj import Transformer


def main():
    site = sys.argv[1] if len(sys.argv) > 1 else 'hajodae'
    if site not in ('hajodae', 'gisamun'):
        raise ValueError('Expected hajodae or gisamun')
    config = json.loads(subprocess.check_output(['node', '--input-type=module', '-e',
        "import {COAST} from './src/constants/coast.ts'; console.log(JSON.stringify(COAST))"], cwd=ROOT))
    if site == 'gisamun':
        config.update(json.loads(subprocess.check_output(['node', '--input-type=module', '-e',
            "import {GISAMUN} from './src/constants/gisamun.ts'; console.log(JSON.stringify(GISAMUN))"], cwd=ROOT)))
    lat, lon = config['latitude'], config['longitude']
    radius, step = config['radiusMeters'], config['spacingMeters']
    size = 2 * radius // step + 1
    projection = f'+proj=aeqd +lat_0={lat} +lon_0={lon} +datum=WGS84 +units=m +no_defs'
    forward = Transformer.from_crs('EPSG:4326', projection, always_xy=True)
    inverse = Transformer.from_crs(projection, 'EPSG:4326', always_xy=True)
    features = json.loads((ROOT / f'data/references/{site}/coastal-features.geojson').read_text(encoding='utf-8'))['features']
    coast = [f for f in features if f['properties']['category'] == 'coastline']
    mainland = next(f['geometry']['coordinates'] for f in coast if f['id'] == 'way/583161999')
    mainland_ids = {'way/583161999'}
    if site == 'gisamun':
        south = next(f['geometry']['coordinates'] for f in coast if f['id'] == 'way/22784123')
        assert south[-1] == mainland[0], 'Mainland segments must join exactly'
        mainland = south[:-1] + mainland
        mainland_ids.add('way/22784123')
    # This source runs south -> north; land is on its west side. Close well west
    # of the crop. These artificial closing edges never cross the published area.
    mainland = mainland + [[128.6, mainland[-1][1]], [128.6, mainland[0][1]], mainland[0]]
    geographic_land = [mainland] + [f['geometry']['coordinates'] for f in coast if f['id'] not in mainland_ids]
    beaches = [f['geometry']['coordinates'][0] for f in features if f['properties']['category'] == 'beach']
    def project(ring):
        return [list(forward.transform(*p)) for p in ring]
    land = [project(r) for r in geographic_land]
    sand = [project(r) for r in beaches]
    if site == 'gisamun':
        # Missing OSM beach polygon: approximate only this known beach segment.
        edge = project([p for p in mainland if config['sandLatitudeMin'] <= p[1] <= config['sandLatitudeMax']])
        assert len(edge) >= 2
        sand.append(edge + [[x-config['sandWidthMeters'], y] for x,y in reversed(edge)] + [edge[0]])
    lines = [project(f['geometry']['coordinates']) for f in coast]
    def shapes(rings):
        return [({'type': 'Polygon', 'coordinates': [r]}, 1) for r in rings]
    def mask(rings, spacing):
        n = 2 * radius // spacing + 1
        return rasterize(shapes(rings), out_shape=(n, n),
            transform=from_origin(-radius-spacing/2, radius+spacing/2, spacing, spacing), dtype='uint8')
    land_grid = mask(land, step)
    beach_grid = mask(sand, step)
    sources = [ROOT / f'data/raw/hajodae/Copernicus_DSM_COG_10_N{n}_00_E128_00_DEM.tif' for n in ([37,38] if site == 'gisamun' else [38])]
    heights = np.full((size, size), np.nan, dtype=np.float32)
    with ExitStack() as stack:
        datasets = [stack.enter_context(rasterio.open(p)) for p in sources]
        mosaic, transform = merge(datasets)
        reproject(mosaic[0], heights, src_transform=transform, src_crs=datasets[0].crs,
            dst_transform=from_origin(-radius-step/2, radius+step/2, step, step),
            dst_crs=projection, dst_nodata=np.nan, resampling=Resampling.bilinear)
    assert np.isfinite(heights).all(), 'Missing DSM coverage'
    # Preserve the resampled Copernicus heights independently of OSM masks.
    heights = heights.astype('<f4')
    assert land_grid[size//2, size//2], 'Origin must lie on mapped land'
    if site == 'gisamun':
        assert beach_grid[size//2, size//2], 'Gisamun origin must lie on the mapped beach'
    output = ROOT / f'public/terrain/{site}'
    output.mkdir(parents=True, exist_ok=True)
    heights.tofile(output/'heights.f32')
    fine_land, fine_sand = mask(land, config['maskSpacingMeters']), mask(sand, config['maskSpacingMeters'])
    # B encodes signed distance to vector coastline (positive inland), not
    # distance to raster stair steps. Work only within each segment's band.
    spacing = config['maskSpacingMeters']
    limit = config['distanceRangeMeters']
    n = fine_land.shape[0]
    fine_distance = np.full((n, n), float(limit))
    for line in lines:
        for (ax, ay), (bx, by) in zip(line, line[1:]):
            dx, dy = bx-ax, by-ay
            length2 = dx*dx+dy*dy
            if not length2:
                continue
            left = max(0, int(np.floor((min(ax,bx)-limit+radius)/spacing)))
            right = min(n, int(np.ceil((max(ax,bx)+limit+radius)/spacing))+1)
            top = max(0, int(np.floor((radius-max(ay,by)-limit)/spacing)))
            bottom = min(n, int(np.ceil((radius-min(ay,by)+limit)/spacing))+1)
            if right <= left or bottom <= top:
                continue
            fx, fy = np.meshgrid(np.arange(left,right)*spacing-radius, radius-np.arange(top,bottom)*spacing)
            t = np.clip(((fx-ax)*dx+(fy-ay)*dy)/length2, 0, 1)
            fine_distance[top:bottom,left:right] = np.minimum(fine_distance[top:bottom,left:right], np.hypot(fx-ax-t*dx,fy-ay-t*dy))
    signed = fine_distance * np.where(fine_land, 1, -1)
    encoded = np.rint(np.clip(.5+signed/(2*limit),0,1)*255).astype('uint8')
    # R land, G sand, B signed coast distance. Existing RG consumers stay valid.
    with rasterio.open(output/'surface.png', 'w', driver='PNG', width=fine_land.shape[1], height=fine_land.shape[0], count=3, dtype='uint8') as png:
        png.write(np.stack([fine_land*255, fine_sand*255, encoded]))
    manifest = json.loads((ROOT/f'data/references/{site}/metadata.json').read_text(encoding='utf-8'))
    used = [s for s in manifest['sources'] if Path(s['file']).name in [p.name for p in sources] or s['file'].endswith('.osm')]
    for item in used:
        assert hashlib.sha256((ROOT/item['file']).read_bytes()).hexdigest() == item['sha256']
    osm_license = {'license': 'ODbL-1.0', 'licenseUrl': 'https://opendatacommons.org/licenses/odbl/1-0/',
        'attribution': '© OpenStreetMap contributors', 'sourceUrl': 'https://www.openstreetmap.org/copyright'}
    meta = {
        'version': 1, 'dataset': f'Copernicus GLO-30 + OSM {site} coast',
        'surfaceType': 'Resampled Copernicus DSM; no OSM height correction', 'sourceResolution': '1 arc-second (~30m); interpolated to 10m',
        'verticalDatum': 'EGM2008 orthometric height (m); sea level approximated as 0m',
        'origin': {'latitude': lat, 'longitude': lon}, 'projection': projection,
        'width': size, 'height': size, 'spacingMeters': step, 'radiusMeters': radius,
        'observerSurfaceElevation': float(heights[size//2,size//2]),
        'minElevation': float(heights.min()), 'maxElevation': float(heights.max()), 'missingSamples': 0,
        'geographicCorners': [inverse.transform(x,y) for x,y in [(-radius,-radius),(radius,-radius),(radius,radius),(-radius,radius)]],
        'sources': used, 'attribution': '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; © OpenStreetMap contributors (ODbL-1.0)',
        'corrections': {'osmHeightCorrection': False, 'note': 'Heights use Copernicus reprojection and bilinear resampling only. OSM masks affect rendering and movement, not stored elevations. Offshore values are not bathymetry. No tide model.'},
        'files': {
            'heights.f32': {'source': 'Copernicus GLO-30', 'processing': 'AEQD projection and bilinear resampling only; no OSM height adjustment',
                'licenseUrl': 'https://dataspace.copernicus.eu/sites/default/files/media/files/2025-06/copernicus_contributing_mission_data_access_v2_cop_dem_licenses.pdf', 'licenseSection': 'GLO-30-F, Articles 4 and 6', 'notice': 'source.html'},
            'coast.json': {**osm_license, 'processing': 'Projected coastline and beach polygons; estimated beach strip for Gisamun'},
            'surface.png': {**osm_license, 'processing': 'R: land mask; G: sand mask; B: signed coast distance; contains no Copernicus elevations'},
        },
    }
    if site == 'gisamun':
        meta['sandApproximation'] = {key: config[key] for key in ['sandLatitudeMin','sandLatitudeMax','sandWidthMeters']}
        meta['sandApproximation']['note'] = 'OSM beach polygon absent. Visual strip west of OSM shoreline; not surveyed beach limits.'
    (output/'metadata.json').write_text(json.dumps(meta, ensure_ascii=False, indent=2)+'\n', encoding='utf-8')
    def local(rings):
        return [[[round(x,3),round(-y,3)] for x,y in r] for r in rings]
    (output/'coast.json').write_text(json.dumps({**osm_license, 'origin': meta['origin'], 'radiusMeters': radius,
        'land': local(land), 'beaches': local(sand), 'shorelines': local(lines)}, separators=(',',':')), encoding='utf-8')
    print(json.dumps({'grid': size, 'elevation': meta['observerSurfaceElevation'], 'files': {p.name:p.stat().st_size for p in output.iterdir()}}, indent=2))


if __name__ == '__main__':
    main()