#!/usr/bin/env python3

"""
Stack present-day, RCP4.5 and RCP8.5 echinoid rasters into multiband COGs.

Processing:
    - Read matching Fa2020 rasters from cwd/input/
    - Stack present-day, RCP4.5 and RCP8.5 as bands 1-3
    - Preserve original raster grid and EPSG:4326 projection
    - Save Cloud Optimised GeoTIFFs to cwd/output/

Expected input naming:
    Fa2020-Austrocidaris_canaliculata_cog.tif
    Fa2020-Austrocidaris_canaliculata-RCP45_cog.tif
    Fa2020-Austrocidaris_canaliculata-RCP85_cog.tif
    ...
    Fa2020-bioregions_cog.tif
    Fa2020-bioregions-RCP45_cog.tif
    Fa2020-bioregions-RCP85_cog.tif

Species outputs:
    Austrocidaris_canaliculata.tif
    Float32, NoData = -999

Bioregions output:
    Echinoid_bioregions.tif
    Byte, NoData = 255

Bands:
    1 = present_day
    2 = RCP45
    3 = RCP85

Requires:
    GDAL Python bindings (osgeo)

Run:
    python 02_stack_tifs.py
"""

from pathlib import Path
import re
import shutil

from osgeo import gdal

gdal.UseExceptions()


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

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

SPECIES_NODATA = -999.0
BIOREGION_NODATA = 255

SCENARIOS = [
    "present_day",
    "RCP45",
    "RCP85",
]

BAND_NAMES = [
    "present_day",
    "RCP45",
    "RCP85",
]

GRID_TOLERANCE = 1e-9


# =============================================================================
# FIND INPUTS
# =============================================================================

def find_inputs():
    """Find and group the three scenario rasters for each taxon."""

    pattern = re.compile(
        r"^Fa2020-(?P<name>.+?)(?:-(?P<scenario>RCP45|RCP85))?_cog\.tif$",
        re.IGNORECASE,
    )

    groups = {}

    for tif in INPUT_DIR.iterdir():

        if not tif.is_file() or tif.suffix.lower() != ".tif":
            continue

        match = pattern.match(tif.name)

        if not match:
            continue

        name = match.group("name")
        scenario = match.group("scenario")

        if scenario is None:
            scenario = "present_day"
        else:
            scenario = scenario.upper()

        groups.setdefault(name, {})

        if scenario in groups[name]:
            raise RuntimeError(
                f"{name}: duplicate {scenario} raster"
            )

        groups[name][scenario] = tif

    if not groups:
        raise RuntimeError(
            f"No matching Fa2020 TIFFs found in {INPUT_DIR}"
        )

    return groups


# =============================================================================
# VALIDATION
# =============================================================================

def check_trio(name, files):
    """Check that all three scenarios exist and use the same raster grid."""

    missing = [
        scenario
        for scenario in SCENARIOS
        if scenario not in files
    ]

    if missing:
        raise RuntimeError(
            f"{name}: missing {', '.join(missing)}"
        )

    datasets = []

    try:

        for scenario in SCENARIOS:

            ds = gdal.Open(
                str(files[scenario]),
                gdal.GA_ReadOnly,
            )

            if ds is None:
                raise RuntimeError(
                    f"Could not open {files[scenario]}"
                )

            if ds.RasterCount != 1:
                raise RuntimeError(
                    f"{files[scenario].name} has "
                    f"{ds.RasterCount} bands; expected 1"
                )

            datasets.append(ds)

        reference = datasets[0]
        ref_gt = reference.GetGeoTransform()
        ref_srs = reference.GetSpatialRef()

        for ds, scenario in zip(
            datasets[1:],
            SCENARIOS[1:],
        ):

            if (
                ds.RasterXSize != reference.RasterXSize
                or ds.RasterYSize != reference.RasterYSize
            ):
                raise RuntimeError(
                    f"{name}: {scenario} raster dimensions "
                    f"do not match present day"
                )

            gt = ds.GetGeoTransform()

            if not all(
                abs(a - b) <= GRID_TOLERANCE
                for a, b in zip(gt, ref_gt)
            ):
                raise RuntimeError(
                    f"{name}: {scenario} raster grid "
                    f"does not match present day"
                )

            srs = ds.GetSpatialRef()

            if ref_srs is not None and srs is not None:

                if not ref_srs.IsSame(srs):
                    raise RuntimeError(
                        f"{name}: {scenario} CRS "
                        f"does not match present day"
                    )

            elif ds.GetProjection() != reference.GetProjection():

                raise RuntimeError(
                    f"{name}: {scenario} CRS "
                    f"does not match present day"
                )

    finally:
        datasets = []


