#!/usr/bin/env python3

# Create bottom-water temperature and salinity climatologies from the
# Yamazaki Southern Ocean monthly climatology.
#
# For each source-grid pixel:
#   - Find deepest depth where all monthly temp/salinity means are valid
#   - Extract temp,salinity and monthly deviations at that same depth
#   - Calculate annual mean and annual standard deviation from the 12
#     monthly climatological means
#
# Resulting products are reprojected to EPSG:3031 (bilinear resampling)
# masked to seabed depths of 2000 m or shallower using IBCSO.
#
# Outputs:
#   cwd/output/bottomWater/SO_Monthly_Climatology_bottom_Temperature.tif
#   cwd/output/bottomWater/SO_Monthly_Climatology_bottom_Salinity.tif
#
# Install requirements:
#   pip install numpy GDAL


from __future__ import annotations

import tempfile
from pathlib import Path

import numpy as np
from osgeo import gdal, osr


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

BASE_DIR = Path.cwd()

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

NC_PATH = INPUT_DIR / "Yamazaki_SO-Monthly-Climatology_v20240604.nc"
IBCSO_PATH = INPUT_DIR / "IBCSO_v2_bed.tif"

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

DST_NODATA = -9999.0

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

# Retain seabed exactly -2000 m and shallower.
# Values below -2000 m are deeper and are masked.
IBCSO_MIN_BED = -2000.0

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

# Exact 66 depth levels in the Yamazaki 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,
]

TEMP_DESCRIPTIONS = (
    [f"temp_{month}" for month in MONTHS]
    + [f"temp_dev_{month}" for month in MONTHS]
    + ["temp_Annual", "temp_dev_Annual"]
)

PSAL_DESCRIPTIONS = (
    [f"psal_{month}" for month in MONTHS]
    + [f"psal_dev_{month}" for month in MONTHS]
    + ["psal_Annual", "psal_dev_Annual"]
)


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

gdal.UseExceptions()


# =============================================================================
# HELPERS
# =============================================================================

def nc_variable(variable: str) -> str:
    """Return a GDAL NetCDF subdataset path."""

    return f'NETCDF:"{NC_PATH}":{variable}'


def srs_to_wkt(srs_name: str) -> str:
    """Convert an SRS identifier to WKT."""

    srs = osr.SpatialReference()
    srs.SetFromUserInput(srs_name)

    return srs.ExportToWkt()


def is_epsg(dataset, epsg: int) -> bool:
    """Return True if a dataset uses the requested EPSG CRS."""

    source = dataset.GetSpatialRef()

    if source is None:
        return False

    target = osr.SpatialReference()
    target.ImportFromEPSG(epsg)

    return bool(
        source.IsSame(target)
    )


def bands_for_depth_index(
    depth_index: int,
) -> list[int]:
    """
    Return the 12 monthly bands for one depth index.

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

    GDAL band numbers are 1-based.
    """

    first_band = (
        depth_index * 12
        + 1
    )

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


def read_bands_masked(
    dataset,
    bands: list[int],
) -> np.ma.MaskedArray:
    """
    Read multiple GDAL bands as one masked array.

    The mask respects the GDAL mask band and source NoData value.
    """

    arrays = []
    masks = []

    for band_number in bands:

        band = dataset.GetRasterBand(
            band_number
        )

        data = band.ReadAsArray()

        mask = (
            band.GetMaskBand()
            .ReadAsArray()
            == 0
        )

        nodata = band.GetNoDataValue()

        if nodata is not None:

            if np.isnan(nodata):
                mask |= np.isnan(data)

            else:
                mask |= data == nodata

        arrays.append(
            data
        )

        masks.append(
            mask
        )

    return np.ma.array(
        np.stack(arrays),
        mask=np.stack(masks),
    )


