#!/usr/bin/env python3

"""
Create 2-band diel abundance GeoTIFFs of mesopelagic fish NetCDF files.

USAGE
-----
1. File all NetCDF files in:
       cwd/input/

   Expected filenames:
       pred_gaus_ba_day.nc
       pred_gaus_ba_night.nc
       pred_gaus_ea_day.nc
       pred_gaus_ea_night.nc
       ...

2. GDAL Python bindings (osgeo) must be installed.

3. Outputs are written to:
       cwd/output/

   Example:
       diel_abundance_Bathylagus_spp.tif
       diel_abundance_Electrona_antarctica.tif

   Bands:
       1 = day
       2 = night
"""

from pathlib import Path
import shutil

from osgeo import gdal


gdal.UseExceptions()


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

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

SRC_NODATA = -3.4e38
DST_NODATA = -9999.0

SRC_SRS = "EPSG:4326"
DST_SRS = "EPSG:3031"

# Filename code -> output taxon name
SPECIES = {
    "ba": "Bathylagus_spp",
    "ea": "Electrona_antarctica",
    "ec": "Electrona_carlsbergi",
    "gb": "Gymnoscopelus_braueri",
    "gf": "Gymnoscopelus_fraseri",
    "gn": "Gymnoscopelus_nicholsi",
    "ka": "Krefftichthys_anderssoni",
    "pb": "Protomyctophum_bolini",
    "pt": "Protomyctophum_tenisoni",
}

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


def process_species(code, species_name):

    day_file = INPUT_DIR / f"pred_gaus_{code}_day.nc"
    night_file = INPUT_DIR / f"pred_gaus_{code}_night.nc"

    if not day_file.exists():
        raise FileNotFoundError(f"Missing: {day_file}")

    if not night_file.exists():
        raise FileNotFoundError(f"Missing: {night_file}")

    print(f"\nProcessing {species_name}...")

    stack_vrt = TEMP_DIR / f"{code}_stack.vrt"
    warped_tif = TEMP_DIR / f"{code}_3031.tif"

    output_tif = OUTPUT_DIR / f"diel_abundance_{species_name}.tif"

    # ------------------------------------------------------------------
    # 1. Stack original day and night NetCDF rasters
    # ------------------------------------------------------------------

    print("  Stacking day and night...")

    vrt = gdal.BuildVRT(
        str(stack_vrt),
        [str(day_file), str(night_file)],
        separate=True,
        srcNodata=SRC_NODATA,
        VRTNodata=SRC_NODATA,
    )

    if vrt is None:
        raise RuntimeError(f"Failed to build VRT for {species_name}")

    vrt.GetRasterBand(1).SetDescription("day")
    vrt.GetRasterBand(2).SetDescription("night")
    vrt.FlushCache()
    vrt = None

    # ------------------------------------------------------------------
    # 2. Reproject to EPSG:3031
    # ------------------------------------------------------------------

    print("  Reprojecting to EPSG:3031...")

    warped = gdal.Warp(
        str(warped_tif),
        str(stack_vrt),
        format="GTiff",
        srcSRS=SRC_SRS,
        dstSRS=DST_SRS,
        srcNodata=SRC_NODATA,
        dstNodata=DST_NODATA,
        resampleAlg=gdal.GRA_Bilinear,
        multithread=True,
        warpOptions=[
            "NUM_THREADS=ALL_CPUS",
            "UNIFIED_SRC_NODATA=YES",
        ],
        outputType=gdal.GDT_Float32,
        creationOptions=[
            "TILED=YES",
            "COMPRESS=DEFLATE",
            "PREDICTOR=3",
            "BIGTIFF=IF_SAFER",
        ],
    )

    if warped is None:
        raise RuntimeError(f"Failed to reproject {species_name}")

    warped.GetRasterBand(1).SetDescription("day")
    warped.GetRasterBand(2).SetDescription("night")

    warped.FlushCache()
    warped = None

    # ------------------------------------------------------------------
    # 3. Convert to COG
    # ------------------------------------------------------------------

    print("  Creating COG...")

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

    if cog is None:
        raise RuntimeError(f"Failed to create COG: {output_tif}")

    cog = None

    print(f"  Output: {output_tif.name}")


def main():

    OUTPUT_DIR.mkdir(exist_ok=True)
    TEMP_DIR.mkdir(exist_ok=True)

    try:
        for code, species_name in SPECIES.items():
            process_species(code, species_name)

    finally:
        # Temporary folder may occasionally be locked by OneDrive/Windows.
        if TEMP_DIR.exists():
            try:
                shutil.rmtree(TEMP_DIR)
            except PermissionError:
                print(f"\nWarning: could not remove temporary folder: {TEMP_DIR}")

    print("\nFinished.")


if __name__ == "__main__":
    main()