#!/usr/bin/env python3

# Reproject downloaded NSIDC ice shelf imagery to EPSG:3031.
#
# Processing:
#   - Find TIFFs recursively in cwd/imagery/<ice_shelf>/orig/
#   - Reproject to EPSG:3031 at 250 m using bilinear resampling
#   - Align output pixels to the 250 m grid
#   - Convert to Int16 with NoData -9999
#   - Build average-resampled overviews
#   - Write COGs to cwd/imagery/<ice_shelf>/int16_3031/
#
# Existing non-empty outputs are skipped.
#
# Install:
#   pip install GDAL

from pathlib import Path

from osgeo import gdal


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

BASE_DIR = Path.cwd()
IMAGERY_DIR = BASE_DIR / "imagery"

INPUT_FOLDER_NAME = "orig"
OUTPUT_FOLDER_NAME = "int16_3031"

TARGET_CRS = "EPSG:3031"
PIXEL_SIZE = 250.0
NODATA = -9999

WARP_RESAMPLING = "bilinear"
OVERVIEW_RESAMPLING = "average"
OVERVIEW_LEVELS = [2, 4, 8, 16, 32, 64]

COG_OPTIONS = [
    "COMPRESS=DEFLATE",
    "PREDICTOR=2",
    "BLOCKSIZE=512",
    "BIGTIFF=IF_SAFER",
    "NUM_THREADS=ALL_CPUS",
]


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

def process_tif(src: Path) -> None:
    shelf_dir = src.parent.parent
    output_dir = shelf_dir / OUTPUT_FOLDER_NAME
    output_dir.mkdir(parents=True, exist_ok=True)

    output_file = output_dir / src.name
    temp_file = output_dir / f"{src.stem}.tmp.tif"

    if output_file.exists() and output_file.stat().st_size > 0:
        print(
            f"[SKIP] "
            f"{output_file.relative_to(BASE_DIR)}"
        )
        return

    if temp_file.exists():
        temp_file.unlink()

    print(
        f"[INPUT]  "
        f"{src.relative_to(BASE_DIR)}"
    )
    print(
        f"[OUTPUT] "
        f"{output_file.relative_to(BASE_DIR)}"
    )

    # Reproject, resample and convert to Int16.
    warped = gdal.Warp(
        str(temp_file),
        str(src),
        options=gdal.WarpOptions(
            format="GTiff",
            dstSRS=TARGET_CRS,
            xRes=PIXEL_SIZE,
            yRes=PIXEL_SIZE,
            targetAlignedPixels=True,
            resampleAlg=WARP_RESAMPLING,
            dstNodata=NODATA,
            outputType=gdal.GDT_Int16,
            creationOptions=[
                "TILED=YES",
                "COMPRESS=DEFLATE",
                "PREDICTOR=2",
                "BIGTIFF=IF_SAFER",
                "BLOCKXSIZE=512",
                "BLOCKYSIZE=512",
            ],
        ),
    )

    if warped is None:
        raise RuntimeError(
            f"Failed to reproject: {src}"
        )

    # Build the same overview pyramid as the original script.
    gdal.SetConfigOption(
        "COMPRESS_OVERVIEW",
        "DEFLATE",
    )
    gdal.SetConfigOption(
        "PREDICTOR_OVERVIEW",
        "2",
    )
    gdal.SetConfigOption(
        "BIGTIFF_OVERVIEW",
        "IF_SAFER",
    )

    result = warped.BuildOverviews(
        OVERVIEW_RESAMPLING,
        OVERVIEW_LEVELS,
    )

    if result != 0:
        warped = None
        raise RuntimeError(
            f"Failed to build overviews: {src}"
        )

    warped = None

    # Convert temporary GeoTIFF to final COG.
    cog = gdal.Translate(
        str(output_file),
        str(temp_file),
        options=gdal.TranslateOptions(
            format="COG",
            creationOptions=COG_OPTIONS,
        ),
    )

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

    cog = None

    temp_file.unlink(missing_ok=True)

    print(
        f"[OK] "
        f"{output_file.relative_to(BASE_DIR)}"
    )


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

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

    tif_files = sorted(
        path
        for path in IMAGERY_DIR.rglob("*.tif")
        if path.parent.name == INPUT_FOLDER_NAME
    )

    if not tif_files:
        raise FileNotFoundError(
            f"No TIFFs found in "
            f"imagery/**/{INPUT_FOLDER_NAME}/"
        )

    print(
        f"Found {len(tif_files)} TIFF(s)"
    )

    for src in tif_files:
        try:
            process_tif(src)

        except Exception as error:
            print(
                f"[FAIL] "
                f"{src.relative_to(BASE_DIR)}: "
                f"{error}"
            )

    print()
    print("FINISHED")


if __name__ == "__main__":
    main()