def valid_mask(
    data: np.ma.MaskedArray,
) -> np.ndarray:
    """
    Return True for genuine source data.

    Invalid values include:
      - GDAL-masked values
      - source NoData
      - NaN / Inf
      - very large fill-like values
    """

    raw = np.asarray(
        data.data
    )

    return (
        ~np.ma.getmaskarray(data)
        & np.isfinite(raw)
        & (
            np.abs(raw)
            <= FILL_THRESHOLD
        )
    )


def clean_raw(
    data: np.ma.MaskedArray,
) -> np.ndarray:
    """Convert source data to Float32 and force invalid values to -9999."""

    raw = np.asarray(
        data.data,
        dtype=np.float32,
    ).copy()

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

    raw[invalid] = np.float32(
        DST_NODATA
    )

    return raw


# =============================================================================
# FIND DEEPEST COMMON TEMPERATURE / SALINITY DEPTH
# =============================================================================

def find_deepest_common_depth() -> np.ndarray:
    """
    Find the deepest depth at every source-grid pixel where:

      temperature Jan-Dec are all valid
      AND
      salinity Jan-Dec are all valid.

    The returned array stores the depth INDEX, not metres.

    -1 = no valid common depth.
    """

    expected_bands = (
        len(ALL_DEPTHS_M)
        * 12
    )

    temp_ds = gdal.Open(
        nc_variable("temp"),
        gdal.GA_ReadOnly,
    )

    psal_ds = gdal.Open(
        nc_variable("psal"),
        gdal.GA_ReadOnly,
    )

    if temp_ds is None:
        raise RuntimeError(
            "Could not open NetCDF variable: temp"
        )

    if psal_ds is None:
        raise RuntimeError(
            "Could not open NetCDF variable: psal"
        )

    if temp_ds.RasterCount != expected_bands:
        raise RuntimeError(
            f"temp: expected {expected_bands} bands, "
            f"found {temp_ds.RasterCount}"
        )

    if psal_ds.RasterCount != expected_bands:
        raise RuntimeError(
            f"psal: expected {expected_bands} bands, "
            f"found {psal_ds.RasterCount}"
        )

    if (
        temp_ds.RasterXSize != psal_ds.RasterXSize
        or temp_ds.RasterYSize != psal_ds.RasterYSize
        or temp_ds.GetGeoTransform()
        != psal_ds.GetGeoTransform()
    ):
        raise RuntimeError(
            "Temperature and salinity grids do not match."
        )

    height = temp_ds.RasterYSize
    width = temp_ds.RasterXSize

    selected = np.full(
        (height, width),
        -1,
        dtype=np.int16,
    )

    # Search deepest -> shallowest.
    # Once a pixel receives a valid depth, leave it unchanged.
    for depth_index in range(
        len(ALL_DEPTHS_M) - 1,
        -1,
        -1,
    ):

        depth_m = ALL_DEPTHS_M[
            depth_index
        ]

        print(
            f"[DEPTH SEARCH] {depth_m:4d} m",
            flush=True,
        )

        bands = bands_for_depth_index(
            depth_index
        )

        temp = read_bands_masked(
            temp_ds,
            bands,
        )

        psal = read_bands_masked(
            psal_ds,
            bands,
        )

        temp_all_months = np.all(
            valid_mask(temp),
            axis=0,
        )

        psal_all_months = np.all(
            valid_mask(psal),
            axis=0,
        )

        valid_here = (
            temp_all_months
            & psal_all_months
        )

        take = (
            (selected < 0)
            & valid_here
        )

        selected[take] = depth_index

        if np.all(
            selected >= 0
        ):
            break

    temp_ds = None
    psal_ds = None

    n_selected = int(
        (selected >= 0).sum()
    )

    print()
    print(
        f"[DEPTH SEARCH] Pixels with common T/S depth: "
        f"{n_selected:,}"
    )

    print()
    print(
        "[SELECTED DEPTH DISTRIBUTION]"
    )

    for depth_index, depth_m in enumerate(
        ALL_DEPTHS_M
    ):

        count = int(
            (selected == depth_index).sum()
        )

        if count:
            print(
                f"  {depth_m:4d} m : {count:,}"
            )

    return selected


