#!/usr/bin/env python3

# Combine CAFE summer NPP mean and standard deviation rasters into
# a 2-band EPSG:3031 GeoTIFF.
#
# Band 1 = mean
# Band 2 = stdev
#
# Input:
#   cwd/input/*_mean.tif
#   cwd/input/*_stdev.tif
#
# Output:
#   cwd/output/CAFE_NPP_Summer.tif
#
#
# Install requirements:
#   pip install rasterio numpy


from pathlib import Path

import numpy as np
import rasterio
from rasterio.crs import CRS


# =============================================================================
# CONFIGURATION
# =============================================================================

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

OUTPUT_FILE = OUTPUT_DIR / "CAFE_NPP_Summer.tif"

OUTPUT_CRS = CRS.from_epsg(3031)
OUTPUT_NODATA = -9999.0


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

def main():

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    # Find input files.
    mean_files = list(INPUT_DIR.glob("*_mean.tif"))
    stdev_files = list(INPUT_DIR.glob("*_stdev.tif"))

    if len(mean_files) != 1:
        raise RuntimeError(
            f"Expected exactly one *_mean.tif in {INPUT_DIR}, "
            f"found {len(mean_files)}."
        )

    if len(stdev_files) != 1:
        raise RuntimeError(
            f"Expected exactly one *_stdev.tif in {INPUT_DIR}, "
            f"found {len(stdev_files)}."
        )

    mean_file = mean_files[0]
    stdev_file = stdev_files[0]

    print(f"Mean:   {mean_file.relative_to(Path.cwd())}")
    print(f"Stdev:  {stdev_file.relative_to(Path.cwd())}")
    print(f"Output: {OUTPUT_FILE.relative_to(Path.cwd())}")
    print()

    # -------------------------------------------------------------------------
    # Read mean raster
    # -------------------------------------------------------------------------

    with rasterio.open(mean_file) as src_mean:

        mean = src_mean.read(1).astype(np.float32)

        transform = src_mean.transform
        width = src_mean.width
        height = src_mean.height

        mean_nodata = src_mean.nodata

        profile = src_mean.profile.copy()

    # -------------------------------------------------------------------------
    # Read standard deviation raster
    # -------------------------------------------------------------------------

    with rasterio.open(stdev_file) as src_stdev:

        # Confirm both rasters use exactly the same grid.
        if src_stdev.width != width or src_stdev.height != height:
            raise ValueError(
                "Mean and stdev rasters have different dimensions."
            )

        if src_stdev.transform != transform:
            raise ValueError(
                "Mean and stdev rasters have different geotransforms."
            )

        stdev = src_stdev.read(1).astype(np.float32)

        stdev_nodata = src_stdev.nodata

    # -------------------------------------------------------------------------
    # Convert source NoData to -9999
    # -------------------------------------------------------------------------

    if mean_nodata is not None:
        mean[mean == mean_nodata] = OUTPUT_NODATA

    if stdev_nodata is not None:
        stdev[stdev == stdev_nodata] = OUTPUT_NODATA

    # Also catch NaN / Inf values.
    mean[~np.isfinite(mean)] = OUTPUT_NODATA
    stdev[~np.isfinite(stdev)] = OUTPUT_NODATA

    # -------------------------------------------------------------------------
    # Write 2-band GeoTIFF
    # -------------------------------------------------------------------------

    profile.update(
        driver="COG",
        count=2,
        dtype="float32",
        crs=OUTPUT_CRS,
        nodata=OUTPUT_NODATA,
        compress="DEFLATE",
        blocksize=512,
        overview_resampling="average",
    )

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

        dst.write(mean, 1)
        dst.write(stdev, 2)

        dst.set_band_description(1, "mean")
        dst.set_band_description(2, "stdev")

    print("Finished.")
    print(f"Created: {OUTPUT_FILE.relative_to(Path.cwd())}")
    print("Bands:")
    print("  1 = mean")
    print("  2 = stdev")
    print("CRS: EPSG:3031")
    print(f"NoData: {OUTPUT_NODATA:g}")


if __name__ == "__main__":
    main()