#!/usr/bin/env python3

# Extract mixed-layer monthly climatologies from Yamazaki et al. NetCDFs.
#
# Creates:
#   - Mixed-layer depth
#   - Mixed-layer temperature
#   - Mixed-layer salinity
#
# Monthly bands are retained and annual mean / standard deviation bands
# are calculated from the 12 monthly climatological means.
#
# Outputs are reprojected to EPSG:3031 using bilinear resampling,
# standardised to NoData -9999, and written as COGs.
#
# Input:
#   cwd/input/*_MLD.nc
#
# Output:
#   cwd/output/mixedLayer/
#
# Install requirements:
#   pip install numpy GDAL


from __future__ import annotations

import tempfile
from pathlib import Path

import numpy as np
from osgeo import gdal


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

BASE_DIR = Path.cwd()

INPUT_DIR = BASE_DIR / "input"
OUTPUT_DIR = BASE_DIR / "output" / "mixedLayer"

SOURCE_SRS = "OGC:CRS84"
TARGET_SRS = "EPSG:3031"

DST_NODATA = -9999.0

# Treat very large NetCDF fill values as NoData.
FILL_THRESHOLD = 1e35

MONTHS = [
    "Jan", "Feb", "Mar", "Apr", "May", "Jun",
    "Jul", "Aug", "Sep", "Oct", "Nov", "Dec",
]

OUTPUTS = [
    {
        "filename": "SO_Monthly_Climatology_ML_Depth.tif",
        "mean_var": "MLD",
        "dev_var": None,
    },
    {
        "filename": "SO_Monthly_Climatology_MLD_Temperature.tif",
        "mean_var": "temp",
        "dev_var": "temp_dev",
    },
    {
        "filename": "SO_Monthly_Climatology_MLD_Salinity.tif",
        "mean_var": "psal",
        "dev_var": "psal_dev",
    },
]


# =============================================================================
# GDAL
# =============================================================================

gdal.UseExceptions()


# =============================================================================
# INPUT
# =============================================================================

def find_nc() -> Path:
    """Find the Yamazaki mixed-layer NetCDF in cwd/input."""

    files = sorted(INPUT_DIR.glob("*_MLD.nc"))

    if not files:
        raise FileNotFoundError(
            f"No *_MLD.nc file found in:\n{INPUT_DIR}"
        )

    if len(files) > 1:
        raise RuntimeError(
            "More than one *_MLD.nc file found:\n"
            + "\n".join(str(path) for path in files)
        )

    return files[0]


# =============================================================================
# NETCDF VARIABLE -> TEMPORARY GEOTIFF
# =============================================================================

def materialise_variable(
    nc_path: Path,
    variable: str,
    output_tif: Path,
) -> None:
    """Materialise one NetCDF variable as a 12-band lon/lat GeoTIFF."""

    source = f'NETCDF:"{nc_path}":{variable}'

    dataset = gdal.Translate(
        str(output_tif),
        source,
        options=gdal.TranslateOptions(
            format="GTiff",
            outputSRS=SOURCE_SRS,
            creationOptions=[
                "TILED=YES",
                "COMPRESS=DEFLATE",
                "BIGTIFF=IF_SAFER",
            ],
        ),
    )

    if dataset is None:
        raise RuntimeError(
            f"Failed to materialise NetCDF variable: {variable}"
        )

    band_count = dataset.RasterCount
    width = dataset.RasterXSize
    height = dataset.RasterYSize

    dataset = None

    if band_count != 12:
        raise RuntimeError(
            f"{variable}: expected 12 bands but found {band_count}"
        )

    print(
        f"[OK] {variable}: "
        f"{width} x {height}, {band_count} bands"
    )


# =============================================================================
# READ AND CLEAN MONTHLY BANDS
# =============================================================================