# =============================================================================
# WRITE SOURCE STACK
# =============================================================================

def write_stack(
    path: Path,
    stack: np.ndarray,
    transform: tuple,
    descriptions: list[str],
) -> None:
    """Write a clean 26-band source-grid GeoTIFF."""

    height = stack.shape[1]
    width = stack.shape[2]

    driver = gdal.GetDriverByName(
        "GTiff"
    )

    dataset = driver.Create(
        str(path),
        width,
        height,
        26,
        gdal.GDT_Float32,
        options=[
            "TILED=YES",
            "COMPRESS=DEFLATE",
            "BIGTIFF=IF_SAFER",
            "INTERLEAVE=BAND",
        ],
    )

    if dataset is None:
        raise RuntimeError(
            f"Could not create: {path}"
        )

    dataset.SetGeoTransform(
        transform
    )

    dataset.SetProjection(
        srs_to_wkt(SOURCE_SRS)
    )

    for band_number in range(
        1,
        27,
    ):

        data = stack[
            band_number - 1
        ].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
            )

        band = dataset.GetRasterBand(
            band_number
        )

        band.WriteArray(
            data
        )

        band.SetNoDataValue(
            DST_NODATA
        )

        band.SetDescription(
            descriptions[
                band_number - 1
            ]
        )

    dataset.FlushCache()
    dataset = None

    print(
        f"[OK] Built source TIFF: {path.name}"
    )


# =============================================================================
# BUILD 26-BAND SOURCE PRODUCTS
# =============================================================================

