# Create present-day and projected Areas of Ecological Significance (AES)
# layers for RAATD marine predators and combine them into a multiband COG.
#
# Present-day AES:
#   The colony-weighted habitat importance raster contains continuous values.
#   The upper 10% of valid values are classified as AES.
#
# RCP layers:
#   The RCP4.5 and RCP8.5 rasters contain the number of climate models
#   (0-8) predicting AES-like conditions.
#
# All three source rasters are in EPSG:4326.
#
# The present-day continuous raster is reprojected to EPSG:3031 using
# bilinear resampling. The RCP model-count rasters are reprojected onto
# exactly the same grid using nearest neighbour resampling.
#
# A one-pixel cyclic overlap is added to all three source rasters before
# reprojection to prevent a seam at the +/-180 degree longitude boundary.
#
# Input:
#   cwd/input/Hi2023-aes_colony_weighted_cog.tif
#   cwd/input/Hi2023-aes_colony_weighted-RCP45_cog.tif
#   cwd/input/Hi2023-aes_colony_weighted-RCP85_cog.tif
#
# Output:
#   cwd/RAATD_marine_predators_AES_GCM_predictions.tif
#
# Output bands:
#   1. present_AES
#          0 = present non-AES
#          1 = present AES
#
#   2. RCP45_AES_modelCount
#   3. RCP85_AES_modelCount
#          0-8 = number of climate models predicting AES-like conditions
#
#   4. RCP45_gainloss
#   5. RCP85_gainloss
#
#       For present non-AES cells:
#           0-8 = number of models predicting gain of AES-like conditions
#
#       For present AES cells:
#           10-18 = 10 + number of models predicting retention of AES
#
#       Therefore:
#           0  = non-AES; 0/8 predict gain
#           8  = non-AES; 8/8 predict gain
#           10 = AES;     0/8 predict retention (8/8 loss)
#           18 = AES;     8/8 predict retention (0/8 loss)
#
# All output bands use NoData = -999.
# Final output datatype is Int16.
# The COG contains no internal overviews.
#
# Install requirements:
#   pip install numpy
#
# GDAL Python bindings must also be available.

from pathlib import Path

import numpy as np
from osgeo import gdal, osr


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

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

PRESENT_FILE = INPUT_DIR / "Hi2023-aes_colony_weighted_cog.tif"
RCP45_FILE = INPUT_DIR / "Hi2023-aes_colony_weighted-RCP45_cog.tif"
RCP85_FILE = INPUT_DIR / "Hi2023-aes_colony_weighted-RCP85_cog.tif"

OUTPUT_FILE = BASE_DIR / "RAATD_marine_predators_AES_GCM_predictions.tif"

SOURCE_EPSG = "EPSG:4326"
TARGET_EPSG = "EPSG:3031"

OUTPUT_NODATA = -999
AES_PERCENTILE = 90


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

def make_cyclic_source(src_path: Path):
    """
    Add one column from the opposite side of the raster to each
    longitude boundary to prevent a seam at +/-180 degrees.
    """

    src = gdal.Open(str(src_path))

    if src is None:
        raise RuntimeError(f"Could not open: {src_path}")

    band = src.GetRasterBand(1)
    array = band.ReadAsArray()
    src_nodata = band.GetNoDataValue()

    # Add one cyclic column to each side.
    array_wrapped = np.concatenate(
        [
            array[:, -1:],
            array,
            array[:, :1],
        ],
        axis=1,
    )

    gt = src.GetGeoTransform()

    # Shift the origin one pixel west for the added column.
    wrapped_gt = (
        gt[0] - gt[1],
        gt[1],
        gt[2],
        gt[3],
        gt[4],
        gt[5],
    )

    wrapped = gdal.GetDriverByName("MEM").Create(
        "",
        src.RasterXSize + 2,
        src.RasterYSize,
        1,
        band.DataType,
    )

    wrapped.SetGeoTransform(wrapped_gt)

    srs = osr.SpatialReference()
    srs.SetFromUserInput(SOURCE_EPSG)
    wrapped.SetProjection(srs.ExportToWkt())

    wrapped_band = wrapped.GetRasterBand(1)
    wrapped_band.WriteArray(array_wrapped)

    if src_nodata is not None:
        wrapped_band.SetNoDataValue(src_nodata)

    src = None

    return wrapped


