#!/usr/bin/env python3

# Process KRILLPODYM Antarctic krill biomass NetCDF files.
#
# Install:
#   pip install numpy
#
# GDAL Python bindings must be installed and compatible with the
# installed GDAL binaries.

from pathlib import Path
import tempfile

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"

NODATA = -9999
SOURCE_CRS = "EPSG:4326"
TARGET_CRS = "EPSG:3031"

MULTIBAND_OUT = OUTPUT_DIR / "krill_biomass_seasonal_stages.tif"
COMBINED_OUT = OUTPUT_DIR / "krill_postlarval_biomass_combined.tif"

BANDS = [
    ("krill_young_larvae", "summer", "summer1_young_larvae"),
    ("krill_free_larvae", "autumn", "autumn_free_larvae"),
    ("krill_late_larvae", "winter", "winter_late_larvae"),
    ("krill_juveniles", "spring", "spring_juveniles"),
    ("krill_adults", "summer", "summer2_adults"),
]

SEASON_MONTHS = {
    "spring": [9, 10, 11],
    "summer": [12, 1, 2],
    "autumn": [3, 4, 5],
    "winter": [6, 7, 8],
}


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

def month_from_fname(path):
    return int(path.name.split("_")[-1][4:6])


def get_input_files():
    return sorted(
        INPUT_DIR.glob("krill_biomass_*.nc"),
        key=month_from_fname,
    )


def sum_subdatasets(files, stage):
    """Sum a NetCDF variable across the supplied files."""

    datasets = [
        gdal.Open(f'NETCDF:"{path}":{stage}')
        for path in files
    ]

    if any(ds is None for ds in datasets):
        raise RuntimeError(f"Could not open NetCDF variable: {stage}")

    reference = datasets[0]

    arrays = []
    valid_masks = []

    for ds in datasets:

        if (
            ds.RasterXSize != reference.RasterXSize
            or ds.RasterYSize != reference.RasterYSize
            or ds.GetGeoTransform() != reference.GetGeoTransform()
        ):
            raise ValueError(
                f"Input grids do not match for {stage}"
            )

        band = ds.GetRasterBand(1)
        array = band.ReadAsArray().astype(np.float32)

        src_nodata = band.GetNoDataValue()

        valid = np.isfinite(array)

        if src_nodata is not None:
            valid &= array != src_nodata

        arrays.append(array)
        valid_masks.append(valid)

    # Match the previous gdal_calc behaviour:
    # a summed cell is valid only where all contributing rasters are valid.
    valid_all = np.logical_and.reduce(valid_masks)

    result = np.full(
        arrays[0].shape,
        NODATA,
        dtype=np.float32,
    )

    summed = np.zeros(
        arrays[0].shape,
        dtype=np.float32,
    )

    for array in arrays:
        summed += array

    result[valid_all] = summed[valid_all]

    mem = gdal.GetDriverByName("MEM").Create(
        "",
        reference.RasterXSize,
        reference.RasterYSize,
        1,
        gdal.GDT_Float32,
    )

    mem.SetGeoTransform(reference.GetGeoTransform())
    mem.SetProjection(reference.GetProjection())

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

    return mem


def reproject(dataset, output_path):
    """Reproject to EPSG:3031 using bilinear resampling."""

    gdal.Warp(
        str(output_path),
        dataset,
        format="GTiff",
        srcSRS=SOURCE_CRS,
        dstSRS=TARGET_CRS,
        resampleAlg=gdal.GRA_Bilinear,
        srcNodata=NODATA,
        dstNodata=NODATA,
        outputType=gdal.GDT_Float32,
        creationOptions=[
            "TILED=YES",
            "COMPRESS=DEFLATE",
            "BIGTIFF=IF_SAFER",
        ],
    )


def seasonal_sum(stage, season, output_path, input_files):

    season_files = [
        path
        for path in input_files
        if month_from_fname(path) in SEASON_MONTHS[season]
    ]

    if not season_files:
        raise ValueError(
            f"No input files found for {stage} / {season}"
        )

    summed = sum_subdatasets(
        season_files,
        stage,
    )

    reproject(
        summed,
        output_path,
    )


def sum_all_months(stage, input_files, output_path):

    summed = sum_subdatasets(
        input_files,
        stage,
    )

    reproject(
        summed,
        output_path,
    )


