#!/usr/bin/env python3

# Create global mean habitat suitability and species richness rasters.
#
# Input:
#   ./input/*.tif
#
# Output:
#   ./output/global/HIMI_allspp_global_mean.tif
#   ./output/global/HIMI_allspp_global_richness.tif
#
# Global mean bands:
#   1 = 2003-2005
#   2 = 2014-2016
#   3 = change
#
# Species richness bands:
#   1 = 2003-2005
#   2 = 2014-2016
#   3 = change
#
# NoData values are excluded from calculations.
# Output NoData = -9999.
#
# GDAL Python bindings must be installed and compatible with the
# installed GDAL binaries.

import math
from array import array
from pathlib import Path

from osgeo import gdal

gdal.UseExceptions()


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

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

OUTPUT_MEAN = OUTPUT_DIR / "HIMI_allspp_global_mean.tif"
OUTPUT_RICHNESS = OUTPUT_DIR / "HIMI_allspp_global_richness.tif"

NODATA = -9999.0

MEAN_COLLECTIONS = [
    ("2003_2005", "2003-2005"),
    ("2014_2016", "2014-2016"),
    ("Difference", "change"),
]

RICHNESS_PERIODS = [
    ("2003_2005", "2003-2005"),
    ("2014_2016", "2014-2016"),
]

RICHNESS_CHANGE_DESCRIPTION = "change"


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

def find_collection_files(prefix):
    """Find all input TIFFs belonging to a collection."""

    files = [
        path
        for path in INPUT_DIR.glob("*.tif")
        if path.stem.startswith(f"{prefix}_")
    ]

    return sorted(files, key=lambda path: path.name.lower())


def is_valid(value, nodata):
    """Check whether an input value is valid."""

    if math.isnan(value):
        return False

    if nodata is None:
        return True

    if isinstance(nodata, float) and math.isnan(nodata):
        return True

    return value != nodata


def is_valid_output(value):
    """Check whether an existing output value is valid."""

    if math.isnan(value):
        return False

    return value != NODATA


def read_float32_row(band, xsize, y):
    """Read one raster row as Float32."""

    data = band.ReadRaster(
        xoff=0,
        yoff=y,
        xsize=xsize,
        ysize=1,
        buf_xsize=xsize,
        buf_ysize=1,
        buf_type=gdal.GDT_Float32,
    )

    row = array("f")
    row.frombytes(data)

    return row


def write_float32_row(band, row, xsize, y):
    """Write one Float32 raster row."""

    band.WriteRaster(
        xoff=0,
        yoff=y,
        xsize=xsize,
        ysize=1,
        buf_string=row.tobytes(),
        buf_xsize=xsize,
        buf_ysize=1,
        buf_type=gdal.GDT_Float32,
    )


def validate_match(ds, ref_info, path):
    """Check that an input raster matches the reference grid."""

    if (
        ds.RasterXSize != ref_info["xsize"]
        or ds.RasterYSize != ref_info["ysize"]
    ):
        raise ValueError(f"Raster size mismatch: {path}")

    if ds.GetGeoTransform() != ref_info["gt"]:
        raise ValueError(f"Geotransform mismatch: {path}")

    if ds.GetProjection() != ref_info["proj"]:
        raise ValueError(f"Projection mismatch: {path}")


def get_reference_info(template_path):
    """Get grid information from the reference raster."""

    ds = gdal.Open(str(template_path))

    if ds is None:
        raise RuntimeError(f"Could not open template raster: {template_path}")

    info = {
        "xsize": ds.RasterXSize,
        "ysize": ds.RasterYSize,
        "gt": ds.GetGeoTransform(),
        "proj": ds.GetProjection(),
    }

    ds = None

    return info


def create_output_raster(output_path, ref_info, band_count):
    """Create an empty Float32 output raster."""

    if output_path.exists():
        output_path.unlink()

    output = gdal.GetDriverByName("GTiff").Create(
        str(output_path),
        ref_info["xsize"],
        ref_info["ysize"],
        band_count,
        gdal.GDT_Float32,
        options=[
            "TILED=YES",
            "COMPRESS=DEFLATE",
            "BIGTIFF=IF_SAFER",
        ],
    )

    if output is None:
        raise RuntimeError(f"Could not create output raster: {output_path}")

    output.SetGeoTransform(ref_info["gt"])
    output.SetProjection(ref_info["proj"])

    return output


