#!/usr/bin/env python3

# Extract selected GRACE trend variables from the NetCDF in cwd/input.
#
# Processing:
#   - Extract univariateTrend and multivariateTrend
#   - Retain the native 50 km grid without reprojection or resampling
#   - Assign EPSG:3031 to the existing projected x/y grid
#   - Standardise NoData to -9999
#   - Write Float64 COGs to cwd/output
#
# Includes Windows extended-path handling for long directory paths.
#
# Install:
#   pip install numpy GDAL

from __future__ import annotations

import os
from pathlib import Path

import numpy as np
from osgeo import gdal, osr


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

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

NC_FILENAME = (
    "GRAVIS-3_2002095-2021090_COSTG_0100_AIS_GRID_TUD_0002_GRACEminusSMB_withtrends.nc"
)

VARIABLES = {
    "univariateTrend": "GRACE_2002-2021_univariateTrend.tif",
    "multivariateTrend": "GRACE_2002-2021_multivariateTrend.tif",
}

TARGET_EPSG = 3031
NODATA = -9999.0

COG_OPTIONS = [
    "COMPRESS=DEFLATE",
    "PREDICTOR=3",
    "BIGTIFF=IF_SAFER",
    "NUM_THREADS=ALL_CPUS",
]


# =============================================================================
# HELPERS
# =============================================================================

def longpath(path: Path) -> str:
    """Return a Windows extended-length path where required."""

    path_str = str(path.resolve())

    if os.name == "nt" and not path_str.startswith("\\\\?\\"):
        return "\\\\?\\" + path_str

    return path_str


def path_exists(path: Path) -> bool:
    return os.path.exists(longpath(path))


def epsg_wkt(epsg: int) -> str:
    srs = osr.SpatialReference()
    srs.ImportFromEPSG(epsg)
    return srs.ExportToWkt()


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

def process_variable(
    nc_path: Path,
    variable: str,
    output_path: Path,
) -> None:

    subdataset = f'NETCDF:"{longpath(nc_path)}":{variable}'

    src = gdal.Open(
        subdataset,
        gdal.GA_ReadOnly,
    )

    if src is None:
        raise RuntimeError(
            f"Could not open NetCDF variable: {variable}"
        )

    if src.RasterCount != 1:
        raise RuntimeError(
            f"Expected one band for {variable}, found {src.RasterCount}"
        )

    band = src.GetRasterBand(1)
    array = band.ReadAsArray()

    if array is None:
        raise RuntimeError(
            f"Could not read variable: {variable}"
        )

    array = array.astype(
        np.float32,
        copy=False,
    )

    src_nodata = band.GetNoDataValue()

    valid = np.isfinite(array)

    if src_nodata is not None:
        valid &= array != src_nodata

    output_array = np.full(
        array.shape,
        NODATA,
        dtype=np.float32,
    )

    output_array[valid] = array[valid]

    # Preserve the source grid exactly.
    geotransform = src.GetGeoTransform()

    mem = gdal.GetDriverByName("MEM").Create(
        "",
        src.RasterXSize,
        src.RasterYSize,
        1,
        gdal.GDT_Float32,
    )

    if mem is None:
        raise RuntimeError(
            f"Could not create temporary raster for {variable}"
        )

    mem.SetGeoTransform(
        geotransform
    )

    # Source grid coordinates are already Antarctic Polar Stereographic;
    # assign the missing CRS without reprojection.
    mem.SetProjection(
        epsg_wkt(TARGET_EPSG)
    )

    out_band = mem.GetRasterBand(1)
    out_band.WriteArray(output_array)
    out_band.SetNoDataValue(NODATA)
    out_band.SetDescription(variable)
    out_band.FlushCache()

    if path_exists(output_path):
        os.remove(longpath(output_path))

    output = gdal.Translate(
        longpath(output_path),
        mem,
        options=gdal.TranslateOptions(
            format="COG",
            creationOptions=COG_OPTIONS,
        ),
    )

    if output is None:
        raise RuntimeError(
            f"Failed to create COG: {output_path}"
        )

    output = None
    mem = None
    src = None


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

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

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

    nc_path = INPUT_DIR / NC_FILENAME

    if not path_exists(nc_path):
        raise FileNotFoundError(
            f"NetCDF not found:\n{nc_path}"
        )

    print(
        f"[INPUT] {nc_path.relative_to(BASE_DIR)}"
    )

    for variable, filename in VARIABLES.items():
        output_path = OUTPUT_DIR / filename

        print(
            f"[WRITE] {variable} "
            f"-> {output_path.relative_to(BASE_DIR)}"
        )

        process_variable(
            nc_path,
            variable,
            output_path,
        )

    print("[OK] Done")


if __name__ == "__main__":
    main()