def warp_present():
    """
    Reproject the continuous present-day habitat importance raster
    to EPSG:3031 using bilinear resampling.
    """

    wrapped = make_cyclic_source(PRESENT_FILE)
    src_nodata = wrapped.GetRasterBand(1).GetNoDataValue()

    warped = gdal.Warp(
        "",
        wrapped,
        options=gdal.WarpOptions(
            srcSRS=SOURCE_EPSG,
            dstSRS=TARGET_EPSG,
            resampleAlg=gdal.GRA_Bilinear,
            srcNodata=src_nodata,
            dstNodata=OUTPUT_NODATA,
            outputType=gdal.GDT_Float32,
            format="MEM",
        ),
    )

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

    return warped


def warp_rcp_to_target(
    src_path: Path,
    target_projection: str,
    target_gt: tuple,
    target_width: int,
    target_height: int,
):
    """
    Reproject an RCP model-count raster onto the exact present-day
    EPSG:3031 grid using nearest neighbour resampling.
    """

    wrapped = make_cyclic_source(src_path)
    src_nodata = wrapped.GetRasterBand(1).GetNoDataValue()

    warped = gdal.Warp(
        "",
        wrapped,
        options=gdal.WarpOptions(
            srcSRS=SOURCE_EPSG,
            dstSRS=target_projection,
            resampleAlg=gdal.GRA_NearestNeighbour,
            srcNodata=src_nodata,
            dstNodata=OUTPUT_NODATA,
            outputBounds=(
                target_gt[0],
                target_gt[3] + target_height * target_gt[5],
                target_gt[0] + target_width * target_gt[1],
                target_gt[3],
            ),
            width=target_width,
            height=target_height,
            outputType=gdal.GDT_Int16,
            format="MEM",
        ),
    )

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

    return (
        warped
        .GetRasterBand(1)
        .ReadAsArray()
        .astype(np.int16)
    )


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

