#!/usr/bin/env python3

# Extract selected ALBMAP variables from NetCDF in cwd/input.
#
# Processing:
#   - Extract firn thickness, snow accumulation and surface air temperature
#   - Build georeferencing from the NetCDF x1/y1 coordinate arrays
#   - Orient arrays north-up where required
#   - Assign EPSG:3031 without reprojection or resampling
#   - Standardise NoData to -9999
#   - Write Float32 COGs to cwd/output
#
# Install:
#   pip install numpy netCDF4 GDAL

from pathlib import Path

import numpy as np
from netCDF4 import Dataset
from osgeo import gdal, osr


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

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

VARIABLES = {
    "firn": "ALBMAP_firn_correction.tif",
    "accr": "ALBMAP_snow_accumulation.tif",
    "temp": "ALBMAP_surface_air_temperature.tif",
}

NODATA = -9999.0

COG_OPTIONS = [
    "COMPRESS=DEFLATE",
    "PREDICTOR=2",
    "BIGTIFF=IF_SAFER",
    "RESAMPLING=NEAREST",
]


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

def find_single_nc() -> Path:
    files = sorted(INPUT_DIR.glob("*.nc"))

    if not files:
        raise FileNotFoundError(f"No NetCDF found in: {INPUT_DIR}")

    if len(files) > 1:
        raise RuntimeError(
            f"Expected one NetCDF in {INPUT_DIR}, found {len(files)}:\n"
            + "\n".join(p.name for p in files)
        )

    return files[0]


def epsg3031_wkt() -> str:
    srs = osr.SpatialReference()
    srs.ImportFromEPSG(3031)
    return srs.ExportToWkt()


def orient_array(
    array: np.ndarray,
    x: np.ndarray,
    y: np.ndarray,
) -> tuple[np.ndarray, tuple]:

    if x.ndim != 1 or y.ndim != 1:
        raise RuntimeError("Expected x1 and y1 to be 1D coordinate arrays")

    if len(x) < 2 or len(y) < 2:
        raise RuntimeError("x1 and y1 must each contain at least two values")

    if array.shape != (len(y), len(x)):
        raise RuntimeError(
            f"Array shape {array.shape} does not match "
            f"coordinates (y={len(y)}, x={len(x)})"
        )

    x = np.asarray(x, dtype=np.float64)
    y = np.asarray(y, dtype=np.float64)

    xres = float(abs(x[1] - x[0]))
    yres = float(abs(y[1] - y[0]))

    if x[0] > x[-1]:
        x = x[::-1]
        array = np.fliplr(array)

    if y[0] < y[-1]:
        y = y[::-1]
        array = np.flipud(array)

    geotransform = (
        float(x[0] - xres / 2),
        xres,
        0.0,
        float(y[0] + yres / 2),
        0.0,
        -yres,
    )

    return array, geotransform


def write_cog(
    array: np.ndarray,
    output_path: Path,
    geotransform: tuple,
    projection: str,
) -> None:

    mem = gdal.GetDriverByName("MEM").Create(
        "",
        array.shape[1],
        array.shape[0],
        1,
        gdal.GDT_Float32,
    )

    mem.SetGeoTransform(geotransform)
    mem.SetProjection(projection)

    band = mem.GetRasterBand(1)
    band.WriteArray(array)
    band.SetNoDataValue(NODATA)
    band.FlushCache()

    if output_path.exists():
        output_path.unlink()

    output = gdal.GetDriverByName("COG").CreateCopy(
        str(output_path),
        mem,
        options=COG_OPTIONS,
    )

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

    output = None
    mem = None


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

def main() -> None:
    gdal.UseExceptions()
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

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

    with Dataset(nc_path, "r") as ds:
        if "x1" not in ds.variables or "y1" not in ds.variables:
            raise RuntimeError("NetCDF does not contain x1 and y1 coordinate variables")

        x = np.asarray(ds.variables["x1"][:], dtype=np.float64)
        y = np.asarray(ds.variables["y1"][:], dtype=np.float64)
        projection = epsg3031_wkt()

        for variable_name, output_name in VARIABLES.items():
            if variable_name not in ds.variables:
                raise RuntimeError(
                    f"Variable '{variable_name}' not found in NetCDF"
                )

            var = ds.variables[variable_name]
            data = var[:]

            if np.ma.isMaskedArray(data):
                array = np.ma.filled(data, NODATA).astype(np.float32)
            else:
                array = np.asarray(data, dtype=np.float32)

            src_nodata = getattr(
                var,
                "_FillValue",
                getattr(var, "missing_value", None),
            )

            array, geotransform = orient_array(array, x, y)

            valid = np.isfinite(array) & (array != NODATA)

            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]

            output_path = OUTPUT_DIR / output_name

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

            write_cog(
                output_array,
                output_path,
                geotransform,
                projection,
            )

    print("[OK] Done")


if __name__ == "__main__":
    main()