# Combine seasonal sea-ice algae productivity GeoTIFFs into one multiband COG.
#
# Usage:
#   1. Place the five input TIFFs in ./input/
#   2. Edit CONFIG below if required.
#   3. Run: python script.py
#   Output is written to ./output/
#
# Required packages:
#   pip install numpy rasterio

from pathlib import Path

import numpy as np
import rasterio as rio
from rasterio.shutil import copy as rio_copy


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

OUTPUT_FILENAME = "Productivity_SeaIceAlgae_seasonal.tif"

# Band order in output
SEASONS = ["annual", "spring", "summer", "autumn", "winter"]

SOURCE_NODATA = -999
OUTPUT_NODATA = -9999

EXPECTED_EPSG = 3031
UNITS = "mgC/m^2/day"

BLOCKSIZE = 512


# =========================
# PATHS
# =========================

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

OUTPUT_DIR.mkdir(exist_ok=True)


# =========================
# FIND INPUT FILES
# =========================

input_files = {}

for season in SEASONS:
    matches = [
        path for path in INPUT_DIR.glob("*.tif")
        if season.lower() in path.name.lower()
    ]

    if len(matches) != 1:
        raise ValueError(
            f"Expected exactly one TIFF containing '{season}' "
            f"in its filename; found {len(matches)}."
        )

    input_files[season] = matches[0]


# =========================
# CHECK INPUTS
# =========================

with rio.open(input_files[SEASONS[0]]) as src:
    profile = src.profile.copy()
    width = src.width
    height = src.height
    transform = src.transform
    crs = src.crs
    dtype = src.dtypes[0]

if crs != rio.CRS.from_epsg(EXPECTED_EPSG):
    raise ValueError(
        f"Expected EPSG:{EXPECTED_EPSG}, found {crs}."
    )

# Inputs must already be on exactly the same grid.
# No reprojection or resampling is performed.
for season, path in input_files.items():
    with rio.open(path) as src:
        if (
            src.width != width
            or src.height != height
            or src.transform != transform
            or src.crs != crs
        ):
            raise ValueError(
                f"{path.name} does not match the grid of the other inputs."
            )


# =========================
# CREATE MULTIBAND TIFF
# =========================

TEMP_PATH = OUTPUT_DIR / f"{OUTPUT_PATH.stem}_TEMP.tif"

profile.update(
    driver="GTiff",
    count=len(SEASONS),
    nodata=OUTPUT_NODATA,
    tiled=True,
    blockxsize=BLOCKSIZE,
    blockysize=BLOCKSIZE,
    compress="DEFLATE",
    predictor=2,
    BIGTIFF="IF_SAFER",
)

with rio.open(TEMP_PATH, "w", **profile) as dst:

    for band_number, season in enumerate(SEASONS, start=1):
        path = input_files[season]

        with rio.open(path) as src:
            data = src.read(1)

        # Change NoData from -999 to -9999
        data = np.where(
            data == SOURCE_NODATA,
            OUTPUT_NODATA,
            data,
        ).astype(dtype)

        dst.write(data, band_number)

        # Band name and units
        dst.set_band_description(band_number, season)
        dst.set_band_unit(band_number, UNITS)

    # Dataset-level metadata
    dst.update_tags(
        units=UNITS,
        description="Sea ice algae productivity"
    )


# =========================
# CONVERT TO COG
# =========================

rio_copy(
    TEMP_PATH,
    OUTPUT_PATH,
    driver="COG",
    BLOCKSIZE=BLOCKSIZE,
    COMPRESS="DEFLATE",
    PREDICTOR=2,
    BIGTIFF="IF_SAFER",
    NUM_THREADS="ALL_CPUS",
    OVERVIEW_RESAMPLING="NEAREST",
)

TEMP_PATH.unlink()

print(f"Created: {OUTPUT_PATH}")
print("Bands:")
for i, season in enumerate(SEASONS, start=1):
    print(f"  {i}: {season}")