#!/usr/bin/env python3

# Convert GlobSed-v3 NetCDF sediment thickness to a 5 km EPSG:3031 COG
# covering the Southern Ocean south of 50°S.
#
# Requires:
#   pip install GDAL
#
# Processing:
#   - Extract NetCDF variable "z"
#   - Assign EPSG:4326
#   - Crop to 50°S–90°S
#   - Reproject to EPSG:3031 at 5 km resolution

from pathlib import Path
from osgeo import gdal

gdal.UseExceptions()


# =========================
# CONFIG
# =========================

BASE_DIR = Path(__file__).resolve().parent

INPUT_FILE = BASE_DIR / "input" / "GlobSed-v3.nc"
OUTPUT_DIR = BASE_DIR / "output"
OUTPUT_FILE = OUTPUT_DIR / "GlobSed-v3_50S.tif"

NETCDF_VARIABLE = "z"

SRC_CRS = "EPSG:4326"
DST_CRS = "EPSG:3031"

# Crop bounds in EPSG:4326:
# upper-left x, upper-left y, lower-right x, lower-right y
CROP_BOUNDS = (-180.0, -50.0, 180.0, -90.0)

TARGET_RESOLUTION = 5000
NODATA = -9999.0

RESAMPLING = "bilinear"


# =========================
# MAIN
# =========================

def main():

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

    if not INPUT_FILE.exists():
        raise FileNotFoundError(f"Input file not found: {INPUT_FILE}")

    source = f'NETCDF:"{INPUT_FILE}":{NETCDF_VARIABLE}'

    temp_4326 = OUTPUT_DIR / "_temp_crop_4326.tif"
    temp_3031 = OUTPUT_DIR / "_temp_3031.tif"

    try:
        # -----------------------------------------------------
        # 1. Extract NetCDF variable and crop south of 50°S
        # -----------------------------------------------------

        print("1. Extracting and cropping NetCDF variable...")

        gdal.Translate(
            str(temp_4326),
            source,
            format="GTiff",
            outputType=gdal.GDT_Float32,
            outputSRS=SRC_CRS,
            projWin=CROP_BOUNDS,
            projWinSRS=SRC_CRS,
            noData=NODATA,
            creationOptions=[
                "TILED=YES",
                "COMPRESS=DEFLATE",
                "PREDICTOR=3",
                "BIGTIFF=IF_SAFER",
            ],
        )

        # -----------------------------------------------------
        # 2. Reproject to EPSG:3031 at 5 km resolution
        # -----------------------------------------------------

        print("2. Reprojecting to EPSG:3031...")

        gdal.Warp(
            str(temp_3031),
            str(temp_4326),
            format="GTiff",
            srcSRS=SRC_CRS,
            dstSRS=DST_CRS,
            xRes=TARGET_RESOLUTION,
            yRes=TARGET_RESOLUTION,
            dstNodata=NODATA,
            outputType=gdal.GDT_Float32,
            resampleAlg=RESAMPLING,
            multithread=True,
            creationOptions=[
                "TILED=YES",
                "COMPRESS=DEFLATE",
                "PREDICTOR=3",
                "BIGTIFF=IF_SAFER",
            ],
        )

        # -----------------------------------------------------
        # 3. Convert to COG
        # -----------------------------------------------------

        print("3. Creating COG...")

        gdal.Translate(
            str(OUTPUT_FILE),
            str(temp_3031),
            format="COG",
            outputType=gdal.GDT_Float32,
            noData=NODATA,
            creationOptions=[
                "COMPRESS=DEFLATE",
                "PREDICTOR=3",
                "OVERVIEWS=AUTO",
                f"RESAMPLING={RESAMPLING}",
                "BIGTIFF=IF_SAFER",
            ],
        )

        print(f"Created: {OUTPUT_FILE}")

    finally:
        temp_4326.unlink(missing_ok=True)
        temp_3031.unlink(missing_ok=True)


if __name__ == "__main__":
    main()