#!/usr/bin/env python3

# Rasterise daily polynya polygons to a frequency raster.
#
# Input:
#   cwd/output/polynyas_daily.shp
#
# Output:
#   cwd/output/polynya_dailyfreq_1979-2022.tif
#
# Each polygon contributes 1 to intersected cells, so overlapping daily
# polygons accumulate to give observation frequency.
#
# Install:
#   pip install GDAL numpy

from pathlib import Path
import math

import numpy as np
from osgeo import gdal, ogr


# =============================================================================
# CONFIGURATION
# =============================================================================

BASE_DIR = Path.cwd()
OUTPUT_DIR = BASE_DIR / "output"

INPUT_SHP = OUTPUT_DIR / "polynyas_daily.shp"
OUTPUT_COG = OUTPUT_DIR / "polynya_dailyfreq_1979-2022.tif"

PIXEL_SIZE = 25000.0
ALL_TOUCHED = False
NODATA = -9999

COG_OPTIONS = [
    "COMPRESS=DEFLATE",
    "BLOCKSIZE=512",
    "PREDICTOR=2",
    "BIGTIFF=IF_SAFER",
    "NUM_THREADS=ALL_CPUS",
    "SPARSE_OK=YES",
]


# =============================================================================
# PROCESSING
# =============================================================================

def main() -> None:
    gdal.UseExceptions()

    if not INPUT_SHP.exists():
        raise FileNotFoundError(f"Input shapefile not found:\n{INPUT_SHP}")

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    ds_vec = ogr.Open(str(INPUT_SHP))
    if ds_vec is None:
        raise RuntimeError(f"Could not open: {INPUT_SHP}")

    layer = ds_vec.GetLayer()
    if layer is None:
        raise RuntimeError("Could not read input shapefile layer")

    minx, maxx, miny, maxy = layer.GetExtent()

    x_size = math.ceil((maxx - minx) / PIXEL_SIZE)
    y_size = math.ceil((maxy - miny) / PIXEL_SIZE)

    if x_size <= 0 or y_size <= 0:
        raise RuntimeError("Invalid raster dimensions")

    print(f"[INPUT]  {INPUT_SHP.relative_to(BASE_DIR)}")
    print(f"[OUTPUT] {OUTPUT_COG.relative_to(BASE_DIR)}")
    print(f"[GRID]   {x_size} x {y_size} cells at {PIXEL_SIZE:g} m")

    # Create an in-memory Int32 raster aligned to the vector extent.
    mem_driver = gdal.GetDriverByName("MEM")
    raster = mem_driver.Create("", x_size, y_size, 1, gdal.GDT_Int32)

    raster.SetGeoTransform(
        (minx, PIXEL_SIZE, 0.0, maxy, 0.0, -PIXEL_SIZE)
    )

    srs = layer.GetSpatialRef()
    if srs is not None:
        raster.SetProjection(srs.ExportToWkt())
    else:
        print("[WARNING] Input shapefile has no CRS")

    band = raster.GetRasterBand(1)
    band.SetNoDataValue(NODATA)
    band.Fill(0)

    options = ["MERGE_ALG=ADD"]
    if ALL_TOUCHED:
        options.append("ALL_TOUCHED=TRUE")

    print("[RASTERIZE] Counting polygon overlaps")

    gdal.RasterizeLayer(
        raster,
        [1],
        layer,
        burn_values=[1],
        options=options,
    )

    # Cells not intersected by any polygon become NoData.
    array = band.ReadAsArray()
    array[array == 0] = NODATA
    band.WriteArray(array.astype(np.int32, copy=False))
    band.FlushCache()

    ds_vec = None

    if OUTPUT_COG.exists():
        OUTPUT_COG.unlink()

    print("[COG] Creating output")

    result = gdal.Translate(
        str(OUTPUT_COG),
        raster,
        options=gdal.TranslateOptions(
            format="COG",
            creationOptions=COG_OPTIONS,
        ),
    )

    if result is None:
        raise RuntimeError("Failed to create output COG")

    result = None
    raster = None

    print(f"[OK] {OUTPUT_COG.relative_to(BASE_DIR)}")


if __name__ == "__main__":
    main()