GeoData/geodata_pipeline/heightmaps.py

from __future__ import annotations

import glob
import math
import os
from typing import Iterable

import numpy as np
from osgeo import gdal

from .config import Config
from .gdal_utils import build_vrt, cleanup_aux_files, ensure_dir, ensure_parent, open_dataset, safe_remove

gdal.UseExceptions()


def _cleanup_patterns(raw_dir: str) -> Iterable[str]:
    return [
        os.path.join("work", "*_tmp.tif"),
        os.path.join("work", "*_tmp.tif.aux.xml"),
        os.path.join("work", "*.aux.xml"),
        os.path.join(raw_dir, "*.aux.xml"),
    ]


def _compute_ds_minmax(ds: gdal.Dataset) -> tuple[float | None, float | None, int]:
    band = ds.GetRasterBand(1)
    nodata = band.GetNoDataValue()
    data = band.ReadAsArray()

    if data is None or data.size == 0:
        return None, None, 0

    valid = np.isfinite(data)
    if nodata is not None and not math.isnan(nodata):
        valid &= data != nodata

    valid_count = int(np.count_nonzero(valid))
    if valid_count == 0:
        return None, None, 0

    tile_min = float(data[valid].min())
    tile_max = float(data[valid].max())
    return tile_min, tile_max, valid_count


def export_heightmaps(cfg: Config, *, force_vrt: bool = False) -> int:
    ensure_dir(cfg.work.work_dir)
    ensure_dir(cfg.export.heightmap_dir)
    ensure_parent(cfg.export.manifest_path)

    tif_paths = sorted(glob.glob(os.path.join(cfg.raw.dgm1_dir, "*.tif")))
    build_vrt(cfg.work.heightmap_vrt, tif_paths, force=force_vrt)
    ds = open_dataset(cfg.work.heightmap_vrt, f"Could not open {cfg.work.heightmap_vrt} after attempting to build it.")

    band = ds.GetRasterBand(1)
    gmin, gmax = band.ComputeRasterMinMax(False)
    print(f"GLOBAL_MIN={gmin}, GLOBAL_MAX={gmax}")

    tile_key_cfg = cfg.tile_key
    tile_key_enabled = tile_key_cfg.enabled
    tile_key_seen: set[str] = set()

    with open(cfg.export.manifest_path, "w", encoding="utf-8") as f:
        f.write("tile_id,xmin,ymin,xmax,ymax,global_min,global_max,out_res,tile_key,tile_min,tile_max\n")

        skipped = 0
        written = 0

        for tif in tif_paths:
            try:
                tds = open_dataset(tif, f"Skipping unreadable {tif}")
            except SystemExit as exc:
                print(exc)
                skipped += 1
                continue

            gt = tds.GetGeoTransform()
            ulx, xres, _, uly, _, yres = gt

            xmax = ulx + xres * tds.RasterXSize
            ymin = uly + yres * tds.RasterYSize
            xmin = ulx
            ymax = uly

            base = os.path.splitext(os.path.basename(tif))[0]
            tile_id = base

            tile_key = ""
            if tile_key_enabled:
                x_key = math.floor((xmin + tile_key_cfg.overlap_x) / tile_key_cfg.tile_size_x)
                y_key = math.floor((ymin + tile_key_cfg.overlap_y) / tile_key_cfg.tile_size_y)
                tile_key = f"{x_key}_{y_key}"
                if tile_key in tile_key_seen:
                    print(f"Warning: duplicate tile_key {tile_key} for tile {tile_id}")
                tile_key_seen.add(tile_key)

            out_path = os.path.join(cfg.export.heightmap_dir, f"{tile_id}.png")

            warp_opts = gdal.WarpOptions(
                format="MEM",
                outputBounds=(xmin, ymin, xmax, ymax),
                width=cfg.heightmap.out_res,
                height=cfg.heightmap.out_res,
                resampleAlg=cfg.heightmap.resample,
                srcNodata=-9999,
                dstNodata=0,  # Use 0 for nodata (safe since valid elevations scale to 1-65535)
            )
            try:
                # Use empty string for destination to signal MEM driver
                tmp_ds = gdal.Warp("", ds, options=warp_opts)
            except RuntimeError as exc:
                print(f"Warp failed for {tile_id}: {exc}")
                skipped += 1
                continue

            try:
                tile_min, tile_max, valid_count = _compute_ds_minmax(tmp_ds)
            except Exception as exc:
                print(f"Min/max computation failed for {tile_id}: {exc}")
                skipped += 1
                continue

            tile_valid = valid_count > 0
            if not tile_valid:
                print(f"Warning: {tile_id} has no valid samples; using global min/max.")
                tile_min = gmin
                tile_max = gmax

            scale_min = gmin
            scale_max = gmax
            if cfg.heightmap.use_tile_minmax and tile_valid:
                scale_min = tile_min
                scale_max = tile_max
                if scale_max <= scale_min:
                    print(f"Warning: {tile_id} has flat elevation range; using epsilon for scaling.")
                    scale_max = scale_min + 1e-3

            trans_opts = gdal.TranslateOptions(
                outputType=gdal.GDT_UInt16,
                scaleParams=[(scale_min, scale_max, 0, 65535)],
                format="PNG",
                creationOptions=["WORLDFILE=YES"],
            )
            try:
                gdal.Translate(out_path, tmp_ds, options=trans_opts)
            except RuntimeError as exc:
                print(f"Translate failed for {tile_id}: {exc}")
                skipped += 1
                continue
            tmp_ds = None

            f.write(
                f"{tile_id},{xmin},{ymin},{xmax},{ymax},{gmin},{gmax},"
                f"{cfg.heightmap.out_res},{tile_key},{tile_min},{tile_max}\n"
            )
            print(f"Wrote {out_path}")
            written += 1

    print(f"Manifest: {cfg.export.manifest_path}")
    print(f"Summary: wrote {written} tiles; skipped {skipped}.")
    removed = cleanup_aux_files(_cleanup_patterns(cfg.raw.dgm1_dir))
    print(f"Cleanup removed {removed} temporary files/sidecars.")
    return 1 if skipped else 0