#!/usr/bin/env python3

# Convert selected variables from a curvilinear NetCDF grid to EPSG:3031 COGs.
#
# Install:
#   pip install numpy
#
# GDAL Python bindings must also be installed and compatible with the
# installed GDAL binaries (on Windows, usually via a matching wheel).

from pathlib import Path

import numpy as np
from osgeo import gdal

gdal.UseExceptions()


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

BASE_DIR = Path.cwd()

CONFIG = [
    {
        "nc": "ecoindex_present_day.nc",
        "vars": [
            (
                "ecoindex_hemispheric",
                "AEV_hemispheric_presentday_continuous.tif",
            ),
            (
                "ecoindex_regional",
                "AEV_regional_presentday_continuous.tif",
            ),
        ],
    },
]

# EPSG:3031 output grid.
XMIN, YMIN, XMAX, YMAX = -4500000, -4500000, 4500000, 4500000
RES = 25000  # 25 km

NODATA = -9999

COG_OPTIONS = [
    "COMPRESS=DEFLATE",
    "PREDICTOR=2",
    "BLOCKSIZE=512",
    "BIGTIFF=IF_SAFER",
]


# ==================================================
# FUNCTIONS
# ==================================================

def make_geolocation_files(nc_path: Path):
    """
    Extract TLAT and create a corrected TLONG raster for GDAL geolocation.
    """

    stem = nc_path.stem

    tlat_path = BASE_DIR / f"TLAT_raw_{stem}.tif"
    tlong_path = BASE_DIR / f"TLONG_fixed_{stem}.tif"

    # Extract latitude.
    if not tlat_path.exists():
        gdal.Translate(
            str(tlat_path),
            f'NETCDF:"{nc_path}":TLAT',
            format="GTiff",
            creationOptions=[
                "TILED=YES",
                "BLOCKXSIZE=512",
                "BLOCKYSIZE=512",
                "COMPRESS=DEFLATE",
            ],
        )

    # Extract longitude and convert values below -180 to -180 to 180.
    if not tlong_path.exists():
        src = gdal.Open(f'NETCDF:"{nc_path}":TLONG')

        if src is None:
            raise RuntimeError(f"Could not read TLONG from {nc_path}")

        array = src.ReadAsArray().astype(np.float32)
        array[~np.isfinite(array)] = np.nan
        array[array < -180.0] += 360.0

        out = gdal.GetDriverByName("GTiff").Create(
            str(tlong_path),
            src.RasterXSize,
            src.RasterYSize,
            1,
            gdal.GDT_Float32,
            options=[
                "TILED=YES",
                "BLOCKXSIZE=512",
                "BLOCKYSIZE=512",
                "COMPRESS=DEFLATE",
            ],
        )

        out.SetGeoTransform((0, 1, 0, 0, 0, 1))

        band = out.GetRasterBand(1)
        band.WriteArray(array)
        band.SetNoDataValue(np.nan)

        out = None
        src = None

    return tlat_path, tlong_path


def build_geolocation_vrt(
    nc_path: Path,
    variable: str,
    tlat_path: Path,
    tlong_path: Path,
    vrt_path: Path,
):
    """
    Create a VRT for a NetCDF variable and attach its latitude/longitude
    geolocation arrays.
    """

    gdal.Translate(
        str(vrt_path),
        f'NETCDF:"{nc_path}":{variable}',
        format="VRT",
    )

    ds = gdal.Open(str(vrt_path), gdal.GA_Update)

    ds.SetMetadata(
        {
            "X_DATASET": tlong_path.name,
            "X_BAND": "1",
            "Y_DATASET": tlat_path.name,
            "Y_BAND": "1",
            "PIXEL_OFFSET": "0",
            "PIXEL_STEP": "1",
            "LINE_OFFSET": "0",
            "LINE_STEP": "1",
        },
        "GEOLOCATION",
    )

    ds = None


def warp_to_cog(vrt_path: Path, output_path: Path):
    """
    Reproject the geolocated raster to the fixed EPSG:3031 grid using
    bilinear resampling and save as a COG.
    """

    temp_path = output_path.with_suffix(".tmp.tif")

    gdal.Warp(
        str(temp_path),
        str(vrt_path),
        options=gdal.WarpOptions(
            format="GTiff",
            srcSRS="EPSG:4326",
            dstSRS="EPSG:3031",
            geoloc=True,
            xRes=RES,
            yRes=RES,
            outputBounds=[XMIN, YMIN, XMAX, YMAX],
            resampleAlg="bilinear",
            dstNodata=NODATA,
            creationOptions=[
                "COMPRESS=DEFLATE",
                "PREDICTOR=2",
                "TILED=YES",
                "BLOCKXSIZE=512",
                "BLOCKYSIZE=512",
                "BIGTIFF=IF_SAFER",
            ],
        ),
    )

    # Remove source metadata before creating the final COG.
    ds = gdal.Open(str(temp_path), gdal.GA_Update)
    ds.SetMetadata({})

    band = ds.GetRasterBand(1)
    band.SetMetadata({})
    band.SetNoDataValue(NODATA)

    ds = None

    gdal.Translate(
        str(output_path),
        str(temp_path),
        format="COG",
        creationOptions=COG_OPTIONS,
    )

    temp_path.unlink()


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

def main():

    for config in CONFIG:

        nc_path = BASE_DIR / config["nc"]

        if not nc_path.exists():
            raise FileNotFoundError(nc_path)

        tlat_path, tlong_path = make_geolocation_files(nc_path)

        for variable, output_name in config["vars"]:

            output_path = BASE_DIR / output_name
            vrt_path = BASE_DIR / f"{variable}_geoloc.vrt"

            build_geolocation_vrt(
                nc_path,
                variable,
                tlat_path,
                tlong_path,
                vrt_path,
            )

            warp_to_cog(vrt_path, output_path)

            print(f"Created: {output_path.name}")

    print("Done.")


if __name__ == "__main__":
    main()