# Sum band 1 from all TIFFs in the current directory and normalise to 0–1.
#
# Band 1 = habitat suitability
# Band 2 = uncertainty (ignored)
#
# Input rasters are aligned to the grid of the first TIFF before summing.
# NoData (-999) is ignored where at least one input has valid data.
# Cells that are NoData in every input remain -999.
#
# Install requirements:
#   pip install rasterio numpy

from pathlib import Path
import numpy as np
import rasterio
from rasterio.warp import reproject, Resampling


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

BASE_DIR = Path.cwd()

OUTPUT_FILE = BASE_DIR / "krill_habitatSuitability_4spp.tif"

NODATA = -999.0


# =========================
# PROCESSING
# =========================

def main():

    tif_files = sorted(
        f for f in BASE_DIR.glob("*.tif")
        if f != OUTPUT_FILE
    )

    if not tif_files:
        raise FileNotFoundError("No TIFF files found in the current directory.")

    # Use the first TIFF as the target grid
    with rasterio.open(tif_files[0]) as src:
        profile = src.profile.copy()
        shape = (src.height, src.width)
        target_transform = src.transform
        target_crs = src.crs

    summed = np.zeros(shape, dtype=np.float64)
    valid_count = np.zeros(shape, dtype=np.uint8)

    for tif in tif_files:
        print(f"Reading: {tif.name}")

        with rasterio.open(tif) as src:

            data = np.full(shape, NODATA, dtype=np.float32)

            reproject(
                source=rasterio.band(src, 1),
                destination=data,
                src_transform=src.transform,
                src_crs=src.crs,
                src_nodata=NODATA,
                dst_transform=target_transform,
                dst_crs=target_crs,
                dst_nodata=NODATA,
                resampling=Resampling.bilinear
            )

        valid = data != NODATA

        summed[valid] += data[valid]
        valid_count[valid] += 1

    valid = valid_count > 0

    output = np.full(shape, NODATA, dtype=np.float32)

    max_value = summed[valid].max()

    if max_value == 0:
        raise ValueError("Maximum summed value is 0; cannot normalise.")

    output[valid] = summed[valid] / max_value

    profile.update(
        driver="COG",
        dtype="float32",
        count=1,
        nodata=NODATA,
        compress="DEFLATE"
    )

    with rasterio.open(OUTPUT_FILE, "w", **profile) as dst:
        dst.write(output, 1)

    print(f"Created: {OUTPUT_FILE}")


if __name__ == "__main__":
    main()