#!/usr/bin/env python3

# Extract monthly temperature and salinity climatologies at selected depths
# from Yamazaki et al. Southern Ocean NetCDFs.
#
# For each selected depth, creates separate temperature and salinity products:
#   Bands  1-12 = monthly mean, Jan-Dec
#   Bands 13-24 = monthly standard deviation, Jan-Dec
#
# Outputs are reprojected to EPSG:3031 using bilinear resampling,
# standardised to NoData -9999, and written as COGs.
#
# Input:
#   cwd/input/Yamazaki_SO-Monthly-Climatology*.nc
#
# Output:
#   cwd/output/byDepth/
#
# 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" / "byDepth"

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

DST_NODATA = -9999.0

# Catch rogue NetCDF fill-like values as well as NaN / Inf.
FILL_THRESHOLD = 1e35

# Target depths.
DEPTHS_M = [
    5,
    50,
    100,
    200,
    300,
    500,
    1000,
    2000,
]

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

# Exact 66 depth levels in the NetCDF.
ALL_DEPTHS_M = [
    5, 10, 15, 20, 25, 30, 35, 40, 45, 50,
    55, 60, 65, 70, 75, 80, 85, 90, 95, 100,
    125, 150, 175, 200, 225, 250, 275, 300,
    325, 350, 375, 400, 425, 450, 475, 500,
    550, 600, 650, 700, 750, 800, 850, 900,
    950, 1000, 1050, 1100, 1150, 1200, 1250,
    1300, 1350, 1400, 1450, 1500, 1550, 1600,
    1650, 1700, 1750, 1800, 1850, 1900, 1950,
    2000,
]

PARAMETERS = [
    {
        "mean_var": "temp",
        "dev_var": "temp_dev",
        "output_name": "Temperature",
    },
    {
        "mean_var": "psal",
        "dev_var": "psal_dev",
        "output_name": "Salinity",
    },
]


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

gdal.UseExceptions()


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

def find_nc() -> Path:
    """
    Find the full Yamazaki monthly climatology NetCDF.

    Excludes _minimal and _MLD files.
    """

    candidates = sorted(
        path
        for path in INPUT_DIR.glob(
            "Yamazaki_SO-Monthly-Climatology*.nc"
        )
        if "_minimal" not in path.name
        and "_MLD" not in path.name
    )

    if not candidates:
        raise FileNotFoundError(
            f"No full Yamazaki climatology NetCDF found in:\n"
            f"{INPUT_DIR}"
        )

    if len(candidates) > 1:
        raise RuntimeError(
            "More than one possible full climatology NetCDF found:\n"
            + "\n".join(str(path) for path in candidates)
        )

    return candidates[0]


# =============================================================================
# BAND MAPPING
# =============================================================================

def source_bands_for_depth(
    depth_m: int,
) -> list[int]:
    """
    Return the 12 source bands for one requested depth.

    NetCDF structure:
        depth 1: Jan-Dec
        depth 2: Jan-Dec
        ...

    GDAL band numbers are 1-based.
    """

    if depth_m not in ALL_DEPTHS_M:
        raise ValueError(
            f"Requested depth {depth_m} m is not present in the NetCDF."
        )

    depth_index = ALL_DEPTHS_M.index(
        depth_m
    )

    first_band = (
        depth_index * 12
        + 1
    )

    return list(
        range(
            first_band,
            first_band + 12,
        )
    )


# =============================================================================
# EXTRACT ONE VARIABLE / DEPTH
# =============================================================================

def extract_variable_depth(
    nc_path: Path,
    variable: str,
    depth_m: int,
    output_tif: Path,
) -> None:
    """
    Extract Jan-Dec for one variable at one depth.

    Produces a temporary 12-band lon/lat GeoTIFF.
    """

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

    band_indices = source_bands_for_depth(
        depth_m
    )

    print(
        f"\n[EXTRACT] {variable} @ {depth_m} m "
        f"-> source bands "
        f"{band_indices[0]}-{band_indices[-1]}"
    )

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

    if dataset is None:
        raise RuntimeError(
            f"Failed to extract {variable} @ {depth_m} m."
        )

    band_count = dataset.RasterCount

    dataset = None

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


# =============================================================================
# BUILD 24-BAND STACK
# =============================================================================

