#!/usr/bin/env python3
"""
Create a 10-band COG of Antarctic ice thickness from PISM equilibrium files.

Expected input files in ./input:
    equi_1k.nc
    ...
    equi_10k.nc

Also accepts names:
    equi_1.0k.nc
    ...

Output:
    ./output/ice_sheet_thickness_1-10C.tif

Band descriptions:
    warming_1C
    warming_2C
    ...
    warming_10C

Non-ice cells are set to -9999 using the PISM mask:
    2 = grounded ice
    3 = floating ice

Requires:
    - numpy
    - GDAL with Python bindings
"""

from pathlib import Path
import re
import sys
import tempfile

import numpy as np
from osgeo import gdal, osr


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

OUTPUT_NAME = "ice_sheet_thickness_1-10C.tif"

EXPECTED_WARMINGS = list(range(1, 11))

NODATA = -9999.0

# PISM mask values to retain
KEEP_MASK_VALUES = {2, 3}  # 2 = grounded ice, 3 = floating ice

# COG options
COG_COMPRESS = "DEFLATE"
COG_BLOCKSIZE = 512
COG_RESAMPLING = "BILINEAR"


def fail(msg: str):
    print(f"ERROR: {msg}", file=sys.stderr)
    sys.exit(1)


def find_input_files(input_dir: Path):
    """
    Find files such as:

        equi_1k.nc
        equi_1.0k.nc
        equi_10k.nc
        equi_10.0k.nc

    Matching is case-insensitive.

    Returns:
        {1: Path(...), 2: Path(...), ..., 10: Path(...)}
    """

    pattern = re.compile(
        r"^equi_(\d+(?:\.\d+)?)k\.nc$",
        re.IGNORECASE
    )

    found = {}

    for path in input_dir.iterdir():

        if not path.is_file():
            continue

        match = pattern.match(path.name)

        if not match:
            continue

        warming_value = float(match.group(1))

        # Only accept whole-number warming levels
        if not warming_value.is_integer():
            continue

        warming = int(warming_value)

        if warming not in EXPECTED_WARMINGS:
            continue

        if warming in found:
            fail(
                f"Duplicate file for {warming}K: "
                f"{found[warming].name} and {path.name}"
            )

        found[warming] = path

    missing = [
        warming
        for warming in EXPECTED_WARMINGS
        if warming not in found
    ]

    if missing:
        fail(
            "Missing expected input files for warming levels: "
            f"{missing}"
        )

    return found


def open_subdataset(nc_path: Path, var_name: str):
    """
    Open a NetCDF variable as a GDAL subdataset.
    """

    subdataset = f'NETCDF:"{nc_path}":{var_name}'

    ds = gdal.Open(subdataset)

    if ds is None:
        fail(
            f"Could not open variable '{var_name}' "
            f"in {nc_path.name}"
        )

    return ds


def get_target_projection_wkt():
    """
    Return EPSG:3031 WKT.
    """

    srs = osr.SpatialReference()
    srs.ImportFromEPSG(3031)

    return srs.ExportToWkt()


def read_masked_thickness(nc_path: Path):
    """
    Read ice thickness (thk) and mask from one PISM file.

    Keep only:
        mask = 2  grounded ice
        mask = 3  floating ice

    All other cells are set to -9999.

    Returns:
        out_arr
        geotransform
        projection
        xsize
        ysize
    """

    thk_ds = open_subdataset(nc_path, "thk")
    mask_ds = open_subdataset(nc_path, "mask")

    xsize = thk_ds.RasterXSize
    ysize = thk_ds.RasterYSize

    if (
        mask_ds.RasterXSize != xsize
        or mask_ds.RasterYSize != ysize
    ):
        fail(
            f"thk/mask size mismatch in {nc_path.name}"
        )

    thk_band = thk_ds.GetRasterBand(1)
    mask_band = mask_ds.GetRasterBand(1)

    thk = thk_band.ReadAsArray().astype(np.float32)
    mask = mask_band.ReadAsArray()

    thk_nodata = thk_band.GetNoDataValue()

    # Keep only grounded or floating ice
    valid = np.isin(mask, list(KEEP_MASK_VALUES))

    # Remove source NoData if defined
    if thk_nodata is not None:
        valid &= ~np.isclose(thk, thk_nodata)

    # Remove NaN / infinity
    valid &= np.isfinite(thk)

    out_arr = np.full(
        (ysize, xsize),
        NODATA,
        dtype=np.float32
    )

    out_arr[valid] = thk[valid]

    geotransform = thk_ds.GetGeoTransform()
    projection = thk_ds.GetProjection()

    # Force EPSG:3031 if source projection is absent
    if not projection:
        projection = get_target_projection_wkt()

    thk_ds = None
    mask_ds = None

    return (
        out_arr,
        geotransform,
        projection,
        xsize,
        ysize,
    )


