#!/usr/bin/env python3

# Convert the Johnson summer chlorophyll-a NetCDF to a georeferenced GeoTIFF.
#
# Constructs the spatial transform explicitly from those coordinate
# variables to generate valid EPSG:3031.
#
# No reprojection or resampling is performed.
#
# Input:
#   cwd/input/chl_summer_climatology_johnson.nc
#
# Output:
#   cwd/output/Johnson_Chlorophyll-a_SummerMean.tif
#
# Install requirements:
#   pip install numpy xarray netCDF4 rasterio


from pathlib import Path

import numpy as np
import rasterio
import xarray as xr
from rasterio.crs import CRS
from rasterio.transform import from_origin


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

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

INPUT_FILE = INPUT_DIR / "chl_summer_climatology_johnson.nc"
OUTPUT_FILE = OUTPUT_DIR / "Johnson_Chlorophyll-a_SummerMean.tif"

VARIABLE_NAME = "mass_concentration_of_chlorophyll_in_sea_water"

X_NAME = "easting"
Y_NAME = "northing"

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


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

def get_spacing(coords, name):
    """Return cell spacing and confirm that the coordinate grid is regular."""

    coords = np.asarray(coords, dtype=np.float64)

    if coords.ndim != 1:
        raise ValueError(f"{name} must be a 1-D coordinate array.")

    if len(coords) < 2:
        raise ValueError(f"{name} contains too few coordinates.")

    diffs = np.diff(coords)
    spacing = float(np.median(np.abs(diffs)))

    if not np.allclose(
        np.abs(diffs),
        spacing,
        rtol=1e-5,
        atol=1e-5,
    ):
        raise ValueError(
            f"{name} coordinates are not regularly spaced."
        )

    return spacing


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

def main():

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

    if not INPUT_FILE.exists():
        raise FileNotFoundError(
            f"Input file not found: {INPUT_FILE}"
        )

    print(f"Input:  {INPUT_FILE.relative_to(Path.cwd())}")
    print(f"Output: {OUTPUT_FILE.relative_to(Path.cwd())}")
    print()

    # -------------------------------------------------------------------------
    # Read NetCDF
    # -------------------------------------------------------------------------

    with xr.open_dataset(
        INPUT_FILE,
        mask_and_scale=True,
    ) as ds:

        if VARIABLE_NAME not in ds:
            raise KeyError(
                f"Variable '{VARIABLE_NAME}' not found.\n"
                f"Available variables: {list(ds.data_vars)}"
            )

        if X_NAME not in ds.variables:
            raise KeyError(
                f"Coordinate '{X_NAME}' not found."
            )

        if Y_NAME not in ds.variables:
            raise KeyError(
                f"Coordinate '{Y_NAME}' not found."
            )

        x = np.asarray(
            ds[X_NAME].values,
            dtype=np.float64,
        )

        y = np.asarray(
            ds[Y_NAME].values,
            dtype=np.float64,
        )

        data_array = ds[VARIABLE_NAME].squeeze(drop=True)

        if data_array.ndim != 2:
            raise ValueError(
                f"Expected a 2-D raster but found dimensions: "
                f"{data_array.dims}"
            )

        # Arrange as rows = northing, columns = easting.
        data_array = data_array.transpose(
            Y_NAME,
            X_NAME,
        )

        data = np.asarray(
            data_array.values,
            dtype=np.float32,
        )

        units = data_array.attrs.get("units", "mg m-3")

        long_name = data_array.attrs.get(
            "long_name",
            VARIABLE_NAME,
        )

    # -------------------------------------------------------------------------
    # Build georeferencing from coordinate arrays
    # -------------------------------------------------------------------------

    x_res = get_spacing(x, X_NAME)
    y_res = get_spacing(y, Y_NAME)

    # GeoTIFF columns must run west -> east.
    if x[0] > x[-1]:
        x = x[::-1]
        data = np.fliplr(data)

    # GeoTIFF rows must run north -> south.
    if y[0] < y[-1]:
        y = y[::-1]
        data = np.flipud(data)

    # NetCDF coordinates represent pixel centres.
    west = x[0] - (x_res / 2)
    north = y[0] + (y_res / 2)

    transform = from_origin(
        west,
        north,
        x_res,
        y_res,
    )

    height, width = data.shape

    # xarray converts the source _FillValue / missing_value to NaN.
    data[~np.isfinite(data)] = OUTPUT_NODATA

    print(f"Raster size:       {width} x {height}")
    print(f"Resolution:        {x_res:g} x {y_res:g} m")
    print("CRS:               EPSG:3031")
    print(f"NoData:            {OUTPUT_NODATA:g}")
    print("Resampling:        none")
    print()

    # -------------------------------------------------------------------------
    # Write cloud-optimised GeoTIFF
    # -------------------------------------------------------------------------

    with rasterio.open(
        OUTPUT_FILE,
        "w",
        driver="COG",
        width=width,
        height=height,
        count=1,
        dtype="float32",
        crs=OUTPUT_CRS,
        transform=transform,
        nodata=OUTPUT_NODATA,
        compress="DEFLATE",
        blocksize=512,
        overview_resampling="average",
    ) as dst:

        dst.write(data, 1)

        dst.set_band_description(
            1,
            "summer_chlorophyll_a_climatology",
        )

        dst.update_tags(
            1,
            long_name=long_name,
            units=units,
        )

    print("Finished.")
    print(
        f"Created: {OUTPUT_FILE.relative_to(Path.cwd())}"
    )


if __name__ == "__main__":
    main()