def read_monthly_data(
    tif_path: Path,
) -> tuple[np.ndarray, dict]:
    """
    Read the 12 monthly bands and convert invalid values to NaN.

    Returned array shape:
        (12, rows, columns)
    """

    dataset = gdal.Open(
        str(tif_path),
        gdal.GA_ReadOnly,
    )

    if dataset is None:
        raise RuntimeError(
            f"Could not open: {tif_path}"
        )

    if dataset.RasterCount != 12:
        raise RuntimeError(
            f"{tif_path.name}: expected 12 bands, "
            f"found {dataset.RasterCount}"
        )

    data = dataset.ReadAsArray().astype(
        np.float32,
        copy=False,
    )

    grid = {
        "width": dataset.RasterXSize,
        "height": dataset.RasterYSize,
        "transform": dataset.GetGeoTransform(),
        "projection": dataset.GetProjection(),
    }

    dataset = None

    invalid = (
        ~np.isfinite(data)
        | (np.abs(data) > FILL_THRESHOLD)
        | (data == np.float32(DST_NODATA))
    )

    data = data.copy()
    data[invalid] = np.nan

    return data, grid


# =============================================================================
# ANNUAL STATISTICS
# =============================================================================

def calculate_annual_stats(
    monthly_data: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Calculate annual mean and standard deviation across the 12 monthly
    climatological means.

    Population standard deviation (ddof=0) is used because the 12 months
    represent the complete annual cycle.

    A pixel is valid only where all 12 months are valid.
    """

    complete = np.all(
        np.isfinite(monthly_data),
        axis=0,
    )

    annual_mean = np.full(
        monthly_data.shape[1:],
        DST_NODATA,
        dtype=np.float32,
    )

    annual_dev = np.full(
        monthly_data.shape[1:],
        DST_NODATA,
        dtype=np.float32,
    )

    annual_mean[complete] = np.mean(
        monthly_data[:, complete],
        axis=0,
    ).astype(np.float32)

    annual_dev[complete] = np.std(
        monthly_data[:, complete],
        axis=0,
        ddof=0,
    ).astype(np.float32)

    return annual_mean, annual_dev


def clean_month_band(
    data: np.ndarray,
) -> np.ndarray:
    """Convert NaN / Inf values to -9999."""

    output = data.copy()

    invalid = ~np.isfinite(output)

    if invalid.any():
        output[invalid] = np.float32(
            DST_NODATA
        )

    return output.astype(
        np.float32,
        copy=False,
    )


# =============================================================================
# OUTPUT WRITING
# =============================================================================

def write_band(
    dataset,
    band_number: int,
    data: np.ndarray,
    description: str,
) -> None:
    """Write one Float32 band with NoData and description."""

    band = dataset.GetRasterBand(
        band_number
    )

    band.WriteArray(
        data
    )

    band.SetNoDataValue(
        DST_NODATA
    )

    band.SetDescription(
        description
    )


def build_output_stack(
    mean_tif: Path,
    mean_var: str,
    dev_tif: Path | None,
    dev_var: str | None,
    output_tif: Path,
) -> list[str]:
    """
    Build the lon/lat output stack.

    MLD:
        1-12   MLD_Jan ... MLD_Dec
        13     MLD_Annual
        14     MLD_Annual_dev

    Temperature / Salinity:
        1-12   monthly means
        13-24  monthly source *_dev
        25     annual mean
        26     annual SD across monthly means
    """

    mean_data, grid = read_monthly_data(
        mean_tif
    )

    annual_mean, annual_dev = calculate_annual_stats(
        mean_data
    )

    total_bands = (
        14
        if dev_var is None
        else 26
    )

    driver = gdal.GetDriverByName(
        "GTiff"
    )

    output = driver.Create(
        str(output_tif),
        grid["width"],
        grid["height"],
        total_bands,
        gdal.GDT_Float32,
        options=[
            "TILED=YES",
            "COMPRESS=DEFLATE",
            "BIGTIFF=IF_SAFER",
            "INTERLEAVE=BAND",
        ],
    )

    if output is None:
        raise RuntimeError(
            f"Could not create: {output_tif}"
        )

    output.SetGeoTransform(
        grid["transform"]
    )

    output.SetProjection(
        grid["projection"]
    )

    descriptions: list[str] = []
    out_band = 1

    # -------------------------------------------------------------------------
    # Monthly means
    # -------------------------------------------------------------------------

    for month_index, month in enumerate(MONTHS):

        description = f"{mean_var}_{month}"

        write_band(
            output,
            out_band,
            clean_month_band(
                mean_data[month_index]
            ),
            description,
        )

        descriptions.append(
            description
        )

        print(
            f"  Band {out_band:02d}: {description}"
        )

        out_band += 1

    # -------------------------------------------------------------------------
    # MLD annual bands
    # -------------------------------------------------------------------------

    if dev_var is None:

        description = f"{mean_var}_Annual"

        write_band(
            output,
            out_band,
            annual_mean,
            description,
        )

        descriptions.append(
            description
        )

        print(
            f"  Band {out_band:02d}: {description}"
        )

        out_band += 1

        description = f"{mean_var}_Annual_dev"

        write_band(
            output,
            out_band,
            annual_dev,
            description,
        )

        descriptions.append(
            description
        )

        print(
            f"  Band {out_band:02d}: {description}"
        )

    # -------------------------------------------------------------------------
    # Temperature / salinity monthly deviations + annual bands
    # -------------------------------------------------------------------------

    else:

        if dev_tif is None:
            raise RuntimeError(
                f"{dev_var}: deviation raster was not supplied."
            )

        dev_data, dev_grid = read_monthly_data(
            dev_tif
        )

        if (
            dev_grid["width"] != grid["width"]
            or dev_grid["height"] != grid["height"]
            or dev_grid["transform"] != grid["transform"]
        ):
            raise RuntimeError(
                f"{dev_var} does not use the same grid as {mean_var}."
            )

        for month_index, month in enumerate(MONTHS):

            description = f"{dev_var}_{month}"

            write_band(
                output,
                out_band,
                clean_month_band(
                    dev_data[month_index]
                ),
                description,
            )

            descriptions.append(
                description
            )

            print(
                f"  Band {out_band:02d}: {description}"
            )

            out_band += 1

        description = f"{mean_var}_Annual"

        write_band(
            output,
            out_band,
            annual_mean,
            description,
        )

        descriptions.append(
            description
        )

        print(
            f"  Band {out_band:02d}: {description}"
        )

        out_band += 1

        description = f"{dev_var}_Annual"

        write_band(
            output,
            out_band,
            annual_dev,
            description,
        )

        descriptions.append(
            description
        )

        print(
            f"  Band {out_band:02d}: {description}"
        )

    output.FlushCache()
    output = None

    return descriptions


# =============================================================================
# REPROJECT TO EPSG:3031
# =============================================================================

def warp_to_3031(
    src_tif: Path,
    dst_tif: Path,
) -> None:
    """Reproject the complete stack to EPSG:3031 using bilinear resampling."""

    dataset = gdal.Warp(
        str(dst_tif),
        str(src_tif),
        options=gdal.WarpOptions(
            format="GTiff",
            srcSRS=SOURCE_SRS,
            dstSRS=TARGET_SRS,
            resampleAlg="bilinear",
            srcNodata=DST_NODATA,
            dstNodata=DST_NODATA,
            multithread=True,
            warpOptions=[
                "UNIFIED_SRC_NODATA=NO",
                "NUM_THREADS=ALL_CPUS",
            ],
            creationOptions=[
                "TILED=YES",
                "COMPRESS=DEFLATE",
                "BIGTIFF=IF_SAFER",
            ],
        ),
    )

    if dataset is None:
        raise RuntimeError(
            f"Failed to reproject: {src_tif}"
        )

    dataset = None


# =============================================================================
# BAND DESCRIPTIONS
# =============================================================================

def set_band_descriptions(
    tif_path: Path,
    descriptions: list[str],
) -> None:
    """Ensure band descriptions are retained after reprojection."""

    dataset = gdal.Open(
        str(tif_path),
        gdal.GA_Update,
    )

    if dataset is None:
        raise RuntimeError(
            f"Could not open for update: {tif_path}"
        )

    if dataset.RasterCount != len(descriptions):
        raise RuntimeError(
            f"{tif_path.name}: expected "
            f"{len(descriptions)} bands, "
            f"found {dataset.RasterCount}"
        )

    for band_number, description in enumerate(
        descriptions,
        start=1,
    ):
        dataset.GetRasterBand(
            band_number
        ).SetDescription(
            description
        )

    dataset.FlushCache()
    dataset = None


# =============================================================================
# FINAL COG
# =============================================================================

def make_cog(
    src_tif: Path,
    dst_tif: Path,
) -> None:
    """Write the final Cloud Optimized GeoTIFF."""

    if dst_tif.exists():
        print(
            f"[DELETE] {dst_tif}"
        )
        dst_tif.unlink()

    # Remove stale GDAL sidecars if present.
    for suffix in (
        ".aux.xml",
        ".ovr",
    ):
        sidecar = Path(
            str(dst_tif) + suffix
        )

        if sidecar.exists():
            print(
                f"[DELETE] {sidecar}"
            )
            sidecar.unlink()

    dataset = gdal.Translate(
        str(dst_tif),
        str(src_tif),
        options=gdal.TranslateOptions(
            format="COG",
            creationOptions=[
                "COMPRESS=DEFLATE",
                "BIGTIFF=IF_SAFER",
                "OVERVIEW_RESAMPLING=BILINEAR",
                "NUM_THREADS=ALL_CPUS",
            ],
        ),
    )

    if dataset is None:
        raise RuntimeError(
            f"Failed to create COG: {dst_tif}"
        )

    dataset = None


# =============================================================================
# PROCESS ONE OUTPUT
# =============================================================================

def process_output(
    nc_path: Path,
    filename: str,
    mean_var: str,
    dev_var: str | None,
) -> None:

    output_path = (
        OUTPUT_DIR
        / filename
    )

    print()
    print("=" * 80)
    print(f"OUTPUT:   {output_path}")
    print(f"MEAN:     {mean_var}")

    if dev_var:
        print(f"DEV:      {dev_var}")

    print("=" * 80)

    with tempfile.TemporaryDirectory() as temp_dir:

        temp_dir = Path(
            temp_dir
        )

        mean_tif = (
            temp_dir
            / f"{mean_var}.tif"
        )

        materialise_variable(
            nc_path,
            mean_var,
            mean_tif,
        )

        dev_tif = None

        if dev_var is not None:

            dev_tif = (
                temp_dir
                / f"{dev_var}.tif"
            )

            materialise_variable(
                nc_path,
                dev_var,
                dev_tif,
            )

        stack_ll = (
            temp_dir
            / "stack_lonlat.tif"
        )

        descriptions = build_output_stack(
            mean_tif=mean_tif,
            mean_var=mean_var,
            dev_tif=dev_tif,
            dev_var=dev_var,
            output_tif=stack_ll,
        )

        stack_3031 = (
            temp_dir
            / "stack_3031.tif"
        )

        warp_to_3031(
            stack_ll,
            stack_3031,
        )

        set_band_descriptions(
            stack_3031,
            descriptions,
        )

        make_cog(
            stack_3031,
            output_path,
        )

    print(
        f"\n[OK] {output_path.relative_to(BASE_DIR)}"
    )


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

def main() -> None:

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

    nc_path = find_nc()

    print(
        f"[INPUT]  {nc_path.relative_to(BASE_DIR)}"
    )

    print(
        f"[OUTPUT] {OUTPUT_DIR.relative_to(BASE_DIR)}"
    )

    print(
        f"[CRS]    {SOURCE_SRS} -> {TARGET_SRS}"
    )

    print(
        f"[NODATA] {DST_NODATA:g}"
    )

    for item in OUTPUTS:

        process_output(
            nc_path=nc_path,
            filename=item["filename"],
            mean_var=item["mean_var"],
            dev_var=item["dev_var"],
        )

    print()
    print("=" * 80)
    print("FINISHED")
    print("=" * 80)


if __name__ == "__main__":
    main()