#!/usr/bin/env python3

# Standardise EnvData rasters for Seamap Antarctica.
#
# Processing:
#   - Read all TIFFs from cwd/input
#   - Remove "envdata_" from the start of each output filename
#   - Assign EPSG:3031 without reprojection/resampling
#   - Convert source NoData / NaN / Inf / fill-like values to -9999
#   - Retain source band descriptions for multiband rasters
#   - Blank the band description for single-band rasters
#   - Output Cloud Optimized GeoTIFFs (COGs) to cwd/output
#
# Run:
#   python <script_name>.py
#
# Install requirements:
#   pip install numpy GDAL


from __future__ import annotations

import tempfile
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"

REMOVE_PREFIX = "envdata_Circumpolar_2kmres_"

OUTPUT_CRS = "EPSG:3031"
OUTPUT_NODATA = -9999.0

# Catch large NetCDF-style fill values such as -3.4e38.
FILL_THRESHOLD = 1e35


# =============================================================================
# GDAL
# =============================================================================

gdal.UseExceptions()


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

def epsg3031_wkt() -> str:
    """Return EPSG:3031 as WKT."""

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

    return srs.ExportToWkt()


def output_name(input_path: Path) -> str:
    """Remove envdata_ from the start of the filename."""

    name = input_path.name

    if name.startswith(REMOVE_PREFIX):
        return name[len(REMOVE_PREFIX):]

    return name


def clean_band(
    band,
) -> np.ndarray:
    """Read one band and convert invalid/fill values to -9999."""

    data = band.ReadAsArray().astype(
        np.float32,
        copy=False,
    )

    invalid = (
        ~np.isfinite(data)
        | (np.abs(data) > FILL_THRESHOLD)
    )

    source_nodata = band.GetNoDataValue()

    if source_nodata is not None:

        if np.isnan(source_nodata):
            invalid |= np.isnan(data)
        else:
            invalid |= data == np.float32(
                source_nodata
            )

    if invalid.any():
        data = data.copy()

        data[invalid] = np.float32(
            OUTPUT_NODATA
        )

    return data


# =============================================================================
# PROCESS ONE RASTER
# =============================================================================

def process_raster(
    input_path: Path,
    output_path: Path,
) -> None:

    print()
    print("=" * 80)
    print(
        f"INPUT:  {input_path.relative_to(BASE_DIR)}"
    )
    print(
        f"OUTPUT: {output_path.relative_to(BASE_DIR)}"
    )

    src = gdal.Open(
        str(input_path),
        gdal.GA_ReadOnly,
    )

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

    width = src.RasterXSize
    height = src.RasterYSize
    band_count = src.RasterCount
    transform = src.GetGeoTransform()

    print(
        f"BANDS:  {band_count}"
    )

    with tempfile.TemporaryDirectory() as temp_dir:

        temp_tif = (
            Path(temp_dir)
            / "standardised.tif"
        )

        driver = gdal.GetDriverByName(
            "GTiff"
        )

        dst = driver.Create(
            str(temp_tif),
            width,
            height,
            band_count,
            gdal.GDT_Float32,
            options=[
                "TILED=YES",
                "COMPRESS=DEFLATE",
                "BIGTIFF=IF_SAFER",
                "INTERLEAVE=BAND",
            ],
        )

        if dst is None:
            raise RuntimeError(
                f"Could not create temporary raster for {input_path.name}"
            )

        # Preserve the existing grid exactly.
        dst.SetGeoTransform(
            transform
        )

        # The supplied CRS definition is EPSG:3031, so define it
        # explicitly without reprojection.
        dst.SetProjection(
            epsg3031_wkt()
        )

        for band_number in range(
            1,
            band_count + 1,
        ):

            src_band = src.GetRasterBand(
                band_number
            )

            dst_band = dst.GetRasterBand(
                band_number
            )

            data = clean_band(
                src_band
            )

            dst_band.WriteArray(
                data
            )

            dst_band.SetNoDataValue(
                OUTPUT_NODATA
            )

            # Retain descriptions for multiband rasters.
            # Single-band outputs deliberately have no description.
            if band_count > 1:
                dst_band.SetDescription(
                    src_band.GetDescription()
                )
            else:
                dst_band.SetDescription("")

        dst.FlushCache()

        dst = None
        src = None

        # Remove an incomplete previous output if present.
        if output_path.exists():
            output_path.unlink()

        # Final Seamap COG.
        cog = gdal.Translate(
            str(output_path),
            str(temp_tif),
            options=gdal.TranslateOptions(
                format="COG",
                creationOptions=[
                    "COMPRESS=DEFLATE",
                    "BIGTIFF=IF_SAFER",
                    "OVERVIEW_RESAMPLING=BILINEAR",
                    "NUM_THREADS=ALL_CPUS",
                ],
            ),
        )

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

        cog = None

    print(
        f"[OK] {output_path.name}"
    )


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

def main() -> None:

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

    input_files = sorted(
        list(INPUT_DIR.glob("*.tif"))
        + list(INPUT_DIR.glob("*.tiff"))
    )

    if not input_files:
        raise FileNotFoundError(
            f"No TIFF files found in:\n{INPUT_DIR}"
        )

    print(
        f"[INPUT]  {INPUT_DIR.relative_to(BASE_DIR)}"
    )
    print(
        f"[OUTPUT] {OUTPUT_DIR.relative_to(BASE_DIR)}"
    )
    print(
        f"[FILES]  {len(input_files)}"
    )
    print(
        f"[CRS]    {OUTPUT_CRS} "
        "(assigned only; no reprojection)"
    )
    print(
        f"[NODATA] {OUTPUT_NODATA:g}"
    )

    for input_path in input_files:

        output_path = (
            OUTPUT_DIR
            / output_name(input_path)
        )

        process_raster(
            input_path,
            output_path,
        )

    print()
    print("=" * 80)
    print(
        f"FINISHED - {len(input_files)} COGs CREATED"
    )
    print("=" * 80)


if __name__ == "__main__":
    main()