def build_stat_band(
    output_band,
    collection_files,
    band_description,
    ref_info,
    statistic,
):
    """
    Calculate either:
        mean = mean predicted occurrence across valid species
        sum  = summed predicted occurrence across valid species
    """

    if statistic not in ("mean", "sum"):
        raise ValueError("statistic must be either 'mean' or 'sum'")

    datasets = []

    try:
        for path in collection_files:

            ds = gdal.Open(str(path))

            if ds is None:
                raise RuntimeError(f"Could not open raster: {path}")

            validate_match(ds, ref_info, path)

            datasets.append(ds)

        xsize = ref_info["xsize"]
        ysize = ref_info["ysize"]

        for y in range(ysize):

            sums = [0.0] * xsize
            counts = [0] * xsize

            for ds in datasets:

                band = ds.GetRasterBand(1)
                src_nodata = band.GetNoDataValue()
                row = read_float32_row(band, xsize, y)

                for i, value in enumerate(row):

                    if is_valid(value, src_nodata):
                        sums[i] += value
                        counts[i] += 1

            output_row = array("f")

            for i in range(xsize):

                if counts[i] > 0:

                    if statistic == "mean":
                        output_row.append(sums[i] / counts[i])

                    else:
                        output_row.append(sums[i])

                else:
                    output_row.append(NODATA)

            write_float32_row(
                output_band,
                output_row,
                xsize,
                y,
            )

        output_band.SetDescription(band_description)
        output_band.SetNoDataValue(NODATA)
        output_band.FlushCache()

    finally:
        datasets = []


def build_change_band(
    output_ds,
    band1_index,
    band2_index,
    change_band_index,
    band_description,
    ref_info,
):
    """Calculate change as band 2 minus band 1."""

    xsize = ref_info["xsize"]
    ysize = ref_info["ysize"]

    band1 = output_ds.GetRasterBand(band1_index)
    band2 = output_ds.GetRasterBand(band2_index)
    change_band = output_ds.GetRasterBand(change_band_index)

    for y in range(ysize):

        row1 = read_float32_row(band1, xsize, y)
        row2 = read_float32_row(band2, xsize, y)

        output_row = array("f")

        for i in range(xsize):

            if (
                is_valid_output(row1[i])
                and is_valid_output(row2[i])
            ):
                output_row.append(row2[i] - row1[i])

            else:
                output_row.append(NODATA)

        write_float32_row(
            change_band,
            output_row,
            xsize,
            y,
        )

    change_band.SetDescription(band_description)
    change_band.SetNoDataValue(NODATA)
    change_band.FlushCache()


def get_collection_files(collections):
    """Build file lists for each requested collection."""

    result = []

    for prefix, description in collections:

        files = find_collection_files(prefix)

        if not files:
            raise RuntimeError(
                f"No TIFFs found for collection: {prefix}"
            )

        result.append(
            (prefix, description, files)
        )

    return result


def create_global_mean(ref_info):
    """Create global mean habitat-suitability raster."""

    collections = get_collection_files(MEAN_COLLECTIONS)

    output = create_output_raster(
        OUTPUT_MEAN,
        ref_info,
        len(MEAN_COLLECTIONS),
    )

    try:
        for band_index, (
            prefix,
            description,
            files,
        ) in enumerate(collections, start=1):

            build_stat_band(
                output.GetRasterBand(band_index),
                files,
                description,
                ref_info,
                statistic="mean",
            )

    finally:
        output.FlushCache()
        output = None

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


def create_species_richness(ref_info):
    """Create species richness and richness-change raster."""

    periods = get_collection_files(RICHNESS_PERIODS)

    output = create_output_raster(
        OUTPUT_RICHNESS,
        ref_info,
        3,
    )

    try:
        # Band 1: 2003-2005 richness
        build_stat_band(
            output.GetRasterBand(1),
            periods[0][2],
            periods[0][1],
            ref_info,
            statistic="sum",
        )

        # Band 2: 2014-2016 richness
        build_stat_band(
            output.GetRasterBand(2),
            periods[1][2],
            periods[1][1],
            ref_info,
            statistic="sum",
        )

        # Band 3: change = 2014-2016 minus 2003-2005
        build_change_band(
            output,
            band1_index=1,
            band2_index=2,
            change_band_index=3,
            band_description=RICHNESS_CHANGE_DESCRIPTION,
            ref_info=ref_info,
        )

    finally:
        output.FlushCache()
        output = None

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


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

def main():

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

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

    # Use the first 2003-2005 raster as the reference grid.
    template_files = find_collection_files("2003_2005")

    if not template_files:
        raise RuntimeError(
            "No 2003_2005 TIFFs found to use as template."
        )

    ref_info = get_reference_info(template_files[0])

    create_global_mean(ref_info)
    create_species_richness(ref_info)

    print("Done.")


if __name__ == "__main__":
    main()