def build_temp_gtiff(temp_tif: Path, warming_files: dict):
    """
    Build an intermediate 10-band GeoTIFF before converting to COG.
    """

    first_arr, gt, proj, xsize, ysize = read_masked_thickness(
        warming_files[1]
    )

    gtiff_driver = gdal.GetDriverByName("GTiff")

    if gtiff_driver is None:
        fail("GTiff driver is not available in GDAL.")

    ds = gtiff_driver.Create(
        str(temp_tif),
        xsize,
        ysize,
        len(EXPECTED_WARMINGS),
        gdal.GDT_Float32,
        options=[
            "TILED=YES",
            "COMPRESS=DEFLATE",
            "BIGTIFF=IF_SAFER",
        ],
    )

    if ds is None:
        fail(
            f"Could not create temporary GeoTIFF: "
            f"{temp_tif}"
        )

    ds.SetGeoTransform(gt)
    ds.SetProjection(get_target_projection_wkt())

    for band_index, warming in enumerate(
        EXPECTED_WARMINGS,
        start=1
    ):

        print(
            f"Processing {warming}K: "
            f"{warming_files[warming].name}"
        )

        arr, gt2, proj2, x2, y2 = read_masked_thickness(
            warming_files[warming]
        )

        if x2 != xsize or y2 != ysize:
            fail(
                f"Raster size mismatch in "
                f"{warming_files[warming].name}"
            )

        band = ds.GetRasterBand(band_index)

        band.WriteArray(arr)
        band.SetNoDataValue(NODATA)

        # Use Celsius notation for the GeoServer/public display
        band.SetDescription(f"warming_{warming}C")

        band.FlushCache()

    ds.FlushCache()
    ds = None


def convert_to_cog(temp_tif: Path, output_tif: Path):
    """
    Convert the intermediate GeoTIFF to a Cloud Optimized GeoTIFF.
    """

    src_ds = gdal.Open(str(temp_tif))

    if src_ds is None:
        fail(
            f"Could not reopen temporary GeoTIFF: "
            f"{temp_tif}"
        )

    cog_driver = gdal.GetDriverByName("COG")

    if cog_driver is None:
        fail("COG driver is not available in GDAL.")

    cog_ds = cog_driver.CreateCopy(
        str(output_tif),
        src_ds,
        options=[
            f"COMPRESS={COG_COMPRESS}",
            f"BLOCKSIZE={COG_BLOCKSIZE}",
            f"RESAMPLING={COG_RESAMPLING}",
            "BIGTIFF=IF_SAFER",
            "OVERVIEWS=AUTO",
        ],
    )

    if cog_ds is None:
        fail(
            f"Could not create COG: {output_tif}"
        )

    cog_ds = None
    src_ds = None


def main():

    gdal.UseExceptions()

    if not INPUT_DIR.exists():
        fail(
            f"Input folder not found: "
            f"{INPUT_DIR}"
        )

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

    warming_files = find_input_files(INPUT_DIR)

    output_path = OUTPUT_DIR / OUTPUT_NAME

    with tempfile.TemporaryDirectory(
        dir=OUTPUT_DIR
    ) as tmpdir:

        temp_tif = (
            Path(tmpdir)
            / "temp_ice_thickness.tif"
        )

        build_temp_gtiff(
            temp_tif,
            warming_files
        )

        convert_to_cog(
            temp_tif,
            output_path
        )

    print(f"\nDone: {output_path}")


if __name__ == "__main__":
    main()