#!/usr/bin/env python3

# Calculate mean pH and aragonite saturation state for the three SOCOML
# depth ranges used in Figure 8, using the Float Grid.
#
# Input:
#   input/SOCOML_Float_GRID_v2025.nc
#
# Outputs:
#   output/SOCOML_pH_byDepth.tif
#   output/SOCOML_Aragonite_byDepth.tif
#
# Run:
#   python 02_process_socoml.py
#
# Install requirements:
#   pip install xarray netCDF4 numpy rasterio

from pathlib import Path

import numpy as np
import rasterio
import xarray as xr
from rasterio.transform import from_origin
from rasterio.warp import calculate_default_transform, reproject, Resampling


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

INPUT_FILE = Path("input") / "SOCOML_Float_GRID_v2025.nc"
OUTPUT_DIR = Path("output")

NODATA = -999.0
TARGET_CRS = "EPSG:3031"

DEPTH_BANDS = [
    ("subsurface_mean_100-300m", 100, 300),
    ("intermediate_mean_1400-2000m", 1400, 2000),
    ("deep_mean_2000-4000m", 2000, 4000),
]

VARIABLES = {
    "pH": "SOCOML_pH_byDepth.tif",
    "OMEGA_Ar": "SOCOML_Aragonite_byDepth.tif",
}


# =========================
# PROCESSING
# =========================

def process_variable(ds, variable, output_name):

    print(f"\nProcessing {variable}")

    depth = ds["DEPTH"].values
    lat = ds["LATITUDE"].values
    lon = ds["LONGITUDE"].values

    # SOCOML is a regular 1-degree grid with cell centres:
    # lon: -179.5 ... 179.5
    # lat:  -30.5 ... -89.5
    src_transform = from_origin(
        -180.0,
        -30.0,
        1.0,
        1.0,
    )

    src_crs = "EPSG:4326"

    src_width = len(lon)
    src_height = len(lat)

    src_bounds = (
        -180.0,
        -90.0,
        180.0,
        -30.0,
    )

    dst_transform, dst_width, dst_height = calculate_default_transform(
        src_crs,
        TARGET_CRS,
        src_width,
        src_height,
        *src_bounds,
    )

    output_path = OUTPUT_DIR / output_name

    profile = {
        "driver": "GTiff",
        "width": dst_width,
        "height": dst_height,
        "count": len(DEPTH_BANDS),
        "dtype": "float32",
        "crs": TARGET_CRS,
        "transform": dst_transform,
        "nodata": NODATA,
        "compress": "DEFLATE",
        "tiled": True,
        "BIGTIFF": "IF_SAFER",
    }

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

        for band_number, (description, min_depth, max_depth) in enumerate(
            DEPTH_BANDS,
            start=1,
        ):

            indices = np.where(
                (depth >= min_depth) &
                (depth <= max_depth)
            )[0]

            selected_depths = depth[indices]

            print(
                f"  Band {band_number}: {description} "
                f"({len(indices)} levels: "
                f"{selected_depths[0]:g}-{selected_depths[-1]:g} m)"
            )

            data = ds[variable].isel(
                depth_level=indices
            )

            # Exclude source NoData from the depth mean.
            data = data.where(data != NODATA)

            mean_data = data.mean(
                dim="depth_level",
                skipna=True,
            )

            # NetCDF dimensions are lon, lat.
            # Raster dimensions must be row=lat, column=lon.
            mean_data = mean_data.transpose("lat", "lon")

            source = mean_data.values.astype(np.float32)

            source = np.where(
                np.isfinite(source),
                source,
                NODATA,
            )

            destination = np.full(
                (dst_height, dst_width),
                NODATA,
                dtype=np.float32,
            )

            reproject(
                source=source,
                destination=destination,
                src_transform=src_transform,
                src_crs=src_crs,
                src_nodata=NODATA,
                dst_transform=dst_transform,
                dst_crs=TARGET_CRS,
                dst_nodata=NODATA,
                resampling=Resampling.bilinear,
            )

            dst.write(destination, band_number)
            dst.set_band_description(
                band_number,
                description,
            )

    print(f"Created: {output_path}")


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

def main():

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

    OUTPUT_DIR.mkdir(exist_ok=True)

    print(f"Input: {INPUT_FILE}")

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

        for variable, output_name in VARIABLES.items():

            if variable not in ds:
                raise KeyError(
                    f"Variable '{variable}' not found."
                )

            process_variable(
                ds,
                variable,
                output_name,
            )

    print("\nDone.")


if __name__ == "__main__":
    main()