#!/usr/bin/env python3

"""
Sum diel abundance across all processed species/taxa.

USAGE
-----
1. Per-species files are read from:
       cwd/output/diel_abundance_*.tif

2. Install requirements:
       pip install numpy

   GDAL Python bindings (osgeo) must also be installed.

5. Output:
       cwd/output/diel_abundance_allSpp.tif

   Bands:
       1 = day
       2 = night
"""

from pathlib import Path

import numpy as np
from osgeo import gdal


gdal.UseExceptions()


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

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

OUTPUT_FILE = OUTPUT_DIR / "diel_abundance_allSpp.tif"
TEMP_FILE = OUTPUT_DIR / "_diel_abundance_allSpp_tmp.tif"

NODATA = -9999.0

# =========================


def main():

    # Find individual taxon rasters, excluding the combined output itself
    input_files = sorted(
        f for f in OUTPUT_DIR.glob("diel_abundance_*.tif")
        if f.name != OUTPUT_FILE.name
    )

    if not input_files:
        raise FileNotFoundError(
            f"No diel_abundance_*.tif files found in {OUTPUT_DIR}"
        )

    print(f"Found {len(input_files)} taxon rasters:")
    for f in input_files:
        print(f"  {f.name}")

    # ---------------------------------------------------------
    # Use first raster as the reference grid
    # ---------------------------------------------------------

    ref = gdal.Open(str(input_files[0]))

    if ref.RasterCount != 2:
        raise ValueError(
            f"{input_files[0].name} does not contain exactly 2 bands."
        )

    width = ref.RasterXSize
    height = ref.RasterYSize
    geotransform = ref.GetGeoTransform()
    projection = ref.GetProjection()

    # Float64 during summing, written as Float32 at the end
    day_sum = np.zeros((height, width), dtype=np.float64)
    night_sum = np.zeros((height, width), dtype=np.float64)

    reference_mask = None

    ref = None

    # ---------------------------------------------------------
    # Sum all taxa
    # ---------------------------------------------------------

    for i, path in enumerate(input_files, start=1):

        print(f"Adding {i}/{len(input_files)}: {path.name}")

        ds = gdal.Open(str(path))

        # Check all rasters use exactly the same grid
        if ds.RasterXSize != width or ds.RasterYSize != height:
            raise ValueError(f"Grid size differs: {path.name}")

        if ds.GetGeoTransform() != geotransform:
            raise ValueError(f"Geotransform differs: {path.name}")

        if ds.GetProjection() != projection:
            raise ValueError(f"Projection differs: {path.name}")

        if ds.RasterCount != 2:
            raise ValueError(f"Expected 2 bands: {path.name}")

        day = ds.GetRasterBand(1).ReadAsArray().astype(np.float64)
        night = ds.GetRasterBand(2).ReadAsArray().astype(np.float64)

        # Both bands should have the same valid-data footprint
        valid = (day != NODATA) & (night != NODATA)

        # All processed taxa should have the same prediction footprint.
        # Fail rather than silently summing only a subset of taxa.
        if reference_mask is None:
            reference_mask = valid.copy()
        elif not np.array_equal(valid, reference_mask):
            raise ValueError(
                f"NoData footprint differs for {path.name}. "
                "Stopped to avoid creating a partial-species sum."
            )

        day_sum[valid] += day[valid]
        night_sum[valid] += night[valid]

        ds = None

    # Set cells outside prediction area to NoData
    day_sum[~reference_mask] = NODATA
    night_sum[~reference_mask] = NODATA

    # ---------------------------------------------------------
    # Write temporary 2-band GeoTIFF
    # ---------------------------------------------------------

    print("Writing combined raster...")

    driver = gdal.GetDriverByName("GTiff")

    out = driver.Create(
        str(TEMP_FILE),
        width,
        height,
        2,
        gdal.GDT_Float32,
        options=[
            "TILED=YES",
            "COMPRESS=DEFLATE",
            "PREDICTOR=3",
            "BIGTIFF=IF_SAFER",
        ],
    )

    out.SetGeoTransform(geotransform)
    out.SetProjection(projection)

    band = out.GetRasterBand(1)
    band.WriteArray(day_sum.astype(np.float32))
    band.SetNoDataValue(NODATA)
    band.SetDescription("day")

    band = out.GetRasterBand(2)
    band.WriteArray(night_sum.astype(np.float32))
    band.SetNoDataValue(NODATA)
    band.SetDescription("night")

    out.FlushCache()
    out = None

    # ---------------------------------------------------------
    # Convert to COG
    # ---------------------------------------------------------

    print("Creating COG...")

    cog = gdal.Translate(
        str(OUTPUT_FILE),
        str(TEMP_FILE),
        format="COG",
        outputType=gdal.GDT_Float32,
        noData=NODATA,
        creationOptions=[
            "COMPRESS=DEFLATE",
            "PREDICTOR=3",
            "BIGTIFF=IF_SAFER",
        ],
    )

    if cog is None:
        raise RuntimeError("Failed to create combined COG.")

    cog = None

    TEMP_FILE.unlink(missing_ok=True)

    print(f"Output: {OUTPUT_FILE}")


if __name__ == "__main__":
    main()