def main():

    gdal.UseExceptions()

    for path in [PRESENT_FILE, RCP45_FILE, RCP85_FILE]:
        if not path.exists():
            raise FileNotFoundError(path)

    # --------------------------------------------------
    # Reproject present-day habitat importance
    # --------------------------------------------------

    print(
        "Reprojecting present-day habitat importance "
        "using bilinear resampling..."
    )

    present_ds = warp_present()

    target_gt = present_ds.GetGeoTransform()
    target_projection = present_ds.GetProjection()
    target_width = present_ds.RasterXSize
    target_height = present_ds.RasterYSize

    present = (
        present_ds
        .GetRasterBand(1)
        .ReadAsArray()
        .astype(np.float32)
    )

    print(f"Target grid: {target_width} x {target_height}")

    # --------------------------------------------------
    # Derive present-day AES
    # --------------------------------------------------

    valid_present = (
        np.isfinite(present)
        & (present != OUTPUT_NODATA)
    )

    aes_threshold = np.percentile(
        present[valid_present],
        AES_PERCENTILE,
    )

    print(
        f"Present-day AES threshold "
        f"({AES_PERCENTILE}th percentile): "
        f"{aes_threshold}"
    )

    present_aes = np.full(
        present.shape,
        OUTPUT_NODATA,
        dtype=np.int16,
    )

    present_aes[valid_present] = 0

    present_aes[
        valid_present & (present >= aes_threshold)
    ] = 1

    present_ds = None

    # --------------------------------------------------
    # Reproject RCP rasters
    # --------------------------------------------------

    print(
        "Reprojecting RCP4.5 "
        "using nearest neighbour resampling..."
    )

    rcp45 = warp_rcp_to_target(
        RCP45_FILE,
        target_projection,
        target_gt,
        target_width,
        target_height,
    )

    print(
        "Reprojecting RCP8.5 "
        "using nearest neighbour resampling..."
    )

    rcp85 = warp_rcp_to_target(
        RCP85_FILE,
        target_projection,
        target_gt,
        target_width,
        target_height,
    )

    print(
        "RCP4.5 values:",
        np.unique(rcp45[rcp45 != OUTPUT_NODATA]),
    )

    print(
        "RCP8.5 values:",
        np.unique(rcp85[rcp85 != OUTPUT_NODATA]),
    )

    # --------------------------------------------------
    # Create gain/loss bands
    # --------------------------------------------------

    print("Creating gain/loss bands...")

    rcp45_gainloss = np.full(
        present.shape,
        OUTPUT_NODATA,
        dtype=np.int16,
    )

    rcp85_gainloss = np.full(
        present.shape,
        OUTPUT_NODATA,
        dtype=np.int16,
    )

    valid45 = (
        valid_present
        & (rcp45 != OUTPUT_NODATA)
    )

    valid85 = (
        valid_present
        & (rcp85 != OUTPUT_NODATA)
    )

    # Present non-AES:
    # retain the number of models predicting future AES.
    mask = valid45 & (present_aes == 0)
    rcp45_gainloss[mask] = rcp45[mask]

    mask = valid85 & (present_aes == 0)
    rcp85_gainloss[mask] = rcp85[mask]

    # Present AES:
    # add 10 to the number of models predicting retention.
    mask = valid45 & (present_aes == 1)
    rcp45_gainloss[mask] = 10 + rcp45[mask]

    mask = valid85 & (present_aes == 1)
    rcp85_gainloss[mask] = 10 + rcp85[mask]

    # --------------------------------------------------
    # Build five-band raster
    # --------------------------------------------------

    print("Building multiband raster...")

    stack = gdal.GetDriverByName("MEM").Create(
        "",
        target_width,
        target_height,
        5,
        gdal.GDT_Int16,
    )

    stack.SetGeoTransform(target_gt)
    stack.SetProjection(target_projection)

    arrays = [
        (present_aes, "present_AES"),
        (rcp45, "RCP45_AES_modelCount"),
        (rcp85, "RCP85_AES_modelCount"),
        (rcp45_gainloss, "RCP45_gainloss"),
        (rcp85_gainloss, "RCP85_gainloss"),
    ]

    for band_number, (array, description) in enumerate(
        arrays,
        start=1,
    ):
        band = stack.GetRasterBand(band_number)
        band.WriteArray(array)
        band.SetNoDataValue(OUTPUT_NODATA)
        band.SetDescription(description)

    # --------------------------------------------------
    # Write final COG
    # --------------------------------------------------

    print(f"Creating: {OUTPUT_FILE.name}")

    output = gdal.Translate(
        str(OUTPUT_FILE),
        stack,
        format="COG",
        creationOptions=[
            "COMPRESS=DEFLATE",
            "PREDICTOR=2",
            "BIGTIFF=IF_SAFER",
            "OVERVIEWS=NONE",
        ],
    )

    if output is None:
        raise RuntimeError(f"Failed to create: {OUTPUT_FILE}")

    output = None
    stack = None

    print("\nDone.")
    print(f"AES threshold: {aes_threshold}")
    print("Band 1: present_AES")
    print("Band 2: RCP45_AES_modelCount")
    print("Band 3: RCP85_AES_modelCount")
    print("Band 4: RCP45_gainloss")
    print("Band 5: RCP85_gainloss")
    print("Data type: Int16")
    print("NoData: -999")
    print("Overviews: none")


if __name__ == "__main__":
    main()