#!/usr/bin/env python3

# Stack the three HIMI rasters for each species into a multiband GeoTIFF.
#
# Input:
#   ./input/*.tif
#
# Expected filename patterns:
#   2003_2005_<species>.tif
#   2014_2016_<species>.tif
#   Difference_<species>.tif
#
# Output:
#   ./output/HIMI_<species>.tif
#
# Bands:
#   1 = 2003-2005
#   2 = 2014-2016
#   3 = change
#
# NaN values are converted to -9999.
#
# GDAL Python bindings must be installed and compatible with the
# installed GDAL binaries.

from pathlib import Path
import shutil

from osgeo import gdal

gdal.UseExceptions()


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

BASE_DIR = Path.cwd()
RAW_DIR = BASE_DIR / "input"
OUT_DIR = BASE_DIR / "output"
TMP_DIR = OUT_DIR / "_tmp_cleaned"

OUTPUT_PREFIX = "HIMI"
NODATA = -9999

COLLECTIONS = {
    "2003_2005": "2003-2005",
    "2014_2016": "2014-2016",
    "Difference": "change",
}

COLLECTION_ORDER = [
    "2003_2005",
    "2014_2016",
    "Difference",
]


# ==================================================
# FUNCTIONS
# ==================================================

def parse_input_filename(stem):
    """Return collection and species from an input filename."""

    for collection in COLLECTION_ORDER:
        prefix = f"{collection}_"

        if stem.startswith(prefix):
            species = stem[len(prefix):]
            return collection, species

    return None, None


def build_species_lookup():
    """Group input rasters by species and collection."""

    species_lookup = {}

    for path in RAW_DIR.glob("*.tif"):
        collection, species = parse_input_filename(path.stem)

        if collection is None:
            continue

        species_lookup.setdefault(species, {})[collection] = path

    return species_lookup


def clean_raster(src_path, dst_path):
    """Convert NaN NoData values to -9999."""

    result = gdal.Warp(
        str(dst_path),
        str(src_path),
        format="GTiff",
        srcNodata="nan",
        dstNodata=NODATA,
        outputType=gdal.GDT_Float32,
        creationOptions=[
            "TILED=YES",
            "COMPRESS=DEFLATE",
            "BIGTIFF=IF_SAFER",
        ],
    )

    if result is None:
        raise RuntimeError(f"Failed to clean raster: {src_path}")

    result = None


def create_species_stack(species, files):

    missing = [
        collection
        for collection in COLLECTION_ORDER
        if collection not in files
    ]

    if missing:
        print(f"Skipping {species} - missing: {', '.join(missing)}")
        return

    cleaned_files = []

    for band_index, collection in enumerate(COLLECTION_ORDER, start=1):

        temp_path = TMP_DIR / f"{band_index:03d}_{species}_{collection}.tif"

        clean_raster(
            files[collection],
            temp_path,
        )

        cleaned_files.append(
            (COLLECTIONS[collection], temp_path)
        )

    vrt_path = TMP_DIR / f"{species}_stack.vrt"
    output_path = OUT_DIR / f"{OUTPUT_PREFIX}_{species}.tif"

    vrt = gdal.BuildVRT(
        str(vrt_path),
        [str(path) for description, path in cleaned_files],
        separate=True,
    )

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

    vrt = None

    output = gdal.Translate(
        str(output_path),
        str(vrt_path),
        format="GTiff",
        outputType=gdal.GDT_Float32,
        noData=NODATA,
        creationOptions=[
            "TILED=YES",
            "COMPRESS=DEFLATE",
            "BIGTIFF=IF_SAFER",
        ],
    )

    if output is None:
        raise RuntimeError(f"Failed to write output: {output_path}")

    for band_index, (description, path) in enumerate(cleaned_files, start=1):
        band = output.GetRasterBand(band_index)
        band.SetDescription(description)
        band.SetNoDataValue(NODATA)

    output = None

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


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

def main():

    if not RAW_DIR.exists():
        raise FileNotFoundError(f"Could not find raw folder: {RAW_DIR}")

    OUT_DIR.mkdir(exist_ok=True)

    if TMP_DIR.exists():
        shutil.rmtree(TMP_DIR)

    TMP_DIR.mkdir()

    try:
        species_lookup = build_species_lookup()

        if not species_lookup:
            raise RuntimeError("No matching TIFFs found in raw folder.")

        for species in sorted(species_lookup, key=str.lower):
            create_species_stack(
                species,
                species_lookup[species],
            )

    finally:
        if TMP_DIR.exists():
            shutil.rmtree(TMP_DIR)

    print("Done.")


if __name__ == "__main__":
    main()