def build_24_band_stack(
    mean_tif: Path,
    dev_tif: Path,
    mean_var: str,
    dev_var: str,
    output_tif: Path,
) -> list[str]:
    """
    Build:
        Bands  1-12 = mean Jan-Dec
        Bands 13-24 = standard deviation Jan-Dec

    NaN / Inf / fill-like values are converted to -9999.
    """

    descriptions = (
        [
            f"{mean_var}_{month}"
            for month in MONTHS
        ]
        +
        [
            f"{dev_var}_{month}"
            for month in MONTHS
        ]
    )

    mean_ds = gdal.Open(
        str(mean_tif),
        gdal.GA_ReadOnly,
    )

    dev_ds = gdal.Open(
        str(dev_tif),
        gdal.GA_ReadOnly,
    )

    if mean_ds is None:
        raise RuntimeError(
            f"Could not open: {mean_tif}"
        )

    if dev_ds is None:
        raise RuntimeError(
            f"Could not open: {dev_tif}"
        )

    if mean_ds.RasterCount != 12:
        raise RuntimeError(
            f"{mean_var}: expected 12 bands, "
            f"found {mean_ds.RasterCount}"
        )

    if dev_ds.RasterCount != 12:
        raise RuntimeError(
            f"{dev_var}: expected 12 bands, "
            f"found {dev_ds.RasterCount}"
        )

    width = mean_ds.RasterXSize
    height = mean_ds.RasterYSize
    transform = mean_ds.GetGeoTransform()

    if (
        dev_ds.RasterXSize != width
        or dev_ds.RasterYSize != height
        or dev_ds.GetGeoTransform() != transform
    ):
        raise RuntimeError(
            f"{dev_var} grid does not match the mean grid."
        )

    driver = gdal.GetDriverByName(
        "GTiff"
    )

    output = driver.Create(
        str(output_tif),
        width,
        height,
        24,
        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(
        transform
    )

    output.SetProjection(
        mean_ds.GetProjection()
    )

    out_band = 1

    for variable, source_ds in (
        (mean_var, mean_ds),
        (dev_var, dev_ds),
    ):

        for src_band, month in enumerate(
            MONTHS,
            start=1,
        ):

            data = (
                source_ds
                .GetRasterBand(src_band)
                .ReadAsArray()
                .astype(
                    np.float32,
                    copy=False,
                )
            )

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

            if invalid.any():
                data = data.copy()

                data[invalid] = np.float32(
                    DST_NODATA
                )

            description = (
                f"{variable}_{month}"
            )

            destination_band = (
                output.GetRasterBand(
                    out_band
                )
            )

            destination_band.WriteArray(
                data
            )

            destination_band.SetNoDataValue(
                DST_NODATA
            )

            destination_band.SetDescription(
                description
            )

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

            out_band += 1

    output.FlushCache()

    output = None
    mean_ds = None
    dev_ds = None

    return descriptions


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

def warp_to_3031(
    src_tif: Path,
    dst_tif: Path,
) -> None:
    """Reproject 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 survive 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 final Cloud Optimized GeoTIFF."""

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

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

        if sidecar.exists():
            print(
                f"[DELETE] stale sidecar: "
                f"{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


# =============================================================================
# VALIDATE FINAL OUTPUT
# =============================================================================

def validate_output(
    output_path: Path,
    expected_descriptions: list[str],
) -> None:
    """Validate CRS, bands, NoData and descriptions."""

    print(
        f"\n[CHECK] {output_path.name}"
    )

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

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

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

    spatial_ref = dataset.GetSpatialRef()

    epsg = (
        spatial_ref.GetAuthorityCode(None)
        if spatial_ref
        else None
    )

    if epsg != "3031":
        raise RuntimeError(
            f"{output_path.name}: "
            f"expected EPSG:3031, got {epsg}"
        )

    nodata = (
        dataset
        .GetRasterBand(1)
        .GetNoDataValue()
    )

    if nodata != DST_NODATA:
        raise RuntimeError(
            f"{output_path.name}: "
            f"expected NoData {DST_NODATA}, "
            f"got {nodata}"
        )

    actual_descriptions = [
        dataset
        .GetRasterBand(band)
        .GetDescription()
        for band in range(
            1,
            dataset.RasterCount + 1,
        )
    ]

    if (
        actual_descriptions
        != expected_descriptions
    ):
        raise RuntimeError(
            f"{output_path.name}: "
            f"band descriptions do not match "
            f"expected values."
        )

    print("  CRS:        EPSG:3031")
    print(
        f"  Bands:      "
        f"{dataset.RasterCount}"
    )
    print(
        f"  NoData:     {nodata:g}"
    )
    print(
        f"  First band: "
        f"{actual_descriptions[0]}"
    )
    print(
        f"  Band 12:    "
        f"{actual_descriptions[11]}"
    )
    print(
        f"  Band 13:    "
        f"{actual_descriptions[12]}"
    )
    print(
        f"  Last band:  "
        f"{actual_descriptions[-1]}"
    )

    dataset = None


# =============================================================================
# PROCESS ONE PARAMETER / DEPTH
# =============================================================================

def process_one(
    nc_path: Path,
    depth_m: int,
    mean_var: str,
    dev_var: str,
    output_name: str,
) -> None:

    output_path = (
        OUTPUT_DIR
        / (
            f"SO_Monthly_Climatology_"
            f"{depth_m}m_{output_name}.tif"
        )
    )

    print()
    print("=" * 90)
    print(f"DEPTH:      {depth_m} m")
    print(f"PARAMETER:  {output_name}")
    print(f"MEAN:       {mean_var}")
    print(f"STD DEV:    {dev_var}")
    print(
        f"OUTPUT:     "
        f"{output_path.relative_to(BASE_DIR)}"
    )
    print("=" * 90)

    with tempfile.TemporaryDirectory() as temp_dir:

        temp_dir = Path(
            temp_dir
        )

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

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

        stack_ll = (
            temp_dir
            / "stack_lonlat.tif"
        )

        stack_3031 = (
            temp_dir
            / "stack_3031.tif"
        )

        # 1. Extract monthly mean and standard deviation.
        extract_variable_depth(
            nc_path=nc_path,
            variable=mean_var,
            depth_m=depth_m,
            output_tif=mean_tif,
        )

        extract_variable_depth(
            nc_path=nc_path,
            variable=dev_var,
            depth_m=depth_m,
            output_tif=dev_tif,
        )

        # 2. Combine into 24-band stack.
        descriptions = build_24_band_stack(
            mean_tif=mean_tif,
            dev_tif=dev_tif,
            mean_var=mean_var,
            dev_var=dev_var,
            output_tif=stack_ll,
        )

        # 3. Reproject to EPSG:3031.
        warp_to_3031(
            stack_ll,
            stack_3031,
        )

        set_band_descriptions(
            stack_3031,
            descriptions,
        )

        # 4. Write final COG.
        make_cog(
            stack_3031,
            output_path,
        )

    validate_output(
        output_path,
        descriptions,
    )

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


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

def main() -> None:

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

    nc_path = find_nc()

    print("=" * 90)
    print(
        "SOUTHERN OCEAN MONTHLY CLIMATOLOGY"
    )
    print("=" * 90)

    print(
        f"Input:   "
        f"{nc_path.relative_to(BASE_DIR)}"
    )

    print(
        f"Output:  "
        f"{OUTPUT_DIR.relative_to(BASE_DIR)}"
    )

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

    print(
        f"NoData:  {DST_NODATA:g}"
    )

    print(
        f"Depths:  {DEPTHS_M}"
    )

    print()

    print(
        "SOURCE BAND RANGES:"
    )

    for depth in DEPTHS_M:

        bands = source_bands_for_depth(
            depth
        )

        print(
            f"  {depth:4d} m -> "
            f"{bands[0]}-{bands[-1]}"
        )

    for depth_m in DEPTHS_M:

        for parameter in PARAMETERS:

            process_one(
                nc_path=nc_path,
                depth_m=depth_m,
                mean_var=parameter["mean_var"],
                dev_var=parameter["dev_var"],
                output_name=parameter["output_name"],
            )

    print()
    print("=" * 90)
    print(
        f"FINISHED - "
        f"{len(DEPTHS_M) * len(PARAMETERS)} "
        f"COGs CREATED"
    )
    print("=" * 90)


if __name__ == "__main__":
    main()