# =============================================================================
# STACK RASTERS
# =============================================================================

def create_stack(name, files):
    """Stack one taxon's three scenarios and create the final COG."""

    check_trio(
        name,
        files,
    )

    is_bioregions = (
        name.lower() == "bioregions"
    )

    print(f"\nProcessing {name}...")
    print(
        "  Stacking present day, RCP45 and RCP85..."
    )

    ordered_files = [
        files["present_day"],
        files["RCP45"],
        files["RCP85"],
    ]

    vrt_file = (
        TEMP_DIR
        / f"{name}_stack.vrt"
    )

    if is_bioregions:

        output_file = (
            OUTPUT_DIR
            / "Echinoid_bioregions.tif"
        )

        output_type = gdal.GDT_Byte
        vrt_nodata = BIOREGION_NODATA
        final_nodata = BIOREGION_NODATA
        overview_resampling = "NEAREST"

    else:

        output_file = (
            OUTPUT_DIR
            / f"{name}.tif"
        )

        output_type = gdal.GDT_Float32
        vrt_nodata = SPECIES_NODATA
        final_nodata = SPECIES_NODATA
        overview_resampling = "BILINEAR"

    # Source rasters all use -999 as NoData.
    # For bioregions, this is remapped to 255 before Byte conversion.
    vrt = gdal.BuildVRT(
        str(vrt_file),
        [str(path) for path in ordered_files],
        separate=True,
        srcNodata=SPECIES_NODATA,
        VRTNodata=vrt_nodata,
    )

    if vrt is None:
        raise RuntimeError(
            f"Failed to build VRT for {name}"
        )

    for band_number, band_name in enumerate(
        BAND_NAMES,
        start=1,
    ):

        vrt.GetRasterBand(
            band_number
        ).SetDescription(
            band_name
        )

    vrt.FlushCache()
    vrt = None

    print("  Creating COG...")

    if output_file.exists():
        output_file.unlink()

    cog = gdal.Translate(
        str(output_file),
        str(vrt_file),
        format="COG",
        outputType=output_type,
        noData=final_nodata,
        creationOptions=[
            "COMPRESS=DEFLATE",
            "BIGTIFF=IF_SAFER",
            f"RESAMPLING={overview_resampling}",
        ],
    )

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

    # Ensure descriptions are present in the final file.
    for band_number, band_name in enumerate(
        BAND_NAMES,
        start=1,
    ):

        cog.GetRasterBand(
            band_number
        ).SetDescription(
            band_name
        )

    cog.FlushCache()
    cog = None

    print(
        f"  Output: {output_file.name}"
    )


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

def main():

    if not INPUT_DIR.exists():
        raise FileNotFoundError(
            f"Input folder does not exist: {INPUT_DIR}"
        )

    OUTPUT_DIR.mkdir(
        exist_ok=True
    )

    TEMP_DIR.mkdir(
        exist_ok=True
    )

    try:

        groups = find_inputs()

        for name in sorted(
            groups,
            key=str.lower,
        ):

            create_stack(
                name,
                groups[name],
            )

    finally:

        if TEMP_DIR.exists():

            try:
                shutil.rmtree(
                    TEMP_DIR
                )

            except FileNotFoundError:
                # A temporary file may already have disappeared.
                pass

            except PermissionError:
                # Windows/GDAL may still briefly hold a temporary file.
                print(
                    f"\nWarning: could not remove "
                    f"temporary folder: {TEMP_DIR}"
                )

    print("\nFinished.")


if __name__ == "__main__":
    main()