#!/usr/bin/env python3

# Sum all species habitat-suitability rasters and rescale the result to 0-1.
#
# species_richness.asc is excluded from the calculation.
#
# Install:
#   pip install numpy
#
# GDAL Python bindings must also be installed and compatible with the
# installed GDAL binaries.

from pathlib import Path

import numpy as np
from osgeo import gdal

gdal.UseExceptions()


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

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

OUTPUT_FILE = OUTPUT_DIR / "allspecies_habitat_suitability.tif"
NODATA = -9999


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

def main():

    species_files = [
        OUTPUT_DIR / f"{path.stem}.tif"
        for path in sorted(INPUT_DIR.glob("*.asc"))
        if path.stem != "species_richness"
    ]

    if not species_files:
        raise FileNotFoundError("No species rasters found.")

    first = gdal.Open(str(species_files[0]))

    summed = np.zeros(
        (first.RasterYSize, first.RasterXSize),
        dtype=np.float32,
    )

    has_data = np.zeros(summed.shape, dtype=bool)

    for path in species_files:
        ds = gdal.Open(str(path))
        array = ds.GetRasterBand(1).ReadAsArray().astype(np.float32)

        valid = np.isfinite(array) & (array != NODATA)

        summed[valid] += array[valid]
        has_data |= valid

        ds = None

    if not np.any(has_data):
        raise RuntimeError("No valid habitat-suitability values found.")

    max_value = float(summed[has_data].max())

    if max_value <= 0:
        raise RuntimeError("Maximum summed habitat-suitability value is not greater than zero.")

    result = np.full(summed.shape, NODATA, dtype=np.float32)
    result[has_data] = summed[has_data] / max_value

    temp_file = OUTPUT_DIR / "_combined_tmp.tif"

    out = gdal.GetDriverByName("GTiff").Create(
        str(temp_file),
        first.RasterXSize,
        first.RasterYSize,
        1,
        gdal.GDT_Float32,
        options=[
            "TILED=YES",
            "COMPRESS=DEFLATE",
            "PREDICTOR=2",
            "BIGTIFF=IF_SAFER",
        ],
    )

    out.SetGeoTransform(first.GetGeoTransform())
    out.SetProjection(first.GetProjection())

    band = out.GetRasterBand(1)
    band.WriteArray(result)
    band.SetNoDataValue(NODATA)

    out = None
    first = None

    gdal.Translate(
        str(OUTPUT_FILE),
        str(temp_file),
        format="COG",
        creationOptions=[
            "COMPRESS=DEFLATE",
            "PREDICTOR=2",
            "BIGTIFF=IF_SAFER",
        ],
    )

    temp_file.unlink()

    print(f"Maximum summed value: {max_value}")
    print(f"Created: {OUTPUT_FILE.name}")


if __name__ == "__main__":
    main()