# Reproject TIFFs to Antarctic Polar Stereographic (EPSG:3031).
#
# Band 1 = probability of presence       -> bilinear resampling
# Band 2 = standard deviation            -> bilinear resampling
# Band 3 = presence/absence (0 or 1)     -> nearest-neighbour resampling
#
# Install requirements:
#   pip install rasterio numpy

from pathlib import Path

import numpy as np
import rasterio
from rasterio.crs import CRS
from rasterio.enums import Resampling
from rasterio.shutil import copy as rio_copy
from rasterio.warp import calculate_default_transform, reproject


# =============================================================================
# CONFIG
# =============================================================================

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

TARGET_CRS = CRS.from_epsg(3031)
TARGET_RESOLUTION = 10000  # metres
NODATA = -999.0


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

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

tif_files = sorted(INPUT_DIR.glob("*.tif"))

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


for input_path in tif_files:

    output_path = OUTPUT_DIR / input_path.name
    temp_path = OUTPUT_DIR / f"{input_path.stem}_3031_temp.tif"

    print(f"\nProcessing: {input_path.name}")

    with rasterio.open(input_path) as src:

        if src.crs is None:
            raise ValueError(
                f"{input_path.name} has no readable CRS definition."
            )

        if src.count != 3:
            raise ValueError(
                f"{input_path.name} has {src.count} bands; expected 3."
            )

        transform, width, height = calculate_default_transform(
            src.crs,
            TARGET_CRS,
            src.width,
            src.height,
            *src.bounds,
            resolution=TARGET_RESOLUTION,
        )

        profile = src.profile.copy()
        profile.update(
            driver="GTiff",
            crs=TARGET_CRS,
            transform=transform,
            width=width,
            height=height,
            dtype="float32",
            nodata=NODATA,
            tiled=True,
            compress="deflate",
            predictor=3,
            BIGTIFF="IF_SAFER",
        )

        with rasterio.open(temp_path, "w", **profile) as dst:

            for band in range(1, src.count + 1):

                if band in (1, 2):
                    resampling = Resampling.bilinear
                else:
                    resampling = Resampling.nearest

                reproject(
                    source=rasterio.band(src, band),
                    destination=rasterio.band(dst, band),
                    src_transform=src.transform,
                    src_crs=src.crs,
                    src_nodata=src.nodata,
                    dst_transform=transform,
                    dst_crs=TARGET_CRS,
                    dst_nodata=NODATA,
                    resampling=resampling,
                )

                # Preserve band descriptions
                if src.descriptions[band - 1]:
                    dst.set_band_description(
                        band,
                        src.descriptions[band - 1],
                    )

    # Convert temporary GeoTIFF to COG
    rio_copy(
        temp_path,
        output_path,
        driver="COG",
        compress="DEFLATE",
        blocksize=512,
        overview_resampling="nearest",
    )

    temp_path.unlink()

    print(f"Created: {output_path.name}")


print("\nDone.")