def combine_rasters(path_a, path_b, output_path):
    """Add two rasters while preserving NoData behaviour."""

    ds_a = gdal.Open(str(path_a))
    ds_b = gdal.Open(str(path_b))

    if (
        ds_a.RasterXSize != ds_b.RasterXSize
        or ds_a.RasterYSize != ds_b.RasterYSize
        or ds_a.GetGeoTransform() != ds_b.GetGeoTransform()
    ):
        raise ValueError(
            "Juvenile and adult reprojected grids do not match."
        )

    band_a = ds_a.GetRasterBand(1)
    band_b = ds_b.GetRasterBand(1)

    array_a = band_a.ReadAsArray().astype(np.float32)
    array_b = band_b.ReadAsArray().astype(np.float32)

    nodata_a = band_a.GetNoDataValue()
    nodata_b = band_b.GetNoDataValue()

    valid = np.isfinite(array_a) & np.isfinite(array_b)

    if nodata_a is not None:
        valid &= array_a != nodata_a

    if nodata_b is not None:
        valid &= array_b != nodata_b

    result = np.full(
        array_a.shape,
        NODATA,
        dtype=np.float32,
    )

    result[valid] = (
        array_a[valid] + array_b[valid]
    )

    mem = gdal.GetDriverByName("MEM").Create(
        "",
        ds_a.RasterXSize,
        ds_a.RasterYSize,
        1,
        gdal.GDT_Float32,
    )

    mem.SetGeoTransform(ds_a.GetGeoTransform())
    mem.SetProjection(ds_a.GetProjection())

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

    gdal.Translate(
        str(output_path),
        mem,
        format="COG",
        noData=NODATA,
        creationOptions=[
            "COMPRESS=DEFLATE",
            "BIGTIFF=IF_SAFER",
        ],
    )


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

def main():

    OUTPUT_DIR.mkdir(exist_ok=True)

    input_files = get_input_files()

    if not input_files:
        raise FileNotFoundError(
            f"No krill_biomass_*.nc files found in {INPUT_DIR}"
        )

    print(f"Found {len(input_files)} input NetCDF files.")

    with tempfile.TemporaryDirectory() as tmp:

        temp_dir = Path(tmp)

        # --------------------------------------------------
        # Seasonal stage raster
        # --------------------------------------------------

        print("\nProcessing seasonal krill biomass...")

        seasonal_outputs = []

        for stage, season, label in BANDS:

            print(f"  Processing {label}...")

            output = temp_dir / f"{label}.tif"

            seasonal_sum(
                stage,
                season,
                output,
                input_files,
            )

            seasonal_outputs.append(output)

            print(f"  Finished {label}")

        print("\nBuilding seasonal multiband raster...")

        vrt = gdal.BuildVRT(
            str(temp_dir / "krill_stack.vrt"),
            [str(path) for path in seasonal_outputs],
            separate=True,
            srcNodata=NODATA,
            VRTNodata=NODATA,
        )

        for index, (_, _, label) in enumerate(
            BANDS,
            start=1,
        ):
            vrt.GetRasterBand(index).SetDescription(label)

        vrt.FlushCache()

        gdal.Translate(
            str(MULTIBAND_OUT),
            vrt,
            format="COG",
            noData=NODATA,
            creationOptions=[
                "COMPRESS=DEFLATE",
                "BIGTIFF=IF_SAFER",
            ],
        )

        vrt = None

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

        # --------------------------------------------------
        # Combined juvenile + adult biomass
        # --------------------------------------------------

        print("\nProcessing combined postlarval biomass...")

        juvenile_sum = temp_dir / "krill_juveniles_3031.tif"
        adult_sum = temp_dir / "krill_adults_3031.tif"

        print("  Summing juvenile biomass across all months...")

        sum_all_months(
            "krill_juveniles",
            input_files,
            juvenile_sum,
        )

        print("  Summing adult biomass across all months...")

        sum_all_months(
            "krill_adults",
            input_files,
            adult_sum,
        )

        print("  Combining juvenile + adult biomass...")

        combine_rasters(
            juvenile_sum,
            adult_sum,
            COMBINED_OUT,
        )

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

        print("\nCleaning up temporary files...")

    print("\nComplete. Final outputs:")
    print(f"  {MULTIBAND_OUT.name}")
    print(f"  {COMBINED_OUT.name}")


if __name__ == "__main__":
    main()