def build_source_products(
    selected_depth: np.ndarray,
    temp_output: Path,
    psal_output: Path,
) -> None:
    """
    Build source-grid 26-band temperature and salinity products.

    Temperature:
       1-12  temp Jan-Dec
      13-24  temp_dev Jan-Dec
      25     annual mean of monthly temp means
      26     SD across monthly temp means

    Salinity:
       1-12  psal Jan-Dec
      13-24  psal_dev Jan-Dec
      25     annual mean of monthly psal means
      26     SD across monthly psal means

    The exact same selected depth is used for temperature, salinity
    and all deviation bands.
    """

    datasets = {
        variable: gdal.Open(
            nc_variable(variable),
            gdal.GA_ReadOnly,
        )
        for variable in (
            "temp",
            "temp_dev",
            "psal",
            "psal_dev",
        )
    }

    for variable, dataset in datasets.items():

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

    reference = datasets[
        "temp"
    ]

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

    if selected_depth.shape != (
        height,
        width,
    ):
        raise RuntimeError(
            "Selected-depth grid does not match NetCDF grid."
        )

    # All four source variables must use the same grid.
    for variable, dataset in datasets.items():

        if (
            dataset.RasterXSize != width
            or dataset.RasterYSize != height
            or dataset.GetGeoTransform()
            != transform
        ):
            raise RuntimeError(
                f"{variable}: source grid does not match temp."
            )

    temp_stack = np.full(
        (26, height, width),
        np.float32(DST_NODATA),
        dtype=np.float32,
    )

    psal_stack = np.full(
        (26, height, width),
        np.float32(DST_NODATA),
        dtype=np.float32,
    )

    # -------------------------------------------------------------------------
    # Fill monthly means and deviations at selected depth
    # -------------------------------------------------------------------------

    for depth_index, depth_m in enumerate(
        ALL_DEPTHS_M
    ):

        pixel_mask = (
            selected_depth
            == depth_index
        )

        n_pixels = int(
            pixel_mask.sum()
        )

        if n_pixels == 0:
            continue

        print(
            f"[EXTRACT] {depth_m:4d} m : "
            f"{n_pixels:,} selected pixels",
            flush=True,
        )

        bands = bands_for_depth_index(
            depth_index
        )

        temp = read_bands_masked(
            datasets["temp"],
            bands,
        )

        temp_dev = read_bands_masked(
            datasets["temp_dev"],
            bands,
        )

        psal = read_bands_masked(
            datasets["psal"],
            bands,
        )

        psal_dev = read_bands_masked(
            datasets["psal_dev"],
            bands,
        )

        # Preserve the original guard: all four variables must be valid
        # at the selected common T/S depth.
        everything_valid = (
            np.all(
                valid_mask(temp),
                axis=0,
            )
            & np.all(
                valid_mask(temp_dev),
                axis=0,
            )
            & np.all(
                valid_mask(psal),
                axis=0,
            )
            & np.all(
                valid_mask(psal_dev),
                axis=0,
            )
        )

        if not np.all(
            everything_valid[
                pixel_mask
            ]
        ):

            bad = int(
                np.sum(
                    pixel_mask
                    & ~everything_valid
                )
            )

            raise RuntimeError(
                f"{bad:,} selected pixels at {depth_m} m "
                f"do not have valid mean/dev data."
            )

        temp_raw = clean_raw(
            temp
        )

        temp_dev_raw = clean_raw(
            temp_dev
        )

        psal_raw = clean_raw(
            psal
        )

        psal_dev_raw = clean_raw(
            psal_dev
        )

        temp_stack[
            0:12,
            pixel_mask,
        ] = temp_raw[
            :,
            pixel_mask,
        ]

        temp_stack[
            12:24,
            pixel_mask,
        ] = temp_dev_raw[
            :,
            pixel_mask,
        ]

        psal_stack[
            0:12,
            pixel_mask,
        ] = psal_raw[
            :,
            pixel_mask,
        ]

        psal_stack[
            12:24,
            pixel_mask,
        ] = psal_dev_raw[
            :,
            pixel_mask,
        ]

    # -------------------------------------------------------------------------
    # Annual mean and population SD from monthly climatological means
    # -------------------------------------------------------------------------

    selected = (
        selected_depth >= 0
    )

    temp_monthly = temp_stack[
        0:12,
        selected,
    ]

    psal_monthly = psal_stack[
        0:12,
        selected,
    ]

    temp_stack[
        24,
        selected,
    ] = np.mean(
        temp_monthly,
        axis=0,
    ).astype(
        np.float32
    )

    psal_stack[
        24,
        selected,
    ] = np.mean(
        psal_monthly,
        axis=0,
    ).astype(
        np.float32
    )

    temp_stack[
        25,
        selected,
    ] = np.std(
        temp_monthly,
        axis=0,
        ddof=0,
    ).astype(
        np.float32
    )

    psal_stack[
        25,
        selected,
    ] = np.std(
        psal_monthly,
        axis=0,
        ddof=0,
    ).astype(
        np.float32
    )

    # Close NetCDF datasets before writing outputs.
    for variable in datasets:
        datasets[variable] = None

    write_stack(
        temp_output,
        temp_stack,
        transform,
        TEMP_DESCRIPTIONS,
    )

    write_stack(
        psal_output,
        psal_stack,
        transform,
        PSAL_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


# =============================================================================
# ALIGN IBCSO TO OUTPUT GRID
# =============================================================================

def build_ibcso_keep_mask(
    target_tif: Path,
) -> np.ndarray:
    """
    Resample IBCSO onto the exact EPSG:3031 Yamazaki output grid.

    True  = retain
    False = mask to -9999

    Seabed exactly -2000 m is retained.
    """

    target = gdal.Open(
        str(target_tif),
        gdal.GA_ReadOnly,
    )

    ibcso = gdal.Open(
        str(IBCSO_PATH),
        gdal.GA_ReadOnly,
    )

    if target is None:
        raise RuntimeError(
            f"Could not open: {target_tif}"
        )

    if ibcso is None:
        raise RuntimeError(
            f"Could not open: {IBCSO_PATH}"
        )

    if not is_epsg(
        target,
        3031,
    ):
        raise RuntimeError(
            "Target grid is not EPSG:3031."
        )

    if not is_epsg(
        ibcso,
        3031,
    ):
        raise RuntimeError(
            "IBCSO is not EPSG:3031."
        )

    width = target.RasterXSize
    height = target.RasterYSize

    transform = target.GetGeoTransform()

    # gdal.Warp output grids are north-up, but check explicitly because
    # the exact-grid calculation below assumes no rotation.
    if (
        transform[2] != 0
        or transform[4] != 0
    ):
        raise RuntimeError(
            "Target raster has a rotated geotransform."
        )

    left = transform[0]
    top = transform[3]

    right = (
        left
        + width * transform[1]
    )

    bottom = (
        top
        + height * transform[5]
    )

    ibcso_nodata = (
        ibcso.GetRasterBand(1)
        .GetNoDataValue()
    )

    print()
    print(
        "[IBCSO] Resampling bathymetry onto "
        "the Yamazaki EPSG:3031 grid..."
    )

    warp_kwargs = {
        "format": "MEM",
        "srcSRS": ibcso.GetProjection(),
        "dstSRS": target.GetProjection(),
        "outputBounds": (
            left,
            bottom,
            right,
            top,
        ),
        "width": width,
        "height": height,
        "resampleAlg": "bilinear",
        "dstNodata": DST_NODATA,
    }

    if ibcso_nodata is not None:
        warp_kwargs[
            "srcNodata"
        ] = ibcso_nodata

    aligned = gdal.Warp(
        "",
        ibcso,
        options=gdal.WarpOptions(
            **warp_kwargs
        ),
    )

    if aligned is None:
        raise RuntimeError(
            "Failed to align IBCSO to target grid."
        )

    # Confirm that GDAL generated exactly the requested target grid.
    if (
        aligned.RasterXSize != width
        or aligned.RasterYSize != height
        or aligned.GetGeoTransform()
        != transform
    ):
        raise RuntimeError(
            "Aligned IBCSO grid does not exactly match target grid."
        )

    bathy = (
        aligned
        .GetRasterBand(1)
        .ReadAsArray()
        .astype(
            np.float32,
            copy=False,
        )
    )

    aligned = None
    target = None
    ibcso = None

    valid_bathy = (
        np.isfinite(bathy)
        & (
            np.abs(bathy)
            <= FILL_THRESHOLD
        )
        & (
            bathy
            != np.float32(
                DST_NODATA
            )
        )
    )

    # -2000 m is retained.
    # Anything less than -2000 m is deeper and removed.
    keep = (
        valid_bathy
        & (
            bathy
            >= IBCSO_MIN_BED
        )
    )

    n_total = bathy.size
    n_valid = int(
        valid_bathy.sum()
    )
    n_keep = int(
        keep.sum()
    )

    n_deep = int(
        (
            valid_bathy
            & (
                bathy
                < IBCSO_MIN_BED
            )
        ).sum()
    )

    n_no_bathy = (
        n_total
        - n_valid
    )

    print()
    print("[IBCSO MASK]")
    print(
        f"  Target pixels:              {n_total:,}"
    )
    print(
        f"  Valid IBCSO pixels:         {n_valid:,}"
    )
    print(
        f"  Retained (>= -2000 m):      {n_keep:,}"
    )
    print(
        f"  Masked deeper than 2000 m:  {n_deep:,}"
    )
    print(
        f"  Masked: no IBCSO data:      {n_no_bathy:,}"
    )

    return keep


# =============================================================================
# APPLY IBCSO MASK
# =============================================================================

def apply_ibcso_mask(
    src_tif: Path,
    dst_tif: Path,
    keep_mask: np.ndarray,
    descriptions: list[str],
) -> None:
    """
    Apply the IBCSO <=2000 m seabed domain.

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

    src = gdal.Open(
        str(src_tif),
        gdal.GA_ReadOnly,
    )

    if src is None:
        raise RuntimeError(
            f"Could not open: {src_tif}"
        )

    if keep_mask.shape != (
        src.RasterYSize,
        src.RasterXSize,
    ):
        raise RuntimeError(
            "IBCSO mask does not match warped raster grid."
        )

    if src.RasterCount != 26:
        raise RuntimeError(
            f"{src_tif.name}: expected 26 bands, "
            f"found {src.RasterCount}"
        )

    driver = gdal.GetDriverByName(
        "GTiff"
    )

    dst = driver.Create(
        str(dst_tif),
        src.RasterXSize,
        src.RasterYSize,
        26,
        gdal.GDT_Float32,
        options=[
            "TILED=YES",
            "COMPRESS=DEFLATE",
            "BIGTIFF=IF_SAFER",
            "INTERLEAVE=BAND",
        ],
    )

    if dst is None:
        raise RuntimeError(
            f"Could not create: {dst_tif}"
        )

    dst.SetGeoTransform(
        src.GetGeoTransform()
    )

    dst.SetProjection(
        src.GetProjection()
    )

    for band_number in range(
        1,
        27,
    ):

        data = (
            src
            .GetRasterBand(
                band_number
            )
            .ReadAsArray()
            .astype(
                np.float32,
                copy=False,
            )
        )

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

        if invalid.any():

            data = data.copy()

            data[invalid] = np.float32(
                DST_NODATA
            )

        band = dst.GetRasterBand(
            band_number
        )

        band.WriteArray(
            data
        )

        band.SetNoDataValue(
            DST_NODATA
        )

        band.SetDescription(
            descriptions[
                band_number - 1
            ]
        )

    dst.FlushCache()

    dst = None
    src = None

    print(
        f"[OK] Applied IBCSO mask: {dst_tif.name}"
    )


# =============================================================================
# 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()

    for suffix in (
        ".aux.xml",
        ".ovr",
    ):

        sidecar = Path(
            str(dst_tif)
            + suffix
        )

        if sidecar.exists():

            print(
                f"[DELETE] stale sidecar: {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,
    descriptions: list[str],
) -> None:
    """Validate CRS, band count, NoData and band descriptions."""

    print()
    print(
        f"[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: {output_path}"
        )

    if dataset.RasterCount != 26:
        raise RuntimeError(
            f"Expected 26 bands, "
            f"found {dataset.RasterCount}"
        )

    if not is_epsg(
        dataset,
        3031,
    ):
        raise RuntimeError(
            "Expected EPSG:3031."
        )

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

    if nodata != DST_NODATA:
        raise RuntimeError(
            f"Expected NoData {DST_NODATA}, "
            f"got {nodata}"
        )

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

    if (
        actual_descriptions
        != descriptions
    ):
        raise RuntimeError(
            "Band descriptions do not match expected values."
        )

    print(
        "  CRS:       EPSG:3031"
    )
    print(
        f"  Bands:     {dataset.RasterCount}"
    )
    print(
        f"  NoData:    {nodata:g}"
    )
    print(
        f"  Band 1:    {actual_descriptions[0]}"
    )
    print(
        f"  Band 12:   {actual_descriptions[11]}"
    )
    print(
        f"  Band 13:   {actual_descriptions[12]}"
    )
    print(
        f"  Band 24:   {actual_descriptions[23]}"
    )
    print(
        f"  Band 25:   {actual_descriptions[24]}"
    )
    print(
        f"  Band 26:   {actual_descriptions[25]}"
    )

    dataset = None


# =============================================================================
# CHECK WARPED GRIDS
# =============================================================================

def check_matching_grids(
    first_path: Path,
    second_path: Path,
) -> None:
    """Confirm two warped products use exactly the same grid."""

    first = gdal.Open(
        str(first_path),
        gdal.GA_ReadOnly,
    )

    second = gdal.Open(
        str(second_path),
        gdal.GA_ReadOnly,
    )

    if first is None or second is None:
        raise RuntimeError(
            "Could not open warped products for grid validation."
        )

    same_crs = False

    first_srs = (
        first.GetSpatialRef()
    )

    second_srs = (
        second.GetSpatialRef()
    )

    if (
        first_srs is not None
        and second_srs is not None
    ):
        same_crs = bool(
            first_srs.IsSame(
                second_srs
            )
        )

    if (
        first.RasterXSize
        != second.RasterXSize
        or first.RasterYSize
        != second.RasterYSize
        or first.GetGeoTransform()
        != second.GetGeoTransform()
        or not same_crs
    ):
        raise RuntimeError(
            "Warped temperature and salinity grids do not match."
        )

    first = None
    second = None


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

def main() -> None:

    if not NC_PATH.exists():
        raise FileNotFoundError(
            f"Could not find:\n{NC_PATH}"
        )

    if not IBCSO_PATH.exists():
        raise FileNotFoundError(
            f"Could not find:\n{IBCSO_PATH}"
        )

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

    final_temp = (
        OUTPUT_DIR
        / "SO_Monthly_Climatology_bottom_Temperature.tif"
    )

    final_psal = (
        OUTPUT_DIR
        / "SO_Monthly_Climatology_bottom_Salinity.tif"
    )

    print(
        f"[NETCDF] {NC_PATH.relative_to(BASE_DIR)}"
    )

    print(
        f"[IBCSO]  {IBCSO_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}"
    )

    # -------------------------------------------------------------------------
    # 1. Determine common deepest temperature / salinity depth
    # -------------------------------------------------------------------------

    selected_depth = (
        find_deepest_common_depth()
    )

    # -------------------------------------------------------------------------
    # 2. Build, warp and mask products
    # -------------------------------------------------------------------------

    with tempfile.TemporaryDirectory(
        prefix="yamazaki_deepest_"
    ) as temp_dir:

        temp_dir = Path(
            temp_dir
        )

        source_temp = (
            temp_dir
            / "temperature_source.tif"
        )

        source_psal = (
            temp_dir
            / "salinity_source.tif"
        )

        warped_temp = (
            temp_dir
            / "temperature_3031.tif"
        )

        warped_psal = (
            temp_dir
            / "salinity_3031.tif"
        )

        masked_temp = (
            temp_dir
            / "temperature_masked.tif"
        )

        masked_psal = (
            temp_dir
            / "salinity_masked.tif"
        )

        build_source_products(
            selected_depth,
            source_temp,
            source_psal,
        )

        print()
        print(
            "[WARP] Temperature"
        )

        warp_to_3031(
            source_temp,
            warped_temp,
        )

        print()
        print(
            "[WARP] Salinity"
        )

        warp_to_3031(
            source_psal,
            warped_psal,
        )

        check_matching_grids(
            warped_temp,
            warped_psal,
        )

        # ---------------------------------------------------------------------
        # 3. Apply IBCSO 2000 m seabed cutoff
        # ---------------------------------------------------------------------

        keep_mask = (
            build_ibcso_keep_mask(
                warped_temp
            )
        )

        apply_ibcso_mask(
            warped_temp,
            masked_temp,
            keep_mask,
            TEMP_DESCRIPTIONS,
        )

        apply_ibcso_mask(
            warped_psal,
            masked_psal,
            keep_mask,
            PSAL_DESCRIPTIONS,
        )

        # ---------------------------------------------------------------------
        # 4. Final COGs
        # ---------------------------------------------------------------------

        make_cog(
            masked_temp,
            final_temp,
        )

        make_cog(
            masked_psal,
            final_psal,
        )

    # -------------------------------------------------------------------------
    # 5. Validate
    # -------------------------------------------------------------------------

    validate_output(
        final_temp,
        TEMP_DESCRIPTIONS,
    )

    validate_output(
        final_psal,
        PSAL_DESCRIPTIONS,
    )

    print()
    print("=" * 70)
    print("DONE")
    print("=" * 70)

    print(
        final_temp.relative_to(
            BASE_DIR
        )
    )

    print(
        final_psal.relative_to(
            BASE_DIR
        )
    )


if __name__ == "__main__":
    main()