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

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

Output:
    ./output/ice_surface_elevation_1-10k.tif

Band descriptions:
    warming_1k
    warming_2k
    ...
    warming_10k

Requires:
    - numpy
    - GDAL with Python bindings

    GDAL Python bindings must also be available.
"""

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_surface_elevation_1-10K.tif"

EXPECTED_WARMINGS = list(range(1, 11))  # 1K to 10K
NODATA = -9999.0
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 matching equi_1.0k.nc ... equi_10.0k.nc (case-insensitive K).
    Returns dict like {1: Path(...), 2: Path(...), ...}
    """
    pattern = re.compile(r"^equi_(\d+)k\.nc$", re.IGNORECASE)
    found = {}

    for path in input_dir.iterdir():
        if not path.is_file():
            continue

        m = pattern.match(path.name)
        if not m:
            continue

        warming_value = float(m.group(1))
        if not warming_value.is_integer():
            continue

        warming_int = int(warming_value)
        if warming_int not in EXPECTED_WARMINGS:
            continue

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

        found[warming_int] = path

    missing = [k for k in EXPECTED_WARMINGS if k not in found]
    if missing:
        fail(f"Missing expected input files for warming levels: {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}' in {nc_path.name}")
    return ds


def get_target_projection_wkt():
    """
    EPSG:3031 WKT.
    """
    srs = osr.SpatialReference()
    srs.ImportFromEPSG(3031)
    return srs.ExportToWkt()


def read_masked_usurf(nc_path: Path):
    """
    Read usurf and mask from one PISM file.
    Keep only mask values 2 and 3; set all other cells to NODATA.
    Returns:
        out_arr (float32)
        geotransform
        projection
        xsize
        ysize
    """
    usurf_ds = open_subdataset(nc_path, "usurf")
    mask_ds = open_subdataset(nc_path, "mask")

    xsize = usurf_ds.RasterXSize
    ysize = usurf_ds.RasterYSize

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

    usurf_band = usurf_ds.GetRasterBand(1)
    mask_band = mask_ds.GetRasterBand(1)

    usurf = usurf_band.ReadAsArray().astype(np.float32)
    mask = mask_band.ReadAsArray()

    usurf_nodata = usurf_band.GetNoDataValue()

    valid = np.isin(mask, list(KEEP_MASK_VALUES))

    if usurf_nodata is not None:
        valid &= ~np.isclose(usurf, usurf_nodata)

    valid &= np.isfinite(usurf)

    out_arr = np.full((ysize, xsize), NODATA, dtype=np.float32)
    out_arr[valid] = usurf[valid]

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

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

    usurf_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 multiband GeoTIFF before converting to COG.
    """
    first_arr, gt, proj, xsize, ysize = read_masked_usurf(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: {temp_tif}")

    ds.SetGeoTransform(gt)
    ds.SetProjection(get_target_projection_wkt())  # enforce EPSG:3031

    for band_index, warming in enumerate(EXPECTED_WARMINGS, start=1):
        print(f"Processing {warming}K: {warming_files[warming].name}")

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

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

        band = ds.GetRasterBand(band_index)
        band.WriteArray(arr)
        band.SetNoDataValue(NODATA)
        band.SetDescription(f"warming_{warming}k")
        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: {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: {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_surface_elevation.tif"

        build_temp_gtiff(temp_tif, warming_files)
        convert_to_cog(temp_tif, output_path)

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


if __name__ == "__main__":
    main()