#!/usr/bin/env python3

# Process Antarctic krill layers and write final EPSG:3031 COGs.
#
# Install:
#   pip install numpy
#
# GDAL Python bindings must also be installed and compatible with
# the installed GDAL binaries.

from pathlib import Path

import numpy as np
from osgeo import gdal, ogr


gdal.UseExceptions()


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

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

LARVAL_INPUT = INPUT_DIR / "larval_antarctickrill.shp"
ADULT_INPUT = INPUT_DIR / "adult_antarctickrill.tif"
ICE_INPUT = INPUT_DIR / "ice_krill.tif"

LARVAL_OUTPUT = OUTPUT_DIR / "larval_antarctickrill.tif"
ADULT_OUTPUT = OUTPUT_DIR / "adult_antarctickrill.tif"
ICE_OUTPUT = OUTPUT_DIR / "ice_krill.tif"

TARGET_CRS = "EPSG:3031"

LARVAL_NODATA = -3.4e38
ADULT_NODATA = -3.4e38
ICE_NODATA = -3.4e38

LARVAL_RESOLUTION = 1000


# ==================================================
# LARVAL ANTARCTIC KRILL
# ==================================================

def process_larval():

    print("Processing larval Antarctic krill...")

    # Reproject vector to EPSG:3031 in memory
    vector = gdal.VectorTranslate(
        "",
        str(LARVAL_INPUT),
        format="Memory",
        dstSRS=TARGET_CRS,
    )

    layer = vector.GetLayer()
    xmin, xmax, ymin, ymax = layer.GetExtent()

    xsize = int(np.ceil((xmax - xmin) / LARVAL_RESOLUTION))
    ysize = int(np.ceil((ymax - ymin) / LARVAL_RESOLUTION))

    raster = gdal.GetDriverByName("MEM").Create(
        "",
        xsize,
        ysize,
        1,
        gdal.GDT_Float32,
    )

    raster.SetProjection(TARGET_CRS)
    raster.SetGeoTransform(
        (
            xmin,
            LARVAL_RESOLUTION,
            0,
            ymax,
            0,
            -LARVAL_RESOLUTION,
        )
    )

    band = raster.GetRasterBand(1)
    band.SetNoDataValue(LARVAL_NODATA)
    band.Fill(LARVAL_NODATA)

    gdal.RasterizeLayer(
        raster,
        [1],
        layer,
        options=["ATTRIBUTE=abundance"],
    )

    gdal.Translate(
        str(LARVAL_OUTPUT),
        raster,
        format="COG",
        creationOptions=[
            "COMPRESS=DEFLATE",
            "BIGTIFF=IF_SAFER",
        ],
    )

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


# ==================================================
# ADULT ANTARCTIC KRILL
# ==================================================

def process_adult():

    print("Processing adult Antarctic krill...")

    src = gdal.Open(str(ADULT_INPUT))
    band = src.GetRasterBand(1)

    values = band.ReadAsArray().astype(np.float32)

    source_nodata = band.GetNoDataValue()

    valid = np.isfinite(values)

    if source_nodata is not None:
        valid &= values != source_nodata

    if not np.any(valid):
        raise ValueError("Adult krill raster contains no valid values.")

    max_value = float(values[valid].max())

    print(f"  Maximum value: {max_value}")

    # Preserve original processing:
    # zero values become NoData, all other valid values are divided by max.
    result = np.full(
        values.shape,
        ADULT_NODATA,
        dtype=np.float32,
    )

    rescale_mask = valid & (values != 0)

    result[rescale_mask] = (
        values[rescale_mask] / max_value
    )

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

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

    out_band = mem.GetRasterBand(1)
    out_band.WriteArray(result)
    out_band.SetNoDataValue(ADULT_NODATA)

    gdal.Warp(
        str(ADULT_OUTPUT),
        mem,
        format="COG",
        dstSRS=TARGET_CRS,
        resampleAlg=gdal.GRA_Bilinear,
        srcNodata=ADULT_NODATA,
        dstNodata=ADULT_NODATA,
        outputType=gdal.GDT_Float32,
        creationOptions=[
            "COMPRESS=DEFLATE",
            "PREDICTOR=2",
            "BIGTIFF=IF_SAFER",
        ],
    )

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


# ==================================================
# ICE KRILL
# ==================================================

def process_ice():

    print("Processing ice krill...")

    src = gdal.Open(str(ICE_INPUT))
    band = src.GetRasterBand(1)

    values = band.ReadAsArray().astype(np.float32)

    source_nodata = band.GetNoDataValue()

    valid = np.isfinite(values)

    if source_nodata is not None:
        valid &= values != source_nodata

    result = np.full(
        values.shape,
        ICE_NODATA,
        dtype=np.float32,
    )

    # Convert log10(individuals/m³) to individuals/m³
    result[valid] = 10 ** values[valid]

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

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

    out_band = mem.GetRasterBand(1)
    out_band.WriteArray(result)
    out_band.SetNoDataValue(ICE_NODATA)

    gdal.Warp(
        str(ICE_OUTPUT),
        mem,
        format="COG",
        dstSRS=TARGET_CRS,
        resampleAlg=gdal.GRA_Bilinear,
        srcNodata=ICE_NODATA,
        dstNodata=ICE_NODATA,
        outputType=gdal.GDT_Float32,
        creationOptions=[
            "COMPRESS=DEFLATE",
            "PREDICTOR=2",
            "BIGTIFF=IF_SAFER",
        ],
    )

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


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

def main():

    OUTPUT_DIR.mkdir(exist_ok=True)

    process_larval()
    process_adult()
    process_ice()

    print("Done.")


if __name__ == "__main__":
    main()