"""This module contains the main setup for the GEB model.

Notes:
- All prices are in nominal USD (face value) for their respective years. That means that the prices are not adjusted for inflation.
"""

import concurrent.futures
import inspect
import logging
import os
import time
from contextlib import contextmanager
from datetime import date, datetime
from pathlib import Path
from typing import Any, Callable, Iterator, Literal

import contextily as ctx
import geopandas as gpd
import matplotlib.pyplot as plt
import networkx
import numpy as np
import numpy.typing as npt
import pandas as pd
import pyflwdir
import rasterio.features
import xarray as xr
import yaml
import zarr
from affine import Affine
from hydromt.data_catalog import DataCatalog
from packaging.version import Version
from rasterio.env import defenv
from scipy.ndimage import binary_dilation
from shapely.geometry import Point, shape
from shapely.ops import unary_union

from geb import GEB_PACKAGE_DIR, __version__
from geb.build.data_catalog import NewDataCatalog
from geb.build.methods import build_method
from geb.build.version_updates import VERSION_UPDATES
from geb.workflows.io import (
    read_params,
    write_array,
    write_geom,
    write_params,
    write_table,
)
from geb.workflows.raster import clip_region, clip_with_grid, full_like, repeat_grid

from ..workflows.io import (
    read_zarr,
    write_zarr,
)
from .modules import (
    Agents,
    Crops,
    Forcing,
    GroundWater,
    Hydrography,
    LandSurface,
    Observations,
)
from .modules.hydrography import (
    create_river_raster_from_river_lines,
    extend_rivers_into_ocean,
)

# Set environment options for robustness
GDAL_HTTP_ENV_OPTS = {
    "GDAL_HTTP_MAX_RETRY": "10",  # Number of retry attempts
    "GDAL_HTTP_RETRY_DELAY": "2",  # Delay (seconds) between retries
    "GDAL_HTTP_TIMEOUT": "30",  # Timeout in seconds
    "GDAL_CACHEMAX": 1 * 1024**3,  # 1 GB cache size
    "GDAL_MAX_BAND_COUNT": "200000",  # Increase max band count
}
defenv(**GDAL_HTTP_ENV_OPTS)

XY_CHUNKSIZE = 3000  # chunksize for xy coordinates

os.environ["AWS_NO_SIGN_REQUEST"] = "YES"


@contextmanager
def suppress_logging_warning(logger: logging.Logger) -> Iterator[None]:
    """A context manager to suppress logging warning messages temporarily.

    Args:
        logger: The logger to suppress warnings for.

    Yields:
        None

    """
    current_level: int = logger.getEffectiveLevel()
    logger.setLevel(logging.ERROR)  # Set level to ERROR to suppress WARNING messages
    try:
        yield
    finally:
        logger.setLevel(current_level)  # Restore the original logging level


def boolean_mask_to_graph(
    mask: npt.NDArray[np.bool_], connectivity: int = 4, **kwargs: npt.NDArray[Any]
) -> networkx.Graph:
    """Convert a boolean mask to an undirected NetworkX graph.

    Additional attributes can be passed as keyword arguments, which
    will be added as attributes to the nodes of the graph.

    Args:
        mask (xarray.DataArray or numpy.ndarray):
            Boolean mask where True values are nodes in the graph
        connectivity:
            4 for von Neumann neighborhood (up, down, left, right)
            8 for Moore neighborhood (includes diagonals)
        **kwargs:
            Additional attributes to be added to the nodes of the graph.
            Must be 2D arrays with the same shape as the mask.

    Returns:
        networkx.Graph:
            Undirected graph where nodes are (y,x) coordinates of True cells

    Raises:
        ValueError: If connectivity is not 4 or 8.
    """
    # check dtypes
    assert isinstance(mask, (np.ndarray))
    for item in kwargs.values():
        assert isinstance(item, (np.ndarray))

    # Create an empty undirected graph
    G = networkx.Graph()

    # Get indices of True cells
    y_indices, x_indices = np.where(mask)

    # Add all True cells as nodes
    for y, x in zip(y_indices, x_indices, strict=True):
        node_attrs = {}
        for key, array in kwargs.items():
            node_attrs[key] = array[y, x].item()
        node_attrs["yx"] = (y, x)
        G.add_node((y, x), **node_attrs)

    # Define neighbor directions for 4-connectivity
    if connectivity == 4:
        directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]  # right, down, left, up
    elif connectivity == 8:  # 8-connectivity
        directions = [
            (0, 1),
            (1, 0),
            (0, -1),
            (-1, 0),  # right, down, left, up
            (1, 1),
            (1, -1),
            (-1, 1),
            (-1, -1),
        ]  # diagonals
    else:
        raise ValueError("Connectivity must be either 4 or 8.")

    # Connect neighboring cells
    for y, x in zip(y_indices, x_indices, strict=True):
        # Check each direction for a neighbor
        for dy, dx in directions:
            ny, nx = y + dy, x + dx

            # Check if neighbor is within bounds and is True
            if 0 <= ny < mask.shape[0] and 0 <= nx < mask.shape[1] and mask[ny, nx]:
                # Add an edge between current cell and neighbor
                G.add_edge((y, x), (ny, nx))

    return G


def get_river_graph(data_catalog: NewDataCatalog) -> networkx.DiGraph:
    """Create a directed graph for the river network.

    Args:
        data_catalog: Data catalog containing the MERIT basins.

    Returns:
        A directed graph where nodes are COMID values and edges point downstream.
    """
    # load river network data
    river_network = (
        data_catalog.fetch("merit_basins_rivers")
        .read(columns=["COMID", "NextDownID"])
        .set_index("COMID")
    )
    assert isinstance(river_network, pd.DataFrame)
    assert river_network.index.name == "COMID", (
        "The index of the river network is not the COMID column"
    )

    # create a directed graph for the river network
    river_graph: networkx.DiGraph = networkx.DiGraph()

    # add rivers with downstream connection
    river_network_with_downstream_connection = river_network[
        river_network["NextDownID"] != 0
    ]

    river_network_with_downstream_connection = (
        river_network_with_downstream_connection.itertuples(index=True, name=None)
    )

    river_graph.add_edges_from(river_network_with_downstream_connection)

    river_network_without_downstream_connection = river_network[
        river_network["NextDownID"] == 0
    ]

    river_graph.add_nodes_from(river_network_without_downstream_connection.index)

    return river_graph


def get_subbasin_id_from_coordinate(
    data_catalog: NewDataCatalog, lon: float, lat: float
) -> int:
    """Find the subbasin ID for a given coordinate.

    Args:
        data_catalog: Data catalog containing the MERIT basins.
        lon: Longitude of the point.
        lat: Latitude of the point.

    Returns:
        The COMID value for the subbasin containing the point.

    Raises:
        ValueError: If the point is not in a basin or in multiple basins.
    """
    # we select the basin that contains the point. To do so
    # we use a bounding box with the point coordinates, thus
    # xmin == xmax and ymin == ymax
    # geoparquet uses < and >, not <= and >=, so we need to add
    # a small value to the coordinates to avoid missing the point
    COMID = (
        data_catalog.fetch("merit_basins_catchments")
        .read(
            bbox=(lon - 10e-6, lat - 10e-6, lon + 10e-6, lat + 10e-6),
        )
        .set_index("COMID")
    )
    assert isinstance(COMID, gpd.GeoDataFrame)

    # get only the points where the point is inside the basin
    COMID = COMID[COMID.geometry.contains(Point(lon, lat))]

    if len(COMID) == 0:
        raise ValueError(
            f"The point is not in a basin. Note, that there are some holes in the MERIT basins dataset ({data_catalog.fetch('merit_basins_catchments').path}), ensure that the point is in a basin."
        )
    assert len(COMID) == 1, "The point is not in a single basin"
    # get the COMID value from the GeoDataFrame
    return COMID.index.values[0]


def get_sink_subbasin_id_for_geom(
    data_catalog: NewDataCatalog, geom: gpd.GeoDataFrame, river_graph: networkx.DiGraph
) -> list[int]:
    """Find all sink subbasins that intersect with the given geometry.

    This can be coastal basins, but also inland basins where the river
    flows out of the geometry. For example, if the geometry is a country
    boundary, all basins that intersect with the country boundary and
    have no downstream basin within the country are returned.

    Args:
        data_catalog: Data catalog containing the MERIT basins.
        geom: GeoDataFrame containing the geometry to find the sink subbasins for.
        river_graph: The river graph containing all subbasins and their connections.

    Returns:
        A list of COMID values for the sink subbasins.
    """
    subbasins = gpd.read_parquet(
        data_catalog.fetch("merit_basins_catchments").path,
        bbox=tuple([float(c) for c in geom.total_bounds]),
    ).set_index("COMID")

    subbasins = subbasins.iloc[
        subbasins.sindex.query(geom.union_all(), predicate="intersects")
    ]

    assert len(subbasins) > 0, "The geometry does not intersect with any subbasin."

    # create a subgraph containing only the selected subbasins
    region_river_graph = river_graph.subgraph(subbasins.index.tolist())

    # get all subbasins with no downstream subbasin (out degree is 0)
    # in the subgraph. These are the sink subbasins
    sink_nodes: list[int] = [
        COMID_ID
        for COMID_ID, out_degree in region_river_graph.out_degree(
            region_river_graph.nodes
        )
        if out_degree == 0
    ]

    return sink_nodes


def get_all_downstream_subbasins_in_geom(
    data_catalog: NewDataCatalog,
    geom: gpd.GeoDataFrame,
    logger: logging.Logger,
) -> list[int]:
    """Find all downstream subbasins (with NextDownID = 0) that intersect with the given geometry.

    Args:
        data_catalog: Data catalog containing the MERIT basins.
        geom: GeoDataFrame containing the geometry to find the downstream subbasins for.
        logger: Logger for progress tracking.

    Returns:
        A list of COMID values for all downstream subbasins in the geometry.
    """
    logger.info("Finding subbasins that intersect with geometry...")

    # Get all subbasins that intersect with the geometry
    subbasins = gpd.read_parquet(
        data_catalog.fetch("merit_basins_catchments").path,
        bbox=tuple([float(c) for c in geom.total_bounds]),
    ).set_index("COMID")

    subbasins = subbasins.iloc[
        subbasins.sindex.query(geom.union_all(), predicate="intersects")
    ]

    assert len(subbasins) > 0, "The geometry does not intersect with any subbasin."

    logger.info(f"Found {len(subbasins)} subbasins intersecting with geometry")
    logger.info("Loading river network data...")

    # Get river network data to find downstream basins (NextDownID = 0)
    river_network = (
        data_catalog.fetch("merit_basins_rivers")
        .read(columns=["COMID", "NextDownID", "uparea"])
        .set_index("COMID")
    )
    assert isinstance(river_network, pd.DataFrame)

    logger.info("Filtering for downstream subbasins (NextDownID = 0)...")

    # Filter for subbasins that are within the geometry and are downstream (NextDownID = 0)
    intersecting_subbasins = subbasins.index.intersection(river_network.index)
    downstream_subbasins = river_network.loc[intersecting_subbasins]
    downstream_subbasins = downstream_subbasins[downstream_subbasins["NextDownID"] == 0]

    logger.info(f"Found {len(downstream_subbasins)} downstream subbasins (outlets)")

    return downstream_subbasins.index.tolist()


def get_subbasin_upstream_areas(
    data_catalog: NewDataCatalog, subbasin_ids: list[int]
) -> dict[int, float]:
    """Get upstream areas for a list of subbasins with optimized batch loading.

    Args:
        data_catalog: Data catalog containing the MERIT basins.
        subbasin_ids: List of COMID values to get upstream areas for.

    Returns:
        Dictionary mapping COMID to upstream area in km2.
    """
    # Use filters to only read the rows we need - much faster than reading all data
    river_network = (
        data_catalog.fetch("merit_basins_rivers")
        .read(columns=["COMID", "uparea"], filters=[("COMID", "in", subbasin_ids)])
        .set_index("COMID")
    )
    assert isinstance(river_network, pd.DataFrame)

    # Convert to dict for faster lookup and handle missing values
    upstream_areas = river_network["uparea"].to_dict()

    # Fill in any missing subbasins with 0.0 area
    for subbasin_id in subbasin_ids:
        if subbasin_id not in upstream_areas:
            upstream_areas[subbasin_id] = 0.0

    return upstream_areas


def calculate_bbox_efficiency(
    outlet_ids: list[int],
    subbasin_to_idx: dict[int, int],
    coords_array: npt.NDArray[np.float64],
    upstream_areas: dict[int, float],
    river_graph: networkx.DiGraph,
    all_subbasin_coords: dict[int, tuple[float, float]],
    all_subbasin_areas: dict[int, float] | None = None,
) -> float:
    """Calculate the bounding box efficiency of a cluster for ERA5 data download optimization.

    This measures how efficiently the cluster's FULL WATERSHED fills its bounding box rectangle,
    considering only LAND AREA (not ocean/sea areas between coastal basins).

    Since ERA5 climate data is downloaded using the min/max lat/lon bounds of ALL subbasins
    (not just outlets), we need to consider the entire upstream watershed extent.

    Efficiency = (Total upstream area of cluster) / (Total land area within bbox)

    This land-only calculation prevents coastal clusters with ocean gaps (e.g., archipelagos,
    fjords, bays) from being unfairly penalized.

    Args:
        outlet_ids: List of outlet basin COMID values in the cluster.
        subbasin_to_idx: Mapping from COMID to row index in coords_array (for outlets only).
        coords_array: Nx2 array with centroid coordinates [lon, lat] for outlets.
        upstream_areas: Dictionary mapping COMID to upstream area in km2.
        river_graph: River graph for finding all upstream subbasins.
        all_subbasin_coords: Dictionary mapping ALL subbasin COMIDs to (lon, lat) tuples.
        all_subbasin_areas: Dictionary mapping ALL subbasin COMIDs to their areas in km2.
            If None, falls back to rectangular bbox area (old behavior).

    Returns:
        Efficiency ratio between 0.0 and 1.0 (higher is better for data download efficiency).
    """
    # Get ALL subbasins in this cluster (outlets + all upstream)
    all_cluster_subbasins = set()
    for outlet_id in outlet_ids:
        # Get all upstream basins for this outlet
        upstream = set(networkx.ancestors(river_graph, outlet_id))
        upstream.add(outlet_id)  # Include the outlet itself
        all_cluster_subbasins.update(upstream)

    # Get coordinates for ALL subbasins in the watershed
    cluster_lons = []
    cluster_lats = []
    for subbasin_id in all_cluster_subbasins:
        if subbasin_id in all_subbasin_coords:
            lon, lat = all_subbasin_coords[subbasin_id]
            cluster_lons.append(lon)
            cluster_lats.append(lat)

    if len(cluster_lons) == 0:
        return 1.0  # Edge case: no coordinates found

    # Calculate bounding box from ALL subbasin centroids
    lon_min, lon_max = min(cluster_lons), max(cluster_lons)
    lat_min, lat_max = min(cluster_lats), max(cluster_lats)

    # Calculate total upstream area of all outlets in cluster
    cluster_total_area_km2 = sum(upstream_areas.get(oid, 0.0) for oid in outlet_ids)

    # Calculate efficiency based on LAND AREA within bbox (not rectangular area)
    if all_subbasin_areas is not None:
        # NEW: Use sum of all subbasin areas within the bbox (land-only)
        # This excludes ocean/sea areas between coastal basins
        bbox_land_area_km2 = sum(
            all_subbasin_areas.get(subbasin_id, 0.0)
            for subbasin_id in all_cluster_subbasins
        )

        # Handle edge case: no valid area data
        if bbox_land_area_km2 < 1.0:
            return 1.0

        # Calculate efficiency ratio using land-only area
        efficiency = min(1.0, cluster_total_area_km2 / bbox_land_area_km2)
    else:
        # FALLBACK: Use rectangular bbox area (old behavior, includes ocean)
        # This is kept for backward compatibility if area data is not provided
        lat_span_km = (lat_max - lat_min) * 111.0
        mean_lat = (lat_min + lat_max) / 2.0
        lon_span_km = (lon_max - lon_min) * (
            111.320 * abs(np.cos(np.deg2rad(mean_lat)))
        )
        bbox_area_km2 = lat_span_km * lon_span_km

        # Handle edge case: very small bbox
        if bbox_area_km2 < 1.0:
            return 1.0

        efficiency = min(1.0, cluster_total_area_km2 / bbox_area_km2)

    return efficiency


def cluster_subbasins_following_coastline(
    data_catalog: NewDataCatalog,
    subbasin_ids: list[int],
    target_area_km2: float,
    logger: logging.Logger,
    river_graph: networkx.DiGraph,
    max_distance_km: float = 100.0,  # Maximum distance in km for cluster growth
    min_bbox_efficiency: float = 0.99,  # Minimum bounding box efficiency (0-1) for ERA5 download optimization
) -> list[list[int]]:
    """Cluster outlet basins by adding entire upstream watersheds of nearby outlets.

    This function creates clusters by:
    1. Starting from an outlet basin (NextDownID=0) - adds its entire upstream watershed
    2. Finding the nearest remaining outlet basin to the current cluster
    3. Adding that outlet's entire upstream watershed to the cluster
    4. Repeating until target area is reached, then starting a new cluster

    The clustering uses bounding box efficiency to prevent thin/diagonal clusters that would
    waste bandwidth when downloading ERA5 climate data (which uses rectangular bounding boxes).

    Args:
        data_catalog: Data catalog containing the MERIT basins.
        subbasin_ids: List of outlet basin COMID values (NextDownID=0) to cluster.
        target_area_km2: Target cumulative upstream area per cluster (e.g., 800,000 km²).
        logger: Logger for progress tracking.
        river_graph: River graph for finding upstream subbasins to calculate true bbox efficiency.
        max_distance_km: Maximum distance in km for adding outlets to a cluster (default: 100km).
        min_bbox_efficiency: Minimum ratio of cluster_area/bbox_area (default: 0.25 = 25%).
            Higher values create more compact clusters. Lower values allow more elongated shapes.

    Returns:
        List of clusters, where each cluster is a list of outlet basin COMID values.
    """
    logger.info(
        f"Clustering {len(subbasin_ids)} outlet basins by upstream watersheds..."
    )
    logger.info(f"Target area per cluster: {target_area_km2:,.0f} km²")
    logger.info(f"Maximum cluster growth distance: {max_distance_km:.0f} km")
    logger.info(
        f"Minimum bounding box efficiency: {min_bbox_efficiency:.1%} (for ERA5 download optimization)"
    )

    # Load outlet basin geometries
    subbasins = gpd.read_parquet(
        data_catalog.fetch("merit_basins_catchments").path,
        filters=[("COMID", "in", subbasin_ids)],
    ).set_index("COMID")

    logger.info("Getting upstream areas for all outlet basins...")
    upstream_areas = get_subbasin_upstream_areas(data_catalog, subbasin_ids)

    # Pre-compute centroids for distance calculations
    logger.info("Computing outlet basin centroids...")
    coords_array = np.zeros((len(subbasin_ids), 2))
    subbasin_to_idx = {}

    for i, outlet_id in enumerate(subbasin_ids):
        centroid = subbasins.loc[outlet_id].geometry.centroid
        coords_array[i] = [centroid.x, centroid.y]
        subbasin_to_idx[outlet_id] = i

    # Pre-load ALL subbasin coordinates for bbox efficiency calculation
    # We need this because ERA5 downloads cover the full watershed extent
    logger.info("Pre-loading ALL subbasin coordinates for bbox efficiency checks...")

    # Get all upstream subbasins for all outlets
    all_upstream_subbasins = set()
    for outlet_id in subbasin_ids:
        upstream = set(networkx.ancestors(river_graph, outlet_id))
        upstream.add(outlet_id)
        all_upstream_subbasins.update(upstream)

    logger.info(
        f"Loading centroids for {len(all_upstream_subbasins)} total subbasins (outlets + upstream)..."
    )

    # Load all subbasin geometries with area data for land-only bbox efficiency
    all_subbasins = gpd.read_parquet(
        data_catalog.fetch("merit_basins_catchments").path,
        filters=[("COMID", "in", list(all_upstream_subbasins))],
        columns=[
            "COMID",
            "geometry",
            "unitarea",
        ],  # Add unitarea for land-only calculation
    ).set_index("COMID")

    # Create coordinate dictionary and area dictionary for ALL subbasins
    all_subbasin_coords = {}
    all_subbasin_areas = {}
    for subbasin_id in all_upstream_subbasins:
        if subbasin_id in all_subbasins.index:
            centroid = all_subbasins.loc[subbasin_id].geometry.centroid
            all_subbasin_coords[subbasin_id] = (centroid.x, centroid.y)
            # unitarea is in km2 in MERIT Basins dataset
            all_subbasin_areas[subbasin_id] = all_subbasins.loc[subbasin_id].unitarea

    logger.info(
        f"Loaded coordinates and areas for {len(all_subbasin_coords)} subbasins"
    )
    logger.info("Using LAND-ONLY area for bbox efficiency (excludes ocean/sea gaps)")

    # Pre-compute distance matrix for fast lookups
    if len(subbasin_ids) < 1000:
        logger.info(f"Pre-computing distance matrix for {len(subbasin_ids)} outlets...")
        diffs = coords_array[:, np.newaxis, :] - coords_array[np.newaxis, :, :]
        distance_matrix = np.sqrt(np.sum(diffs**2, axis=2))
    else:
        logger.info("Computing distances on-demand (large dataset)")
        distance_matrix = None

    # Maximum distance threshold for cluster growth
    MAX_DISTANCE_DEGREES = max_distance_km / 111.0  # Convert km to degrees

    logger.info(
        f"Distance threshold: {MAX_DISTANCE_DEGREES:.2f}° (~{max_distance_km:.0f} km)"
    )
    logger.info("Strategy: Add nearest outlet basins (with full upstream watersheds)")

    clusters = []
    remaining_outlets = set(subbasin_ids)
    cluster_number = 1
    start_time = time.time()

    while remaining_outlets:
        # Progress tracking
        processed = len(subbasin_ids) - len(remaining_outlets)
        progress_pct = (processed / len(subbasin_ids)) * 100
        logger.info(
            f"\nCluster {cluster_number}: {len(remaining_outlets)} outlets remaining ({progress_pct:.1f}% complete)"
        )

        # Select starting outlet for this cluster
        if cluster_number == 1:
            # First cluster: start from top-right corner (highest lat, then easternmost lon)
            outlet_coords = [
                (oid, coords_array[subbasin_to_idx[oid]]) for oid in remaining_outlets
            ]
            outlet_coords.sort(key=lambda x: (-x[1][1], -x[1][0]))  # Sort by -lat, -lon
            start_outlet = outlet_coords[0][0]
            logger.info(f"  Starting outlet: {start_outlet} (top-right corner)")
        else:
            # Subsequent clusters: find nearest remaining outlet to current cluster
            # Use vectorized operations for speed
            cluster_indices = np.array(
                [subbasin_to_idx[oid] for oid in current_cluster]
            )
            remaining_list = list(remaining_outlets)
            remaining_indices = np.array(
                [subbasin_to_idx[oid] for oid in remaining_list]
            )

            if distance_matrix is not None:
                relevant_distances = distance_matrix[
                    np.ix_(cluster_indices, remaining_indices)
                ]
                min_distances = relevant_distances.min(axis=0)
            else:
                cluster_coords = coords_array[cluster_indices]
                remaining_coords = coords_array[remaining_indices]
                diffs = (
                    cluster_coords[:, np.newaxis, :]
                    - remaining_coords[np.newaxis, :, :]
                )
                distances = np.sqrt(np.sum(diffs**2, axis=2))
                min_distances = distances.min(axis=0)

            closest_idx = min_distances.argmin()
            min_dist = float(min_distances[closest_idx])
            start_outlet = remaining_list[closest_idx]

            logger.info(
                f"  Starting outlet: {start_outlet} (nearest to previous cluster, dist={min_dist:.3f}°)"
            )

        # Initialize cluster with starting outlet
        current_cluster = [start_outlet]
        current_area = upstream_areas.get(start_outlet, 0.0)
        remaining_outlets.remove(start_outlet)
        logger.info(f"  Initial area: {current_area:,.0f} km²")

        # Grow cluster by adding nearest outlets until target area reached
        while current_area < target_area_km2 and remaining_outlets:
            # Check if current cluster already has poor bounding box efficiency
            current_efficiency = calculate_bbox_efficiency(
                current_cluster,
                subbasin_to_idx,
                coords_array,
                upstream_areas,
                river_graph,
                all_subbasin_coords,
                all_subbasin_areas,
            )

            if current_efficiency < min_bbox_efficiency:
                logger.info(
                    f"  Current cluster bbox efficiency too low ({current_efficiency:.1%} < {min_bbox_efficiency:.1%}), stopping cluster"
                )
                break

            # Find nearest remaining outlet to ANY outlet in current cluster using vectorized ops
            # Get current cluster indices for distance calculations
            cluster_indices = np.array(
                [subbasin_to_idx[oid] for oid in current_cluster]
            )

            # Convert to lists for indexing
            remaining_list = list(remaining_outlets)
            remaining_indices = np.array(
                [subbasin_to_idx[oid] for oid in remaining_list]
            )

            if distance_matrix is not None:
                # Use pre-computed distance matrix (much faster!)
                # Extract relevant submatrix: cluster_outlets × remaining_outlets
                relevant_distances = distance_matrix[
                    np.ix_(cluster_indices, remaining_indices)
                ]
                # Find minimum distance from ANY cluster outlet to each remaining outlet
                min_distances_to_remaining = relevant_distances.min(axis=0)
            else:
                # On-demand vectorized calculation
                cluster_coords = coords_array[cluster_indices]
                remaining_coords = coords_array[remaining_indices]

                # Broadcasting: (n_cluster, 1, 2) - (1, n_remaining, 2) = (n_cluster, n_remaining, 2)
                diffs = (
                    cluster_coords[:, np.newaxis, :]
                    - remaining_coords[np.newaxis, :, :]
                )
                distances = np.sqrt(np.sum(diffs**2, axis=2))
                # Min distance from ANY cluster outlet to each remaining outlet
                min_distances_to_remaining = distances.min(axis=0)

            # Find the outlet with minimum distance
            closest_idx = min_distances_to_remaining.argmin()
            min_dist = float(min_distances_to_remaining[closest_idx])
            nearest_outlet = remaining_list[closest_idx]

            # Check distance threshold
            if min_dist > MAX_DISTANCE_DEGREES:
                logger.info(
                    f"  Nearest outlet is {min_dist:.3f}° away (>{MAX_DISTANCE_DEGREES:.3f}°), stopping cluster"
                )
                break

            # Check bounding box efficiency if we add this outlet
            # This prevents thin/diagonal clusters that waste ERA5 download bandwidth
            all_outlets_with_candidate = current_cluster + [nearest_outlet]
            candidate_efficiency = calculate_bbox_efficiency(
                all_outlets_with_candidate,
                subbasin_to_idx,
                coords_array,
                upstream_areas,
                river_graph,
                all_subbasin_coords,
                all_subbasin_areas,
            )

            if candidate_efficiency < min_bbox_efficiency:
                logger.info(
                    f"  Adding outlet {nearest_outlet} would reduce bbox efficiency to {candidate_efficiency:.1%} (< {min_bbox_efficiency:.1%}), stopping cluster"
                )
                break

            # Add the nearest outlet to cluster
            outlet_area = upstream_areas.get(nearest_outlet, 0.0)
            current_cluster.append(nearest_outlet)
            current_area += outlet_area
            remaining_outlets.remove(nearest_outlet)

            # Log progress occasionally
            if len(current_cluster) <= 10 or len(current_cluster) % 5 == 0:
                logger.info(
                    f"  Added outlet {nearest_outlet} (area: {outlet_area:,.0f} km², dist: {min_dist:.3f}°) - Total: {current_area:,.0f} km² [{len(current_cluster)} outlets]"
                )

        # Cluster complete
        clusters.append(current_cluster)
        logger.info(
            f"Completed cluster {cluster_number}: {len(current_cluster)} outlets, {current_area:,.0f} km²"
        )
        cluster_number += 1

    total_time = time.time() - start_time
    logger.info(
        f"\nClustering complete! Created {len(clusters)} clusters in {total_time:.1f}s ({total_time / 60:.1f} min)"
    )

    # Post-processing configuration
    MERGE_DISTANCE_KM = 100.0
    SMALL_CLUSTER_THRESHOLD_PERCENT = 0.10  # 10% of target area

    logger.info("\n=== POST-PROCESSING ===")
    logger.info(f"Merge distance: {MERGE_DISTANCE_KM:.0f} km")
    logger.info(
        f"Small cluster threshold: {SMALL_CLUSTER_THRESHOLD_PERCENT:.0%} of target area ({SMALL_CLUSTER_THRESHOLD_PERCENT * target_area_km2:,.0f} km²)"
    )

    # Post-processing 1: Merge any clusters <10% of target area with nearest neighbor
    logger.info("\n[1/2] Merging small clusters (<10% of target area)...")

    # Calculate cluster areas
    cluster_areas = {}
    small_clusters = []  # Clusters < 10% of target area
    SMALL_CLUSTER_THRESHOLD_KM2 = SMALL_CLUSTER_THRESHOLD_PERCENT * target_area_km2

    for cluster_idx, cluster in enumerate(clusters):
        cluster_area = sum(upstream_areas.get(outlet_id, 0.0) for outlet_id in cluster)
        cluster_areas[cluster_idx] = cluster_area

        if cluster_area < SMALL_CLUSTER_THRESHOLD_KM2:
            small_clusters.append(cluster_idx)

    logger.info(
        f"  Found {len(small_clusters)} small clusters (<{SMALL_CLUSTER_THRESHOLD_KM2:,.0f} km²)"
    )

    if small_clusters:
        MERGE_DISTANCE_DEGREES = MERGE_DISTANCE_KM / 111.0
        clusters_to_remove = set()

        for small_idx in small_clusters:
            # Skip if this cluster was already merged into another cluster
            if small_idx in clusters_to_remove:
                continue

            small_cluster = clusters[small_idx]
            small_area = cluster_areas[small_idx]

            # Get coordinates of small cluster outlets
            small_indices = np.array([subbasin_to_idx[oid] for oid in small_cluster])
            small_coords = coords_array[small_indices]

            # Find nearest cluster using minimum distance between ANY subbasins (not centroids!)
            # This ensures we find truly adjacent clusters even if centroids are far apart
            min_distance = float("inf")
            nearest_cluster_idx = None

            for other_idx, other_cluster in enumerate(clusters):
                if other_idx == small_idx or other_idx in clusters_to_remove:
                    continue

                other_indices = np.array(
                    [subbasin_to_idx[oid] for oid in other_cluster]
                )
                other_coords = coords_array[other_indices]

                # Calculate minimum distance between ANY outlet in small cluster and ANY outlet in target cluster
                # Broadcasting: (n_small, 1, 2) - (1, n_other, 2) = (n_small, n_other, 2)
                diffs = small_coords[:, np.newaxis, :] - other_coords[np.newaxis, :, :]
                distances = np.sqrt(np.sum(diffs**2, axis=2))
                min_dist = float(distances.min())  # Closest subbasin pair distance

                if min_dist < min_distance:
                    min_distance = min_dist
                    nearest_cluster_idx = other_idx

            # Simple rule: if cluster is <10% of target, ALWAYS merge with nearest neighbor
            # No checks - small clusters must be merged to avoid inefficient tiny runs
            if nearest_cluster_idx is not None:
                target_area = cluster_areas[nearest_cluster_idx]
                merged_area = target_area + small_area

                # Force merge - log the result
                logger.info(
                    f"    Merging small cluster {small_idx} (area: {small_area:,.0f} km², {small_area / target_area_km2:.1%} of target) "
                    f"into cluster {nearest_cluster_idx} (distance: {min_distance * 111.0:.1f} km, merged area: {merged_area:,.0f} km²)"
                )

                # Update both the cluster list AND the cluster_areas dictionary
                clusters[nearest_cluster_idx].extend(small_cluster)
                cluster_areas[nearest_cluster_idx] = merged_area
                clusters_to_remove.add(small_idx)
            else:
                logger.warning(
                    f"    WARNING: Cannot find merge target for cluster {small_idx} (area: {small_area:,.0f} km²) - this should never happen!"
                )

        # Remove merged small clusters
        if clusters_to_remove:
            logger.info(
                f"  Merged and removing {len(clusters_to_remove)} small clusters..."
            )
            clusters = [
                cluster
                for idx, cluster in enumerate(clusters)
                if idx not in clusters_to_remove
            ]
            logger.info(f"  Cluster count after small cluster merging: {len(clusters)}")
    else:
        logger.info("  No small clusters found")

    # Post-processing 2: Assign any remaining unclustered basins
    logger.info("\n[2/2] Checking for unclustered outlet basins...")

    all_clustered_outlets = set()
    for cluster in clusters:
        all_clustered_outlets.update(cluster)

    unclustered_outlets = set(subbasin_ids) - all_clustered_outlets

    if unclustered_outlets:
        logger.info(f"  Found {len(unclustered_outlets)} unclustered outlet basins")

        for unclustered_outlet in unclustered_outlets:
            unclustered_idx = subbasin_to_idx[unclustered_outlet]
            unclustered_coords = coords_array[unclustered_idx]

            # Find nearest cluster
            min_distance = float("inf")
            nearest_cluster_idx = 0

            for cluster_idx, cluster in enumerate(clusters):
                cluster_indices = np.array([subbasin_to_idx[oid] for oid in cluster])
                cluster_coords = coords_array[cluster_indices]

                diffs = cluster_coords - unclustered_coords
                distances = np.sqrt(np.sum(diffs * diffs, axis=1))
                min_dist_to_cluster = float(distances.min())

                if min_dist_to_cluster < min_distance:
                    min_distance = min_dist_to_cluster
                    nearest_cluster_idx = cluster_idx

            clusters[nearest_cluster_idx].append(unclustered_outlet)
            logger.info(
                f"    Assigned outlet {unclustered_outlet} to cluster {nearest_cluster_idx} "
                f"(distance: {min_distance * 111.0:.1f} km, area: {upstream_areas.get(unclustered_outlet, 0.0):,.0f} km²)"
            )

        logger.info(
            f"  Successfully assigned all {len(unclustered_outlets)} unclustered basins"
        )
    else:
        logger.info("  No unclustered basins found - all outlets assigned")

    logger.info(f"\n=== FINAL RESULT: {len(clusters)} clusters ===")

    return clusters


def save_clusters_to_geoparquet(
    clusters: list[list[int]],
    data_catalog: NewDataCatalog,
    output_path: Path,
    cluster_prefix: str = "cluster",
) -> None:
    """Save clusters to a geoparquet file with cluster IDs.

    Args:
        clusters: List of clusters, where each cluster is a list of COMID values.
        data_catalog: Data catalog containing the MERIT basins.
        output_path: Path where to save the geoparquet file.
        cluster_prefix: Prefix for cluster names.
    """
    print(f"Saving clusters to geoparquet: {output_path}")

    # Get all subbasin IDs using list comprehension for better performance
    all_subbasin_ids = [sid for cluster in clusters for sid in cluster]

    # Load subbasin geometries and upstream areas in parallel
    print(f"Loading geometries for {len(all_subbasin_ids)} subbasins...")
    subbasins = gpd.read_parquet(
        data_catalog.fetch("merit_basins_catchments").path,
        filters=[("COMID", "in", all_subbasin_ids)],
        columns=["COMID", "geometry"],  # Only load needed columns
    ).set_index("COMID")

    # Simplify geometries for faster QGIS rendering
    print("Simplifying geometries for QGIS performance...")
    simplify_tolerance = 0.001  # ~111m at equator
    subbasins["geometry"] = subbasins["geometry"].simplify(
        simplify_tolerance, preserve_topology=True
    )

    # Get upstream areas for display - this is already optimized with filters
    print("Loading upstream areas...")
    upstream_areas = get_subbasin_upstream_areas(data_catalog, all_subbasin_ids)

    # Pre-calculate cluster areas for better performance
    print("Pre-calculating cluster areas...")
    cluster_areas = {}
    for cluster_idx, cluster_subbasins in enumerate(clusters):
        cluster_id = f"{cluster_prefix}_{cluster_idx:03d}"
        cluster_areas[cluster_idx] = sum(
            upstream_areas.get(sid, 0) for sid in cluster_subbasins
        )

    # Create cluster assignments using vectorized operations
    print("Creating cluster assignments...")
    cluster_data = []
    for cluster_idx, cluster_subbasins in enumerate(clusters):
        cluster_id = f"{cluster_prefix}_{cluster_idx:03d}"
        cluster_area_km2 = cluster_areas[cluster_idx]

        # Create all records for this cluster at once
        for subbasin_id in cluster_subbasins:
            cluster_data.append(
                {
                    "COMID": subbasin_id,
                    "cluster_id": cluster_id,
                    "cluster_number": cluster_idx,
                    "cluster_area_km2": cluster_area_km2,
                    "subbasin_area_km2": upstream_areas.get(subbasin_id, 0),
                    "geometry": subbasins.loc[subbasin_id, "geometry"],
                }
            )

    # Create GeoDataFrame
    print("Creating GeoDataFrame...")
    cluster_gdf = gpd.GeoDataFrame(cluster_data, crs=subbasins.crs)

    # Save to geoparquet
    print("Saving to geoparquet...")
    write_geom(cluster_gdf, output_path)
    print(
        f"Saved {len(cluster_data)} subbasins in {len(clusters)} clusters to {output_path}"
    )


def create_cluster_visualization_map(
    clusters: list[list[int]],
    data_catalog: NewDataCatalog,
    river_graph: networkx.DiGraph,
    output_path: str | Path,
    cluster_prefix: str = "cluster",
    figsize: tuple[int, int] = (20, 16),
) -> None:
    """Create a visualization map showing dissolved cluster boundaries with distinct colors.

    This creates a clean map where each cluster is shown as a single merged polygon
    with distinct colors and labeled numbers for easy identification.

    Args:
        clusters: List of clusters, where each cluster is a list of downstream COMID values.
        data_catalog: Data catalog containing the MERIT basins.
        river_graph: River graph for finding upstream subbasins.
        output_path: Path where to save the PNG file.
        cluster_prefix: Prefix for cluster names.
        figsize: Figure size in inches (width, height).
    """
    from matplotlib.patches import Patch

    print(f"Creating cluster visualization map: {output_path}")

    # Get upstream subbasins for each cluster (with caching for speed)
    print("Computing upstream subbasins for all clusters (with caching)...")

    # Build a cache of upstream subbasins to avoid redundant networkx.ancestors calls
    upstream_cache = {}

    def get_all_upstream_cached(node_id: int) -> set[int]:
        """Get all upstream subbasins with caching.

        Args:
            node_id: The COMID to get upstream subbasins for.

        Returns:
            Set of all upstream subbasin COMIDs including the node itself.
        """
        if node_id not in upstream_cache:
            upstream = set(networkx.ancestors(river_graph, node_id))
            upstream.add(node_id)
            upstream_cache[node_id] = upstream
        return upstream_cache[node_id]

    all_cluster_subbasins = []
    for cluster_idx, downstream_subbasins in enumerate(clusters):
        cluster_all_subbasins = set()
        for downstream_subbasin in downstream_subbasins:
            cluster_all_subbasins.update(get_all_upstream_cached(downstream_subbasin))
        all_cluster_subbasins.append(list(cluster_all_subbasins))

    total_subbasin_count = sum(len(c) for c in all_cluster_subbasins)
    print(f"Loading geometries for {total_subbasin_count} unique subbasins...")
    all_geometries = gpd.read_parquet(
        data_catalog.fetch("merit_basins_catchments").path,
        filters=[("COMID", "in", [sid for c in all_cluster_subbasins for sid in c])],
        columns=["COMID", "geometry"],
    ).set_index("COMID")

    # simplify geometries before plotting
    print(
        "Pre-simplifying geometries for faster dissolving (tolerance: 0.01° ≈ 1km)..."
    )
    all_geometries["geometry"] = all_geometries["geometry"].simplify(
        0.01, preserve_topology=True
    )
    print("Geometry complexity reduced significantly")

    # Dissolve clusters sequentially to create merged boundaries for visualization
    print("Dissolving clusters into merged boundaries...")

    # Use a colormap with very distinct colors
    colors = plt.cm.tab20(np.linspace(0, 1, min(len(clusters), 20)))  # type: ignore[attr-defined]
    if len(clusters) > 20:
        colors = plt.cm.hsv(np.linspace(0, 1, len(clusters)))  # type: ignore[attr-defined]

    # Process clusters sequentially
    merged_cluster_data = []
    for cluster_idx, subbasins_to_merge in enumerate(all_cluster_subbasins):
        cluster_id = f"{cluster_prefix}_{cluster_idx:03d}"
        color = colors[cluster_idx]

        print(
            f"  Dissolving cluster {cluster_idx}: {len(subbasins_to_merge)} subbasins"
        )

        # Get geometries for this cluster
        cluster_geometries = all_geometries.loc[subbasins_to_merge]

        # Optimized dissolving for large clusters: use chunked unary_union
        if len(subbasins_to_merge) > 10000:
            print(f"    Large cluster detected - using chunked dissolving...")
            # Process in chunks of 5000 geometries
            chunk_size = 5000
            geometries = cluster_geometries.geometry.values
            chunks = [
                geometries[i : i + chunk_size]
                for i in range(0, len(geometries), chunk_size)
            ]

            # First pass: dissolve each chunk
            print(f"    Dissolving {len(chunks)} chunks...")
            chunk_unions = []
            for i, chunk in enumerate(chunks):
                if (i + 1) % 5 == 0 or (i + 1) == len(chunks):
                    print(f"      Chunk {i + 1}/{len(chunks)}...")
                chunk_union = unary_union(chunk)
                chunk_unions.append(chunk_union)

            # Second pass: dissolve the chunk results
            print(f"    Merging {len(chunk_unions)} chunk results...")
            merged_geometry = unary_union(chunk_unions)
        else:
            # Small cluster: direct dissolve
            merged_geometry = unary_union(cluster_geometries.geometry.values)

        # Clean up internal boundaries
        merged_geometry = merged_geometry.buffer(0)

        # Simplify for cleaner visualization
        merged_geometry = merged_geometry.simplify(0.01, preserve_topology=True)

        merged_cluster_data.append(
            {
                "cluster_id": cluster_id,
                "cluster_number": cluster_idx,
                "color": color,
                "geometry": merged_geometry,
            }
        )

    # Create GeoDataFrame with dissolved cluster boundaries
    cluster_gdf = gpd.GeoDataFrame(merged_cluster_data, crs="EPSG:4326")

    # Convert to Web Mercator for contextily
    cluster_gdf_mercator = cluster_gdf.to_crs(epsg=3857)

    # Create the plot
    fig, ax = plt.subplots(1, 1, figsize=figsize)

    # Plot each cluster with distinct colors and thick boundaries
    print("Plotting clusters...")
    for idx, row in cluster_gdf_mercator.iterrows():
        gpd.GeoSeries([row.geometry], crs=cluster_gdf_mercator.crs).plot(
            ax=ax,
            color=row.color,
            alpha=0.6,
            edgecolor="black",
            linewidth=2.5,
        )

        # Add cluster number labels at centroids
        if row.geometry.geom_type == "MultiPolygon":
            centroid = row.geometry.centroid
        else:
            centroid = row.geometry.centroid

        ax.annotate(
            f"{row.cluster_number}",
            (centroid.x, centroid.y),
            fontsize=20,
            fontweight="bold",
            ha="center",
            va="center",
            color="white",
            bbox=dict(
                boxstyle="round,pad=0.5",
                facecolor="black",
                alpha=0.8,
                edgecolor="white",
                linewidth=2,
            ),
        )

    # Add background map AFTER plotting data (so axis limits are set)
    print("Adding background map...")
    ctx.add_basemap(
        ax,
        crs=cluster_gdf_mercator.crs,
        source=ctx.providers.CartoDB.Positron,  # Light, clean background # ty:ignore[unresolved-attribute]
        alpha=0.5,
        zoom="auto",
    )
    print("  Added CartoDB Positron background")
    ax.set_title(
        f"Europe Basin Clusters - {len(clusters)} Clusters",
        fontsize=22,
        fontweight="bold",
        pad=20,
    )
    ax.set_xlabel("Longitude", fontsize=16)
    ax.set_ylabel("Latitude", fontsize=16)

    # Add legend with cluster information
    legend_elements = [
        Patch(
            facecolor=colors[i],
            alpha=0.6,
            edgecolor="black",
            linewidth=2,
            label=f"Cluster {i}",
        )
        for i in range(len(clusters))
    ]
    legend = ax.legend(
        handles=legend_elements,
        loc="upper left",
        bbox_to_anchor=(1.02, 1),
        fontsize=12,
        framealpha=0.9,
        title="Clusters",
        title_fontsize=14,
    )

    # Remove axes ticks for cleaner look
    ax.set_xticks([])
    ax.set_yticks([])

    # Remove spines
    for spine in ax.spines.values():
        spine.set_visible(False)

    # Tight layout
    plt.tight_layout()

    # Save the figure
    print(f"Saving map to {output_path}...")
    plt.savefig(output_path, dpi=300, bbox_inches="tight", facecolor="white")
    plt.close()

    print(f"Saved cluster visualization map to {output_path}")


def create_multi_basin_configs(
    clusters: list[list[int]],
    working_directory: Path,
    cluster_prefix: str = "cluster",
    data_catalog: NewDataCatalog | None = None,
    river_graph: networkx.DiGraph | None = None,
) -> list[Path]:
    """Create separate config files and directories for each cluster of subbasins.

    Args:
        clusters: List of clusters, where each cluster is a list of COMID values.
        working_directory: Working directory for the models (init_multiple_dir).
        cluster_prefix: Prefix for cluster directory names.
        data_catalog: Data catalog for calculating basin areas (optional).
        river_graph: River graph for finding upstream subbasins (optional).

    Returns:
        List of paths to created cluster directories.
    """
    print(f"Creating configuration files for {len(clusters)} clusters...")

    # Create full build.yml in init_multiple_dir directory
    print("Creating build.yml in main init multiple directory...")
    build_config_path = working_directory / "build.yml"
    # Read build config from geul example and automatically copy it
    geul_build_path = GEB_PACKAGE_DIR / "examples" / "geul" / "build.yml"

    print(f"Reading build configuration from: {geul_build_path}")

    # Copy geul build.yml content directly to init_multiple_dir build.yml
    with open(geul_build_path, "r") as src, open(build_config_path, "w") as dst:
        dst.write(src.read())

    print(f"Created build.yml in {working_directory}")

    # Create model.yml in init_multiple_dir directory that inherits from reasonable default
    print("Creating model.yml in init_multiple_dir directory...")
    model_config_path = working_directory / "model.yml"

    # Define the model configuration content with inheritance from reasonable default
    model_config_content = """inherits: "{GEB_PACKAGE_DIR}/reasonable_default_config.yml"
"""

    with open(model_config_path, "w") as f:
        f.write(model_config_content)

    print(f"Created model.yml in {working_directory}")

    cluster_directories = []

    # Pre-calculate all basin areas at once if needed
    basin_areas = {}
    if data_catalog is not None and river_graph is not None:
        print("Pre-calculating basin areas for all clusters...")
        from functools import partial

        # Create a partial function for area calculation
        calculate_area_fn = partial(
            calculate_cluster_basin_area,
            data_catalog=data_catalog,
            river_graph=river_graph,
        )

        # Use ThreadPoolExecutor for I/O-bound parquet reading
        with concurrent.futures.ThreadPoolExecutor(
            max_workers=min(8, len(clusters))
        ) as executor:
            # Submit all area calculation tasks
            future_to_cluster = {
                executor.submit(calculate_area_fn, cluster): i
                for i, cluster in enumerate(clusters)
            }

            # Collect results as they complete
            for future in concurrent.futures.as_completed(future_to_cluster):
                cluster_idx = future_to_cluster[future]
                try:
                    basin_areas[cluster_idx] = future.result()
                except Exception as exc:
                    print(f"Cluster {cluster_idx} area calculation failed: {exc}")
                    basin_areas[cluster_idx] = None

    # Create directories and config files
    print("Creating cluster directories and configuration files...")
    for i, cluster in enumerate(clusters):
        print(
            f"Creating cluster {i + 1}/{len(clusters)}: {cluster_prefix}_{i:03d} ({len(cluster)} subbasins)"
        )

        # Create cluster directory and base subdirectory
        cluster_dir = working_directory / f"{cluster_prefix}_{i:03d}"
        base_dir = cluster_dir / "base"
        base_dir.mkdir(parents=True, exist_ok=True)
        cluster_directories.append(cluster_dir)

        # Create build.yml with inheritance in base folder (inherit from init_multiple_dir build.yml)
        build_config_path = base_dir / "build.yml"

        # Create the build configuration dictionary
        build_config = {"inherits": "../../build.yml"}

        with open(build_config_path, "w") as f:
            yaml.dump(build_config, f, default_flow_style=False, sort_keys=False)

        # Create model.yml with inheritance and cluster-specific subbasins in base folder
        model_config_path = base_dir / "model.yml"

        # Convert all cluster values to regular Python integers
        cluster_ints = [int(subbasin_id) for subbasin_id in cluster]

        # Use pre-calculated basin area
        total_basin_area_km2 = basin_areas.get(i)
        if total_basin_area_km2 is not None:
            print(f"  Basin area: {total_basin_area_km2:,.0f} km²")

        # Create the configuration dictionary
        model_config = {
            "inherits": "../../model.yml",
            "general": {"region": {"subbasin": cluster_ints}},
        }

        # Add basin area if calculated
        if total_basin_area_km2 is not None:
            model_config["basin"] = {"total_area_km2": round(total_basin_area_km2, 2)}

        with open(model_config_path, "w") as f:
            yaml.dump(model_config, f, default_flow_style=False, sort_keys=False)

        print(
            f"  Created configuration files in {base_dir.relative_to(working_directory)}"
        )

    print(
        f"Successfully created {len(clusters)} cluster configurations in {working_directory}"
    )

    return cluster_directories


def calculate_cluster_basin_area(
    cluster: list[int],
    data_catalog: DataCatalog,
    river_graph: networkx.DiGraph,
) -> float:
    """Calculate the total basin area for a cluster of downstream subbasins.

    Args:
        cluster: List of downstream COMID values for the cluster.
        data_catalog: Data catalog containing the MERIT basins.
        river_graph: River graph for finding upstream subbasins.

    Returns:
        Total basin area in km².
    """
    # Find all upstream subbasins for this cluster using set operations for efficiency
    cluster_all_subbasins = set()
    for downstream_subbasin in cluster:
        upstream_subbasins = set(networkx.ancestors(river_graph, downstream_subbasin))
        upstream_subbasins.add(downstream_subbasin)
        cluster_all_subbasins.update(upstream_subbasins)

    subbasins_to_merge = list(cluster_all_subbasins)

    # Load geometries for subbasins - only load geometry and area columns for performance
    cluster_geometries = gpd.read_parquet(
        data_catalog.fetch("merit_basins_catchments").path,  # type: ignore[attr-defined]
        filters=[("COMID", "in", subbasins_to_merge)],
        columns=["COMID", "geometry"],  # Only load required columns
    ).set_index("COMID")

    # Union all geometries efficiently using cascaded union
    merged_geometry = cluster_geometries.geometry.union_all()

    # Calculate area using equal area projection (EPSG:6933) with optimized GeoSeries creation
    gdf_temp = gpd.GeoDataFrame([{"geometry": merged_geometry}], crs="EPSG:4326")
    total_area_m2 = float(gdf_temp.to_crs("EPSG:6933").geometry.area.iloc[0])

    return float(total_area_m2 / 1e6)  # Convert m² to km² and ensure Python float


def save_clusters_as_merged_geometries(
    clusters: list[list[int]],
    data_catalog: NewDataCatalog,
    river_graph: networkx.DiGraph,
    output_path: Path,
    cluster_prefix: str = "cluster",
    pre_simplify: bool = True,
    buffer_distance_km: float = 5.0,
) -> None:
    """Save clusters as merged geometries - one polygon outline per cluster.

    This creates the most compact representation possible: each cluster becomes a single
    dissolved polygon geometry representing the entire basin area. For clusters with many
    disconnected pieces (islands, archipelagos), a buffer-dissolve-unbuffer operation is
    applied to merge nearby polygons into single features for better performance.

    Args:
        clusters: List of clusters, where each cluster is a list of downstream COMID values.
        data_catalog: Data catalog containing the MERIT basins.
        river_graph: River graph for finding upstream subbasins.
        output_path: Path where to save the geoparquet file.
        cluster_prefix: Prefix for cluster names.
        pre_simplify: If True, simplify geometries before dissolving (much faster). Default: True.
        buffer_distance_km: Buffer distance in km to merge nearby polygons. Default: 5.0 km.
    """
    print(f"Creating merged cluster geometries: {output_path}")
    if pre_simplify:
        print("  Using optimized mode: simplify before dissolve (much faster)")
    print(f"  Buffer distance for merging polygons: {buffer_distance_km} km")

    # Pre-compute all upstream subbasins for all clusters using caching
    print("Pre-computing upstream subbasins for all clusters (with caching)...")

    # Build a cache of upstream subbasins to avoid redundant networkx.ancestors calls
    upstream_cache = {}

    def get_all_upstream_cached(node_id: int) -> set[int]:
        """Get all upstream subbasins with caching.

        Args:
            node_id: The COMID to get upstream subbasins for.

        Returns:
            Set of all upstream subbasin COMIDs including the node itself.
        """
        if node_id not in upstream_cache:
            upstream = set(networkx.ancestors(river_graph, node_id))
            upstream.add(node_id)
            upstream_cache[node_id] = upstream
        return upstream_cache[node_id]

    all_cluster_subbasins = []
    for cluster_idx, downstream_subbasins in enumerate(clusters):
        cluster_all_subbasins = set()
        for downstream_subbasin in downstream_subbasins:
            cluster_all_subbasins.update(get_all_upstream_cached(downstream_subbasin))
        all_cluster_subbasins.append(list(cluster_all_subbasins))
        print(f"  Cluster {cluster_idx}: {len(cluster_all_subbasins)} subbasins")

    total_subbasin_count = sum(len(c) for c in all_cluster_subbasins)
    print(
        f"Loading geometries for {total_subbasin_count} unique subbasins (single batch read)..."
    )
    all_geometries = gpd.read_parquet(
        data_catalog.fetch("merit_basins_catchments").path,
        filters=[("COMID", "in", [sid for c in all_cluster_subbasins for sid in c])],
        columns=["COMID", "geometry"],  # Only load what we need
    ).set_index("COMID")

    # Optimization: Simplify geometries BEFORE dissolving (much faster)
    if pre_simplify:
        simplify_tolerance = 0.01  # ~1km at equator
        print(
            f"Pre-simplifying {len(all_geometries)} geometries (tolerance: {simplify_tolerance}° ≈ 1km)..."
        )
        all_geometries["geometry"] = all_geometries["geometry"].simplify(
            simplify_tolerance, preserve_topology=True
        )
        print(f"  Geometry complexity reduced significantly for faster dissolving")

    # Process each cluster sequentially (optimized with pre-loaded geometries)
    print(f"Processing {len(clusters)} clusters...")
    merged_cluster_data = []

    for cluster_idx, (downstream_subbasins, subbasins_to_merge) in enumerate(
        zip(clusters, all_cluster_subbasins)
    ):
        cluster_id = f"{cluster_prefix}_{cluster_idx:03d}"
        print(
            f"  Processing cluster {cluster_idx + 1}/{len(clusters)}: {cluster_id} ({len(subbasins_to_merge)} subbasins)"
        )

        # Extract geometries for this cluster from pre-loaded data
        cluster_geometries = all_geometries.loc[subbasins_to_merge]

        # Optimized dissolving for large clusters: use chunked unary_union
        if len(subbasins_to_merge) > 10000:
            print(f"    Large cluster detected - using chunked dissolving...")
            # Process in chunks of 5000 geometries
            chunk_size = 5000
            geometries = cluster_geometries.geometry.values
            chunks = [
                geometries[i : i + chunk_size]
                for i in range(0, len(geometries), chunk_size)
            ]

            # First pass: dissolve each chunk
            print(f"    Dissolving {len(chunks)} chunks...")
            chunk_unions = []
            for i, chunk in enumerate(chunks):
                if (i + 1) % 5 == 0 or (i + 1) == len(chunks):
                    print(f"      Chunk {i + 1}/{len(chunks)}...")
                chunk_union = unary_union(chunk)
                chunk_unions.append(chunk_union)

            # Second pass: dissolve the chunk results
            print(f"    Merging {len(chunk_unions)} chunk results...")
            merged_geometry = unary_union(chunk_unions)
        else:
            # Small cluster: direct dissolve
            merged_geometry = unary_union(cluster_geometries.geometry.values)

        # Clean up internal boundaries
        merged_geometry = merged_geometry.buffer(0)

        # For MultiPolygons with many pieces, use buffer-dissolve-unbuffer to merge nearby islands
        # This significantly reduces polygon count for coastal clusters with archipelagos
        if hasattr(merged_geometry, "geoms") and len(merged_geometry.geoms) > 10:
            print(
                f"    Cluster has {len(merged_geometry.geoms)} separate polygons - applying buffer merge..."
            )

            # Convert buffer distance from km to degrees (approximate at mid-latitude)
            # At 45° latitude: 1° longitude ≈ 78 km, 1° latitude ≈ 111 km
            # Use conservative estimate: 1° ≈ 100 km
            buffer_degrees = buffer_distance_km / 100.0

            # Buffer → Dissolve → Negative buffer
            # This merges polygons within buffer_distance of each other
            buffered = merged_geometry.buffer(buffer_degrees)
            dissolved = (
                unary_union([buffered]) if hasattr(buffered, "__iter__") else buffered
            )
            merged_geometry = dissolved.buffer(-buffer_degrees)

            # Count final polygons
            final_count = (
                len(merged_geometry.geoms) if hasattr(merged_geometry, "geoms") else 1
            )
            print(f"    Reduced to {final_count} polygons after buffer merge")

        # Apply post-dissolve simplification for smoother rendering
        post_simplify_tolerance = 0.01  # ~1km at equator
        merged_geometry = merged_geometry.simplify(
            post_simplify_tolerance, preserve_topology=True
        )

        # Fast area calculation using equal area projection
        gdf_temp = gpd.GeoDataFrame([{"geometry": merged_geometry}], crs="EPSG:4326")
        total_area_m2 = float(gdf_temp.to_crs("EPSG:6933").geometry.area.iloc[0])
        total_area_km2 = float(total_area_m2 / 1e6)

        merged_cluster_data.append(
            {
                "cluster_id": cluster_id,
                "cluster_number": cluster_idx,
                "num_total_subbasins": len(subbasins_to_merge),
                "total_basin_area_km2": total_area_km2,
                "geometry_type": merged_geometry.geom_type,
                "geometry": merged_geometry,
            }
        )

        print(f"    Complete: {total_area_km2:,.0f} km²")

    # Create GeoDataFrame with merged geometries
    print("Creating final GeoDataFrame with merged geometries...")
    merged_gdf = gpd.GeoDataFrame(merged_cluster_data, crs="EPSG:4326")

    # Save to geoparquet
    print(f"Saving to {output_path}...")
    merged_gdf.to_parquet(output_path)


def get_touching_subbasins(
    data_catalog: DataCatalog, subbasins: gpd.GeoDataFrame
) -> gpd.GeoDataFrame:
    """Find all subbasins that touch the given subbasins.

    Args:
        data_catalog: Data catalog containing the MERIT basins.
        subbasins: GeoDataFrame containing the subbasins to find touching subbasins for.

    Returns:
        A GeoDataFrame containing all subbasins that touch the given subbasins.
    """
    bbox = subbasins.total_bounds
    buffer: float = 0.1
    buffered_bbox = (
        bbox[0] - buffer,
        bbox[1] - buffer,
        bbox[2] + buffer,
        bbox[3] + buffer,
    )
    potentially_touching_basins = gpd.read_parquet(
        data_catalog.get_source("MERIT_Basins_cat").path,
        bbox=buffered_bbox,
        filters=[
            ("COMID", "not in", subbasins.index.tolist()),
        ],
    )
    # get all touching subbasins
    touching_subbasins = potentially_touching_basins[
        potentially_touching_basins.geometry.touches(subbasins.union_all())
    ]

    return touching_subbasins.set_index("COMID")


def get_coastline_nodes(
    coastline_graph: networkx.Graph,
    riverine_mask: xr.DataArray,
    STUDY_AREA_OUTFLOW: int,
    NEARBY_OUTFLOW: int,
) -> set:
    """Get all coastline nodes that are part of the coastline for the study area.

    Args:
        coastline_graph: The graph containing all coastline nodes.
        riverine_mask: A DataArray containing the riverine mask.
        STUDY_AREA_OUTFLOW: The outflow type value for outflows within the study area.
        NEARBY_OUTFLOW: The outflow type value for outflows outside the study area, but close enough to influence the coastline.

    Returns:
        A set of all coastline nodes that are part of the coastline for the study area.

    Raises:
        AssertionError: If a coastal segment has both a study area outflow and a nearby outflow, but not exactly one of each.
    """
    coastline_nodes = set()

    for island in networkx.connected_components(coastline_graph):
        island = coastline_graph.subgraph(island)

        # extract the outflow nodes from the island, these should be included for sure
        island_outflow_nodes_study_area = set(
            [
                node
                for node, attrs in island.nodes(data=True)
                if attrs["outflow_type"] == STUDY_AREA_OUTFLOW
            ]
        )
        island_outflow_nodes_nearby = set(
            [
                node
                for node, attrs in island.nodes(data=True)
                if attrs["outflow_type"] == NEARBY_OUTFLOW
            ]
        )

        # no outflows found, we can skip this island
        if len(island_outflow_nodes_study_area) == 0:
            continue

        # if there is only one outflow node in the study area and no nearby
        # outflow nodes, it is an island, with only this outflow node
        # so we can add all nodes to the coastline
        if (
            len(island_outflow_nodes_study_area) == 1
            and len(island_outflow_nodes_nearby) == 0
        ):
            coastline_nodes.update(island.nodes)
            continue

        for node in island.nodes:
            island.nodes[node]["neighbor_of_study_area_outflow"] = False
            island.nodes[node]["neighbor_of_nearby_outflow"] = False

        island_with_outflow_cuts = island.copy()
        for node in island_outflow_nodes_study_area:
            neighbors = island.neighbors(node)
            for neighbor in neighbors:
                island_with_outflow_cuts.nodes[neighbor][
                    "neighbor_of_study_area_outflow"
                ] = True
                island_with_outflow_cuts.nodes[neighbor]["study_area_outflow_node"] = (
                    node
                )

        for node in island_outflow_nodes_nearby:
            neighbors = island.neighbors(node)
            for neighbor in neighbors:
                island_with_outflow_cuts.nodes[neighbor][
                    "neighbor_of_nearby_outflow"
                ] = True
                island_with_outflow_cuts.nodes[neighbor]["nearby_outflow_node"] = node

        island_with_outflow_cuts.remove_nodes_from(island_outflow_nodes_study_area)
        island_with_outflow_cuts.remove_nodes_from(island_outflow_nodes_nearby)

        for coastal_segment_nodes in networkx.connected_components(
            island_with_outflow_cuts
        ):
            coastal_segment = island_with_outflow_cuts.subgraph(
                coastal_segment_nodes
            ).copy()
            assert len(coastal_segment) == coastal_segment.number_of_nodes()

            study_area_nodes = [
                (node, attr)
                for node, attr in coastal_segment.nodes(data=True)
                if attr["neighbor_of_study_area_outflow"] is True
            ]

            nearby_nodes = [
                (node, attr)
                for node, attr in coastal_segment.nodes(data=True)
                if attr["neighbor_of_nearby_outflow"] is True
            ]

            outflow_neighbor_types = [
                attrs["neighbor_type"]
                for _, attrs in coastal_segment.nodes(data=True)
                if "neighbor_type" in attrs
            ]
            assert len(outflow_neighbor_types) <= 2

            # in case the segment has both a study area outflow and a nearby outflow
            # we divide the segment in a part that is closer to the study area outflow
            # and a part that is not
            if study_area_nodes and nearby_nodes:
                if len(study_area_nodes) != 1 or len(nearby_nodes) != 1:
                    # create a diagnostic visualization grid that marks, for this coastal
                    # segment, which cells neighbor the study area outflow, which neighbor
                    # the nearby outflow, and which belong to neither. This is written to
                    # disk to help debug cases where there is not exactly one study area
                    # outflow and one nearby outflow per coastal segment.
                    nodes_grid: xr.DataArray = riverine_mask.copy().astype(np.int32)
                    nodes_grid.attrs["_FillValue"] = 0
                    nodes_grid.values[:] = 0
                    for node, node_attributes in coastal_segment.nodes(data=True):
                        node_y, node_x = node_attributes["yx"]
                        if node_attributes["neighbor_of_nearby_outflow"] is True:
                            nodes_grid.values[node_y, node_x] = NEARBY_OUTFLOW
                        elif node_attributes["neighbor_of_study_area_outflow"] is True:
                            nodes_grid.values[node_y, node_x] = STUDY_AREA_OUTFLOW
                        else:
                            nodes_grid.values[node_y, node_x] = -1

                    # clip to area of interest for smaller output
                    nodes_grid, _ = clip_with_grid(nodes_grid, mask=nodes_grid != 0)

                    debug_file = Path("debug_coastal_segment.zarr")
                    write_zarr(
                        nodes_grid,
                        debug_file,
                        crs=nodes_grid.rio.crs,
                    )
                    raise AssertionError(
                        f"There should only be one study area outflow and one nearby outflow per coastal segment, "
                        f"found {len(study_area_nodes)} study area outflows and {len(nearby_nodes)} nearby outflows. "
                        f"Debug output written to {debug_file}."
                    )

                study_area_node = study_area_nodes[0]
                nearby_node = nearby_nodes[0]

                for node in coastal_segment.nodes:
                    # find distance to the nearest outflow node. If the path
                    # to the study area node is shorter, we add the node to the
                    # coastline nodes. If the path to the nearby node is shorter,
                    # we don't add the node to the coastline nodes.
                    path_length_to_study_area_node = len(
                        networkx.shortest_path(
                            coastline_graph, source=node, target=study_area_node[0]
                        )
                    )
                    path_length_to_nearby_node = len(
                        networkx.shortest_path(
                            coastline_graph, source=node, target=nearby_node[0]
                        )
                    )

                    # In case of break even, we need to take some special steps. We
                    # only want to add the node in one of the two datasets, if we
                    # also build another model for the nearby outflow. To make this
                    # decision, we need to use parameters that behave identically
                    # for both models. Therefore we use the relative xy location
                    # of the outflow nodes.
                    if path_length_to_study_area_node == path_length_to_nearby_node:
                        study_area_outflow_node = study_area_node[1][
                            "study_area_outflow_node"
                        ]
                        nearby_outflow_node = nearby_node[1]["nearby_outflow_node"]
                        if study_area_outflow_node[0] >= nearby_outflow_node[0]:
                            if study_area_outflow_node[0] == nearby_outflow_node[0]:
                                if study_area_outflow_node[0] > nearby_outflow_node[0]:
                                    coastline_nodes.add(node)
                                else:
                                    pass
                            else:
                                coastline_nodes.add(node)
                        else:
                            pass

                    elif path_length_to_study_area_node < path_length_to_nearby_node:
                        coastline_nodes.add(node)
                    else:
                        pass

            # if the segment has only a study area outflow, we add all nodes to
            # the coastline nodes
            elif study_area_nodes:
                coastline_nodes.update(coastal_segment.nodes)
            # if the segment has only a nearby outflows, we don't add any nodes to
            # the coastline nodes
            else:
                pass

    return coastline_nodes


def create_riverine_mask(
    ldd: xr.DataArray, ldd_network: pyflwdir.FlwdirRaster, geom: gpd.GeoDataFrame
) -> xr.DataArray:
    """Create a riverine mask from the ldd and the river network.

    Args:
        ldd : The local drainage direction (ldd) data array.
        ldd_network : The flow direction raster created from the ldd data.
        geom : The geometry of the riverine basin, which is used to clip the mask.

    Returns:
        A boolean mask where True indicates riverine cells and False indicates non-riverine cells.

    """
    riverine_mask = full_like(
        ldd,
        fill_value=True,
        nodata=False,
        dtype=bool,
    ).compute()

    riverine_mask = riverine_mask.rio.clip([geom.union_all()], drop=False)

    # MERIT-Basins rivers don't always extend all the way into coastline pits.
    # To extend these rivers, we first find all downstream cells of the area
    # within the initial riverine mask. Then we find all pits within these
    # downstream cells, and set these pits to be part of the riverine mask.
    downstream_indices_of_area_in_mask = ldd_network.idxs_ds[
        riverine_mask.values.ravel()
    ]
    all_pits_in_area = ldd_network.idxs_pit
    downstream_indices_that_are_pits = np.intersect1d(
        downstream_indices_of_area_in_mask, all_pits_in_area
    )
    riverine_mask.values.ravel()[downstream_indices_that_are_pits] = True

    # because the basin mask from the source is not perfect and has some holes
    # we need to extend the riverine mask to include all cells upstream of any
    # riverine cell. This is done by creating a flow raster from the masked
    # ldd, find all the pits within the riverine mask, and then get all upstream
    # cells from these pits.
    ldd_network_masked = pyflwdir.from_array(
        ldd.values,
        ftype="d8",
        mask=riverine_mask.values,
        transform=ldd.rio.transform(recalc=True),
        latlon=True,  # hydrography is specified in latlon
    )

    # set all cells that are upstream of a pit to True
    riverine_mask.values[ldd_network.basins(idxs=ldd_network_masked.idxs_pit) > 0] = (
        True
    )

    return riverine_mask


class DelayedReader(dict):
    """A dictionary that reads data from files only when accessed.

    This is useful because some datasets are very large and reading them
    all at once would require a lot of memory. Furthermore, when updating the model
    it is usually not required to have all data in memory. This class allows to
    read data only when it is actually needed.

    When setting an item, we should not set the actual data, but the file path.
    """

    def __init__(self, reader: Any) -> None:
        """Initialize the DelayedReader with a reader function.

        Args:
            reader: A function that takes a file path and returns the data.
        """
        self.reader: Any = reader

    def __getitem__(self, key: str) -> Any:
        """Get item from the dictionary using reader.

        Args:
            key: dictionary key

        Returns:
            The data read from the file.
        """
        fp: str | Path = super().__getitem__(key)
        return self.reader(fp)

    def __setitem__(self, key: str, value: Any) -> None:
        """Set item in the dictionary using a file path that refers to the actual data.

        The file path is stored in the dictionary, not the actual data.

        Args:
            key: dictionary key
            value: file path to the data

        Raises:
            ValueError: If the value is not a string or Path.
        """
        if isinstance(value, (str, Path)):
            super().__setitem__(key, value)
        else:
            raise ValueError("Value must be a file path (str or Path).")


class GEBModel(
    Hydrography, Forcing, Crops, LandSurface, Agents, GroundWater, Observations
):
    """Main GEB model build class.

    This class contains:
    - methods to setup the model region and grid
    - all general methods for example for saving and loading data, calling methods etc.
    - subclasses all build modules, which contain methods for building specific parts of the model.
    """

    def __init__(
        self,
        logger: logging.Logger,
        root: Path,
        data_catalog: str | None = None,
        epsg: int = 4326,
        data_provider: str = "default",
    ) -> None:
        """Initialize the GEB build model.

        Args:
            logger: Logger to use for logging.
            root: Root directory for the model build. If None, the current working directory is used.
            data_catalog: List of data catalogs to use. If None, the default data catalogs are used.
            epsg: EPSG code for the model grid. Default is 4326 (WGS84).
            data_provider: Data provider to use for the data catalog. Default is "default".
        """
        self._logger = logger
        build_method.logger = logger
        self._old_data_catalog = DataCatalog(
            data_libs=[data_catalog], logger=self._logger, fallback_lib=None
        )

        Hydrography.__init__(self)
        Forcing.__init__(self)
        Crops.__init__(self)
        LandSurface.__init__(self)
        Agents.__init__(self)
        GroundWater.__init__(self)
        Observations.__init__(self)

        self.root = root
        self.epsg = epsg
        self.data_provider = data_provider
        self._data_catalog = NewDataCatalog()

        # the grid, subgrid, and region subgrids are all datasets, which should
        # have exactly matching coordinates
        self.grid = xr.Dataset()
        self.subgrid = xr.Dataset()
        self.region_subgrid = xr.Dataset()

        # all other data types are dictionaries because these entries don't
        # necessarily match the grid coordinates, shapes etc.
        self.geom = DelayedReader(reader=gpd.read_parquet)
        self.table = DelayedReader(reader=pd.read_parquet)
        self.array = DelayedReader(zarr.load)
        self.params = DelayedReader(reader=read_params)
        self.other = DelayedReader(reader=read_zarr)
        self.files = {}

        self.maybe_update_version()

    def set_current_version(self) -> None:
        """Set the current version in the version file."""
        self.logger.info(
            f"Setting version in version file to current version: {__version__}"
        )
        self.version_path.write_text(__version__)

    def version_is_current(self) -> bool:
        """Check if the version in the version file is the same as the current version.

        Returns:
            True if the version is current, False otherwise.
        """
        if not self.version_path.exists():
            return False
        version_info = self.version_path.read_text()
        return version_info == __version__

    def maybe_update_version(self) -> None:
        """Check if the version in the version file is the same as the current version.

        If the version is not current, print a warning with the updates that need to be made to update to the current version.

        Raises:
            RuntimeError: If the version is not current and updates need to be made.
        """
        # No version file exists, so we create one with the current version
        if not self.version_path.exists():
            self.set_current_version()
            return
        version_info = self.version_path.read_text()
        if self.version_is_current():
            self.logger.info("Version is already current.")
        else:
            self.logger.warning(
                f"Version mismatch: version file contains {version_info}, but current version is {__version__}."
            )
            # Find and print all updates between the stored version and the current version
            current_v = Version(__version__)
            stored_v = Version(version_info)

            versions = sorted(VERSION_UPDATES.keys(), key=Version)
            updates_to_print = []
            for v_str in versions:
                v = Version(v_str)
                if v > stored_v and v <= current_v:
                    updates_to_print.extend(VERSION_UPDATES[v_str])

            if updates_to_print:
                updates_msg = "\n- ".join(updates_to_print)
                self.set_current_version()
                error = f"\n\nIMPORTANT: Make the following changes to update to this version:\n\n- {updates_msg}\n\nTHIS WARNING WILL ONLY BE GIVEN ONCE. If you already did this, you can ignore this.\n"
                self.logger.error(error)
                raise RuntimeError(error)
            else:
                self.logger.info(
                    "No specific updates found for this version. Updated version file."
                )
                self.set_current_version()

    @property
    def logger(self) -> logging.Logger:
        """Get the logger."""
        return self._logger

    @logger.setter
    def logger(self, value: logging.Logger) -> None:
        self._logger = value
        build_method.logger = value
        # Ensure that child classes use the updated logger
        if hasattr(self, "_old_data_catalog"):
            self._old_data_catalog.logger = value

    @property
    def old_data_catalog(self) -> DataCatalog:
        """Get the data catalog."""
        return self._old_data_catalog

    @old_data_catalog.setter
    def old_data_catalog(self, value: DataCatalog) -> None:
        self._old_data_catalog: DataCatalog = value

    @property
    def data_catalog(self) -> NewDataCatalog:
        """Get the new data catalog."""
        return self._data_catalog

    @data_catalog.setter
    def data_catalog(self, value: NewDataCatalog) -> None:
        self._data_catalog = value

    @property
    def grid(self) -> xr.Dataset:
        """Get the grid."""
        return self._grid

    @grid.setter
    def grid(self, value: xr.Dataset) -> None:
        self._grid = value

    @property
    def subgrid(self) -> xr.Dataset:
        """Get the subgrid."""
        return self._subgrid

    @subgrid.setter
    def subgrid(self, value: xr.Dataset) -> None:
        self._subgrid = value

    @property
    def region_subgrid(self) -> xr.Dataset:
        """Get the region subgrid."""
        return self._region_subgrid

    @region_subgrid.setter
    def region_subgrid(self, value: xr.Dataset) -> None:
        self._region_subgrid = value

    @property
    def geom(self) -> DelayedReader:
        """Get the geometry reader."""
        return self._geom

    @geom.setter
    def geom(self, value: DelayedReader) -> None:
        self._geom = value

    @property
    def table(self) -> DelayedReader:
        """Get the table reader."""
        return self._table

    @table.setter
    def table(self, value: DelayedReader) -> None:
        self._table = value

    @property
    def array(self) -> DelayedReader:
        """Get the array reader."""
        return self._array

    @array.setter
    def array(self, value: DelayedReader) -> None:
        self._array = value

    @property
    def params(self) -> DelayedReader:
        """Get the params reader."""
        return self._params

    @params.setter
    def params(self, value: DelayedReader) -> None:
        self._params = value

    @property
    def other(self) -> DelayedReader:
        """Get the other reader."""
        return self._other

    @other.setter
    def other(self, value: DelayedReader) -> None:
        self._other = value

    @property
    def files(self) -> dict:
        """Get the files dictionary."""
        return self._files

    @files.setter
    def files(self, value: dict) -> None:
        self._files = value

    @build_method(required=True)
    def setup_region(
        self,
        region: dict,
        subgrid_factor: int,
        resolution_arcsec: int = 30,
        include_coastal_area: bool = True,
    ) -> None:
        """Creates a 2D regular grid or reads an existing grid.

        An 2D regular grid will be created from a geometry (geom_fn) or bbox. If an existing
        grid is given, then no new grid will be generated.

        Adds/Updates model layers:
        * **grid** grid mask: add grid mask to grid object

        Args:
            region: Dictionary describing region of interest, e.g.:
                * {'basin': [x, y]}

                Region must be of kind [basin, subbasin].
            subgrid_factor: GEB implements a subgrid. This parameter determines the factor by which the subgrid is smaller than the original grid.
            resolution_arcsec: Resolution of the grid in arcseconds. Must be a multiple of 3 to align with MERIT.
            include_coastal_area: Because the subbasins are delinated using a minimum upstream area small basins near the coast are not included.
                If this parameter is set to True, the coastal area will be included in the riverine mask by automatically extending the riverine mask to the coastal area,
                by finding all coastal basins between the outlets within the study area and half the distance to the nearest outlet outside the study area.
                All cells upstream of these coastal basins will be included in the riverine mask.

        Raises:
            ValueError: If region is not understood.
        """
        assert resolution_arcsec % 3 == 0, (
            "resolution_arcsec must be a multiple of 3 to align with MERIT"
        )
        assert subgrid_factor >= 2

        self.logger.info("Loading river network.")
        river_graph = get_river_graph(self.data_catalog)

        self.logger.info("Finding sinks in river network of requested region.")
        if "subbasin" in region:
            if isinstance(region["subbasin"], list):
                sink_subbasin_ids = region["subbasin"]
            else:
                sink_subbasin_ids = [region["subbasin"]]
        elif "outflow" in region:
            lat, lon = region["outflow"]["lat"], region["outflow"]["lon"]
            sink_subbasin_ids = [
                get_subbasin_id_from_coordinate(self.data_catalog, lon, lat)
            ]
        elif "geom" in region:
            regions = self.data_catalog.fetch(region["geom"]["source"]).read()
            assert isinstance(regions, gpd.GeoDataFrame)
            regions = regions[
                regions[region["geom"]["column"]] == region["geom"]["key"]
            ]
            sink_subbasin_ids = get_sink_subbasin_id_for_geom(
                self.data_catalog, regions, river_graph
            )
        else:
            raise ValueError(f"Region {region} not understood.")

        self.logger.info(
            f"Found {len(sink_subbasin_ids)} sink subbasins in region {region}."
        )
        rivers: gpd.GeoDataFrame = self.get_rivers(river_graph, sink_subbasin_ids)

        buffer = 1.0  # buffer in degrees
        xmin, ymin, xmax, ymax = rivers[~rivers["is_downstream_outflow"]].total_bounds
        xmin -= buffer
        ymin -= buffer
        xmax += buffer
        ymax += buffer

        ldd = (
            self.data_catalog.fetch(
                "merit_hydro_dir",
                xmin=xmin,
                xmax=xmax,
                ymin=ymin,
                ymax=ymax,
            )
            .read()
            .compute()
        )
        assert isinstance(ldd, xr.DataArray), "Expected ldd to be an xarray DataArray."

        # We remove all pits that are not directly adjacent to valid flow directions
        # Identify cells with flow directions
        cells_with_flow_directions = (ldd.values > 0) & (ldd.values != 247)
        # grow valid values mask by one cell
        valid_values_and_coastline = binary_dilation(
            cells_with_flow_directions, structure=np.ones((3, 3))
        )
        # and use said mask which includes valid values and its neighbors to set ldd to no data
        ldd.values[~valid_values_and_coastline] = 247

        ldd_network = pyflwdir.from_array(
            ldd.values,
            ftype="d8",
            transform=ldd.rio.transform(recalc=True),
            latlon=True,
        )

        rivers: gpd.GeoDataFrame = extend_rivers_into_ocean(rivers, ldd_network)

        self.logger.info("Preparing 2D grid.")
        if "outflow" in region:
            # get basin geometry
            riverine_mask = full_like(
                ldd,
                fill_value=False,
                nodata=False,
                dtype=bool,
            )
            riverine_mask.values[ldd_network.basins(xy=(lon, lat)) > 0] = True
            rivers = rivers[~rivers["is_upstream_of_downstream_basin"]]
        elif "subbasin" in region or "geom" in region:
            rivers_outlets_for_basins = rivers[~rivers["is_further_downstream_outflow"]]
            outlet_lonlats = rivers_outlets_for_basins.geometry.apply(
                lambda geom: geom.coords[-2]
            ).tolist()
            subbasins_grid = ldd_network.basins(
                xy=(
                    [lon for lon, lat in outlet_lonlats],
                    [lat for lon, lat in outlet_lonlats],
                ),
                ids=rivers_outlets_for_basins.index,
            ).astype(np.int32)

            # we want to remove the areas upstream of the downstream outflow basins
            # that are not part of the study area.
            is_upstream_of_downstream_basin = rivers["is_upstream_of_downstream_basin"]

            # by setting these basins to 0 in the subbasins grid, they will be removed
            subbasins_grid[
                np.isin(subbasins_grid, rivers[is_upstream_of_downstream_basin].index)
            ] = 0

            # and we can now remove them from the rivers geodataframe as well. They have had
            # their purpose now.
            rivers = rivers[~is_upstream_of_downstream_basin]

            subbasins: list[tuple[dict[str, Any], float]] = list(
                rasterio.features.shapes(
                    subbasins_grid,
                    mask=subbasins_grid != 0,
                    transform=ldd_network.transform,
                    connectivity=8,
                )
            )
            subbasins: gpd.GeoDataFrame = (
                gpd.GeoDataFrame.from_records(
                    [
                        {
                            "COMID": int(value),
                            "geometry": shape(geom),
                        }
                        for geom, value in subbasins
                    ],
                )
                .set_index("COMID")
                .set_crs(4326)
            )
            subbasins["is_downstream_outflow"] = subbasins.index.isin(
                rivers[rivers["is_downstream_outflow"]].index
            )
            subbasins["is_coastal"] = subbasins.apply(
                lambda subbasin: rivers.loc[subbasin.name, "downstream_ID"] == -1,
                axis=1,
            )

            geom = gpd.GeoDataFrame(
                geometry=[subbasins[~subbasins["is_downstream_outflow"]].union_all()],
                crs=subbasins.crs,
            )
            # ESPG 6933 (WGS 84 / NSIDC EASE-Grid 2.0 Global) is an equal area projection
            # while the shape of the polygons becomes vastly different, the area is preserved mostly.
            # usable between 86°S and 86°N.
            self.logger.info(
                f"Approximate riverine basin size: {round(geom.to_crs(epsg=6933).area.sum() / 1e6, 2)} km2"
            )

            riverine_mask = create_riverine_mask(ldd, ldd_network, geom)
            assert not riverine_mask.attrs["_FillValue"]
        else:
            raise ValueError(f"Region {region} not understood.")

        self.set_geom(rivers, name="routing/rivers")
        self.set_geom(subbasins, name="routing/subbasins")

        if include_coastal_area and subbasins["is_coastal"].any():
            mask: xr.DataArray = self.extend_mask_to_coastal_area(
                ldd, riverine_mask, subbasins
            )
        else:
            mask: xr.DataArray = riverine_mask

        mask.attrs["_FillValue"] = None
        self.set_other(mask, name="drainage/mask")

        ldd: xr.DataArray = xr.where(
            mask,
            ldd,
            ldd.attrs["_FillValue"],
        )

        ldd_elevation = self.data_catalog.fetch(
            "merit_hydro_elv",
            xmin=xmin,
            xmax=xmax,
            ymin=ymin,
            ymax=ymax,
        ).read()
        assert isinstance(ldd_elevation, xr.DataArray), (
            "Expected ldd_elevation to be an xarray DataArray."
        )

        assert ldd_elevation.shape == ldd.shape == mask.shape

        ldd_elevation = xr.where(
            mask,
            ldd_elevation,
            ldd_elevation.attrs["_FillValue"],
        )

        mask, ldd, ldd_elevation = clip_region(
            mask, ldd, ldd_elevation, align=30 / 60 / 60
        )

        self.set_other(ldd_elevation, name="drainage/original_d8_elevation")

        ldd: xr.DataArray = xr.where(
            mask,
            ldd,
            ldd.attrs["_FillValue"],
            keep_attrs=True,
        )
        ldd: xr.DataArray = self.set_other(
            ldd, name="drainage/original_d8_flow_directions"
        )

        self.derive_mask(ldd, ldd.rio.transform(recalc=True), resolution_arcsec)

        self.create_subgrid(subgrid_factor)

    def extend_mask_to_coastal_area(
        self,
        ldd: xr.DataArray,
        riverine_mask: xr.DataArray,
        subbasins: gpd.GeoDataFrame,
    ) -> xr.DataArray:
        """Extend the riverine mask to include coastal areas.

        This is done by finding all coastal basins between the outlets within the study area
        and half the distance to the nearest outlet outside the study area. All cells upstream of these
        coastal basins will be included in the riverine mask.

        Args:
            ldd: The local drainage direction (ldd) data array.
            riverine_mask: The initial riverine mask, which will be extended to the coastal area.
            subbasins: The subbasins used to create the initial riverine mask.

        Returns:
            The extended riverine mask.
        """
        flow_raster: pyflwdir.FlwdirRaster = pyflwdir.from_array(
            ldd.values,
            ftype="d8",
            transform=ldd.rio.transform(recalc=True),
            latlon=True,
        )

        STUDY_AREA_OUTFLOW: Literal[1] = 1
        NEARBY_OUTFLOW: Literal[2] = 2

        rivers = (
            self.data_catalog.fetch(
                "merit_basins_rivers",
            )
            .read(
                columns=["COMID", "lengthkm", "uparea", "maxup", "geometry"],
                bbox=ldd.rio.bounds(),
            )
            .set_index("COMID")
        )
        assert isinstance(rivers, gpd.GeoDataFrame), (
            "Expected rivers to be a GeoDataFrame."
        )

        rivers["outflow_type"] = rivers.apply(
            lambda row: (
                STUDY_AREA_OUTFLOW if row.name in subbasins.index else NEARBY_OUTFLOW
            ),
            axis=1,
        )

        river_raster_outflow_type = create_river_raster_from_river_lines(
            rivers, ldd, column="outflow_type"
        )
        river_raster_ID = create_river_raster_from_river_lines(rivers, ldd, index=True)
        river_raster_ID = river_raster_ID.ravel()

        downstream_indices = flow_raster.idxs_ds

        # MERIT-Basins rivers don't always extend all the way into coastline pits.
        # To extend these rivers, we first find the river (which have a river ID),
        # but have a no downstream cell.
        river_cells_with_downstream_pits = (downstream_indices != -1) & (
            river_raster_ID != -1
        )
        # For those cells, we extract the river ID, and set this river ID to the downstream cells.
        river_raster_ID[downstream_indices[river_cells_with_downstream_pits]] = (
            river_raster_ID[river_cells_with_downstream_pits]
        )
        river_raster_ID = river_raster_ID.reshape(ldd.shape)

        pits = ldd == 0
        # TODO: Filter non-coastline pits
        coastline = pits
        coastline.attrs["_FillValue"] = None

        # save coastline, for now mostly for debugging purposes
        self.set_other(coastline, name="drainage/full_coastline_in_bounding_box")

        pits = flow_raster.idxs_pit

        downstream_indices_study_area = flow_raster.idxs_ds[
            river_raster_outflow_type.ravel() == STUDY_AREA_OUTFLOW
        ]
        outflow_pits_study_area = np.intersect1d(pits, downstream_indices_study_area)

        downstream_indices_nearby = flow_raster.idxs_ds[
            river_raster_outflow_type.ravel() == NEARBY_OUTFLOW
        ]
        outflow_pits_nearby = np.intersect1d(pits, downstream_indices_nearby)

        outflows = np.full(ldd.size, -1, dtype=np.int8)
        outflows[outflow_pits_study_area] = STUDY_AREA_OUTFLOW
        outflows[outflow_pits_nearby] = NEARBY_OUTFLOW
        outflows = outflows.reshape(ldd.shape)

        outflow_da = self.full_like(
            ldd,
            fill_value=False,
            nodata=-1,
            dtype=bool,
        )
        outflow_da.values = outflows
        self.set_other(outflow_da, name="drainage/outflows")

        # first we create a graph with all coastline cells. Neighbouring cells
        # are connected. From this graph we want to select all coastline
        # cells in between river outlets. The riverine mask does not include
        # any of the coastal cells. So we first grow the riverine mask by one cell
        # and use this to set the outflow nodes in the graph.
        coastline_graph = boolean_mask_to_graph(
            coastline.values,
            connectivity=4,
            outflow_type=outflows,
            river_id=river_raster_ID,
        )
        coastline_nodes = get_coastline_nodes(
            coastline_graph,
            riverine_mask,
            STUDY_AREA_OUTFLOW=STUDY_AREA_OUTFLOW,
            NEARBY_OUTFLOW=NEARBY_OUTFLOW,
        )

        # here we go from the graph back to a mask. We do this by creating a new mask
        # and setting coastline cells to true
        simulated_coastline_da = self.full_like(
            ldd,
            fill_value=False,
            nodata=None,
            dtype=bool,
        )
        simulated_coastline = np.zeros(
            ldd.shape,
            dtype=bool,
        )

        for node in coastline_nodes:
            yx = coastline_graph.nodes[node]["yx"]
            simulated_coastline[yx] = True

        simulated_coastline_da.values = simulated_coastline

        # save the connected coastline, for now mostly for debugging purposes
        self.set_other(simulated_coastline_da, name="drainage/simulated_coastline")

        # get all upstream cells from the selected coastline
        flow_raster = pyflwdir.from_array(ldd.values, ftype="d8")
        coastal_mask = (
            flow_raster.basins(idxs=np.where(simulated_coastline_da.values.ravel())[0])
            > 0
        )

        # return the combination of the riverine mask and the coastal mask
        return riverine_mask | coastal_mask

    def derive_mask(
        self,
        d8_original: xr.DataArray,
        transform: Affine,
        target_resolution_arcsec: int,
    ) -> None:
        """Derives the model grid mask from the original D8 flow directions.

        The model grid mask is derived by upscaling the original D8 flow directions
        to the model resolution using the IHU method (Iterative Hydrography Upscaling).
        The model grid mask is True for all cells that are part of the upscaled flow
        directions, and False for all other cells.

        Args:
            d8_original: The original D8 flow directions.
            transform: The affine transform of the original D8 flow directions.
            target_resolution_arcsec: The resolution of the target low resolution model grid in arcseconds.
        """
        assert d8_original.dtype == np.uint8

        d8_original_data = d8_original.values
        flow_raster = pyflwdir.from_array(
            d8_original_data,
            ftype="d8",
            transform=transform,
            latlon=True,  # hydrography is specified in latlon
            mask=d8_original_data
            != d8_original.attrs["_FillValue"],  # this mask is True within study area
        )

        # the original resolution is 3 arcseconds, thus we divide by 3
        scale_factor: int = target_resolution_arcsec // 3

        self.logger.info("Coarsening hydrography")
        # IHU = Iterative hydrography upscaling method, see https://doi.org/10.5194/hess-25-5287-2021
        flow_raster_upscaled, idxs_out = flow_raster.upscale(
            scale_factor=scale_factor,
            method="ihu",
        )
        flow_raster_upscaled.repair_loops()

        mask = xr.DataArray(
            ~flow_raster_upscaled.mask.reshape(flow_raster_upscaled.shape),
            coords={
                "y": flow_raster_upscaled.transform.f
                + flow_raster_upscaled.transform.e
                * (np.arange(flow_raster_upscaled.shape[0]) + 0.5),
                "x": flow_raster_upscaled.transform.c
                + flow_raster_upscaled.transform.a
                * (np.arange(flow_raster_upscaled.shape[1]) + 0.5),
            },
            dims=("y", "x"),
            name="mask",
            attrs={
                "_FillValue": None,
            },
        )
        self.set_grid(mask, name="mask")

        mask_geom = list(
            rasterio.features.shapes(
                mask.astype(np.uint8),
                mask=~mask,
                connectivity=8,
                transform=mask.rio.transform(recalc=True),
            ),
        )
        mask_geom = gpd.GeoDataFrame.from_features(
            [{"geometry": geom[0], "properties": {}} for geom in mask_geom],
            crs=4326,
        )

        self.set_geom(mask_geom, name="mask")

        flow_raster_idxs_ds = self.full_like(
            self.grid["mask"],
            fill_value=-1,
            nodata=-1,
            dtype=np.int32,
        )
        flow_raster_idxs_ds.name = "flow_raster_idxs_ds"
        flow_raster_idxs_ds.data = flow_raster_upscaled.idxs_ds.reshape(
            flow_raster_upscaled.shape
        )
        self.set_grid(flow_raster_idxs_ds, name="flow_raster_idxs_ds")

        idxs_out_da = self.full_like(
            self.grid["mask"],
            fill_value=-1,
            nodata=-1,
            dtype=np.int32,
        )
        idxs_out_da.name = "idxs_outflow"
        idxs_out_da.data = idxs_out
        self.set_grid(idxs_out_da, name="idxs_outflow")

    def create_subgrid(self, subgrid_factor: int) -> None:
        """Creates the model subgrid.Affine.

        The model subgrid is a higher resolution grid that is used for representing
        subgrid processes, such as different land uses within a grid cell, and
        unique agents.

        Args:
            subgrid_factor: The factor by which the subgrid is smaller than the original grid.
        """
        mask: xr.DataArray = self.grid["mask"]
        dst_transform: Affine = mask.rio.transform(recalc=True) * Affine.scale(
            1 / subgrid_factor
        )

        submask: xr.DataArray = xr.DataArray(
            data=repeat_grid(mask.data, subgrid_factor),
            coords={
                "y": dst_transform.f
                + np.arange(mask.shape[0] * subgrid_factor) * dst_transform.e,
                "x": dst_transform.c
                + np.arange(mask.shape[1] * subgrid_factor) * dst_transform.a,
            },
            attrs={"_FillValue": None},
        )

        self.set_subgrid(submask, name="mask")

    @build_method(required=True)
    def set_time_range(self, start_date: date, end_date: date) -> None:
        """Sets the time range for the build model.

        This time range is used to ensure that all datasets with a time dimension
        cover at least this time range.

        Start date must be on or after 1960, because of data availability. End date can be in the future.

        Args:
            start_date: The start date of the model.
            end_date: The end date of the model.

        Raises:
            ValueError: If the start date is not before the end date.
            ValueError: If the start date is before 1960, because of data availability.
        """
        if not start_date < end_date:
            raise ValueError("Start date must be before end date.")

        if start_date.year < 1960:
            raise ValueError(
                "Start date must be on or after 1960, because of data availability."
            )
        self.set_params(
            {"start_date": start_date, "end_date": end_date},
            name="model_time_range",
        )

    @property
    def start_date(self) -> datetime:
        """Gets the start date of the build model.

        All datasets with a time dimension should cover at least from this date.

        Returns:
            The start date of the model.
        """
        start_date = self.params["model_time_range"]["start_date"]

        # TODO: This can be removed in 2026
        if isinstance(start_date, str):
            start_date = datetime.fromisoformat(start_date)
        return start_date

    @property
    def end_date(self) -> datetime:
        """Gets the end date of the build model.

        All datasets with a time dimension should cover at least until this date.

        Returns:
            The end date of the model.
        """
        end_date = self.params["model_time_range"]["end_date"]

        # TODO: This can be removed in 2026
        if isinstance(end_date, str):
            end_date = datetime.fromisoformat(end_date)
        return end_date

    @build_method(required=True)
    def set_ssp(self, ssp: str) -> None:
        """Sets the SSP name for the model.

        Args:
            ssp: The SSP name. Supported SSPs are: ssp1, ssp3, ssp5.
        """
        assert ssp in ["ssp1", "ssp3", "ssp5"], (
            f"SSP {ssp} not supported. Supported SSPs are: ssp1, ssp3, ssp5."
        )
        self.set_params({"ssp": ssp}, name="ssp")

    @property
    def ssp(self) -> str:
        """Gets the SSP name that was set using set_ssp, or the default SSP of "ssp3".

        Returns:
            The SSP name. If no SSP was set, returns "ssp3" (middle of the road).
        """
        return self.params["ssp"]["ssp"] if "ssp" in self.params else "ssp3"

    @property
    def ISIMIP_ssp(self) -> str:
        """Returns the ISIMIP SSP name.

        Raises:
            ValueError: If the SSP is not supported.
        """
        if self.ssp == "ssp1":
            return "ssp126"
        elif self.ssp == "ssp3":
            return "ssp370"
        elif self.ssp == "ssp5":
            return "ssp585"
        else:
            raise ValueError(f"SSP {self.ssp} not supported.")

    def setup_coastal_water_levels(
        self,
    ) -> None:
        """Sets up coastal water level data from the GTSM dataset.

        Filters the dataset to include only stations within the model bounds,
        and ensures that the time dimension is consistent.
        """
        water_levels = self.old_data_catalog.get_dataset("GTSM")
        assert isinstance(water_levels, xr.DataArray)
        assert (
            water_levels.time.diff("time").astype(np.int64)
            == (water_levels.time[1] - water_levels.time[0]).astype(np.int64)
        ).all()
        # convert to geodataframe
        stations = gpd.GeoDataFrame(
            water_levels.stations,
            geometry=gpd.points_from_xy(
                water_levels.station_x_coordinate, water_levels.station_y_coordinate
            ),
        )
        # filter all stations within the bounds, considering a buffer
        station_ids = stations.cx[
            self.bounds[0] - 0.1 : self.bounds[2] + 0.1,
            self.bounds[1] - 0.1 : self.bounds[3] + 0.1,
        ].index.values

        water_levels = water_levels.sel(stations=station_ids)

        assert len(water_levels.stations) > 0, (
            "No stations found in the region. If no stations should be set, set include_coastal_area=False"
        )

        self.set_other(
            water_levels,
            name="waterlevels",
            time_chunksize=24 * 6,  # 10 minute data
        )

    @build_method(required=True)
    def setup_damage_parameters(
        self,
        parameters: dict[
            str, dict[str, dict[str, dict[str, list[tuple[float, float]] | float]]]
        ],
    ) -> None:
        """Sets up damage parameters for different hazards and asset types.

        Args:
            parameters: A nested dictionary containing damage parameters. The structure is:
                {
                    "hazard_type": {
                        "asset_type": {
                            "component": {
                                "curve": [(severity, damage_ratio), ...],
                                "maximum_damage": float
                            },
                            ...
                        },
                        ...
                    },
                    ...
                }
        """
        for hazard, hazard_parameters in parameters.items():
            for asset_type, asset_parameters in hazard_parameters.items():
                for component, asset_compontents in asset_parameters.items():
                    curve = pd.DataFrame(
                        asset_compontents["curve"],
                        columns=np.array(["severity", "damage_ratio"]),
                    )

                    self.set_table(
                        curve,
                        name=f"damage_parameters/{hazard}/{asset_type}/{component}/curve",
                    )

                    maximum_damage = {
                        "maximum_damage": asset_compontents["maximum_damage"]
                    }

                    self.set_params(
                        maximum_damage,
                        name=f"damage_parameters/{hazard}/{asset_type}/{component}/maximum_damage",
                    )

    @build_method(required=False)
    def setup_precipitation_scaling_factors_for_return_periods(
        self, risk_scaling_factors: list[tuple[float, float]]
    ) -> None:
        """Sets up precipitation scaling factors for different return periods.

        Args:
            risk_scaling_factors: A list of tuples containing the exceedance probability
                and the corresponding scaling factor.
        """
        risk_scaling_factors_df = pd.DataFrame(
            risk_scaling_factors,
            columns=np.array(["exceedance_probability", "scaling_factor"]),
        )
        self.set_table(risk_scaling_factors_df, name="precipitation_scaling_factors")

    def set_table(self, table: pd.DataFrame, name: str, write: bool = True) -> None:
        """Set a table and save it to disk.

        Args:
            table: The table to save.
            name: The name of the table.
            write: Whether to write the table to disk. If False, the table
                is only added to the file library, but not written to disk.
        """
        fp: Path = Path("table") / (name + ".parquet")
        fp_with_root: Path = Path(self.root, fp)
        if write:
            self.logger.info(f"Writing file {fp}")

            self.files["table"][name] = fp

            fp_with_root.parent.mkdir(parents=True, exist_ok=True)
            write_table(table, fp_with_root)

        self.table[name] = fp_with_root

    def set_array(self, data: np.ndarray, name: str, write: bool = True) -> None:
        """Set an array and save it to disk.

        Args:
            data: The array to save.
            name: The name of the array.
            write: Whether to write the array to disk. If False, the array
                is only added to the file library, but not written to disk.
        """
        assert isinstance(data, np.ndarray)
        fp: Path = Path("array") / (name + ".zarr")
        fp_with_root: Path = Path(self.root, fp)

        if write:
            self.logger.info(f"Writing file {fp}")
            self.files["array"][name] = fp
            fp_with_root.parent.mkdir(parents=True, exist_ok=True)
            write_array(data, fp_with_root)

        self.array[name] = fp_with_root

    def set_params(self, data: dict, name: str, write: bool = True) -> None:
        """Set a dictionary and save it to disk.

        Args:
            data: The dictionary to save.
            name: The name of the dictionary.
            write: Whether to write the dictionary to disk. If False, the dictionary
                is only added to the file library, but not written to disk.
        """
        fp: Path = Path("dict") / (name + ".yml")
        fp_with_root: Path = Path(self.root) / fp
        fp_with_root.parent.mkdir(parents=True, exist_ok=True)
        if write:
            self.logger.info(f"Writing file {fp}")

            self.files["dict"][name] = fp

            write_params(data, fp_with_root)

        self.params[name] = fp_with_root

    def set_geom(self, geom: gpd.GeoDataFrame, name: str, write: bool = True) -> None:
        """Set a geometry and save it to disk.

        Args:
            geom: The geometry to save.
            name: The name of the geometry.
            write: Whether to write the geometry to disk. If False, the geometry
                is only added to the file library, but not written to disk.
        """
        assert isinstance(geom, gpd.GeoDataFrame)
        fp: Path = Path("geom") / (name + ".geoparquet")
        fp_with_root: Path = self.root / fp
        if write:
            self.logger.info(f"Writing file {fp}")
            self.files["geom"][name] = fp
            fp_with_root.parent.mkdir(parents=True, exist_ok=True)
            write_geom(geom, fp_with_root)

        self.geom[name] = fp_with_root

    @property
    def files_path(self) -> Path:
        """Path to the files.yml file that contains the file library."""
        return Path(self.root, "files.yml")

    @property
    def progress_path(self) -> Path:
        """Path to the progress file that contains the build progress."""
        return Path(self.root, "progress.txt")

    @property
    def build_complete_path(self) -> Path:
        """Path to the file that indicates that the build is complete."""
        return Path(self.root, "build_complete.txt")

    @property
    def version_path(self) -> Path:
        """Path to the version file that contains the build version."""
        return Path(self.root, "version.txt")

    def write_file_library(self) -> None:
        """Writes the file library to disk.

        Rather than overwriting the file library, we read the existing
        file library from disk and merge it with the new files, prioritizing
        the new files. This way, if another process has added files to the
        file library, we don't overwrite them.
        """
        file_library: dict = self.read_or_create_file_library()
        # merge file library from disk with new files, prioritizing new files
        for type_name, type_files in self.files.items():
            if type_name not in file_library:
                file_library[type_name] = type_files
            else:
                file_library[type_name].update(type_files)

        write_params(file_library, self.files_path)

    def read_or_create_file_library(self) -> dict:
        """Reads the file library from disk.

        If the file library does not exist, an empty file library is returned, but
        with all expected categories.

        Returns:
            A dictionary with the file library.
        """
        fp: Path = Path(self.files_path)
        if not fp.exists():
            return {
                "geom": {},
                "array": {},
                "table": {},
                "dict": {},
                "grid": {},
                "subgrid": {},
                "region_subgrid": {},
                "other": {},
            }
        else:
            files = read_params(self.files_path)

            # geoms was renamed to geom in the file library. To upgrade old models,
            # we check if "geoms" is in the files and rename it to "geom"
            # this line can be removed in august 2026 (also in geb/model.py)
            if "geoms" in files:
                files["geom"] = files.pop("geoms", {})
        return files

    def read_geom(self) -> None:
        """Reads all geometries from disk based on the file library."""
        for name, fn in self.files["geom"].items():
            self.geom[name] = Path(self.root, fn)

    def read_array(self) -> None:
        """Reads all arrays from disk based on the file library."""
        for name, fn in self.files["array"].items():
            self.array[name] = Path(self.root, fn)

    def read_table(self) -> None:
        """Reads all tables from disk based on the file library."""
        for name, fn in self.files["table"].items():
            self.table[name] = Path(self.root, fn)

    def read_params(self) -> None:
        """Reads all dictionaries from disk based on the file library."""
        for name, fn in self.files["dict"].items():
            self.params[name] = Path(self.root, fn)

    def read_grid(self) -> None:
        """Reads all grid data arrays from disk based on the file library."""
        # first read and set the mask. This is required.
        grid_files: dict[str, Path] = self.files["grid"]
        if len(grid_files) == 0:
            return
        mask: xr.DataArray = read_zarr(Path(self.root) / grid_files["mask"])
        self.set_grid(mask, name="mask", write=False)

        for name, fn in self.files["grid"].items():
            if name == "mask":  # mask already read
                continue
            data: xr.DataArray = read_zarr(Path(self.root) / fn)
            self.set_grid(data, name=name, write=False)

    def read_subgrid(self) -> None:
        """Reads all subgrid data arrays from disk based on the file library."""
        # first read and set the mask. This is required.
        subgrid_files: dict[str, Path] = self.files["subgrid"]
        if len(subgrid_files) == 0:
            return
        mask: xr.DataArray = read_zarr(Path(self.root) / subgrid_files["mask"])
        self.set_subgrid(mask, name="mask", write=False)
        for name, fn in self.files["subgrid"].items():
            if name == "mask":  # mask already read
                continue
            data: xr.DataArray = read_zarr(Path(self.root) / fn)
            self.set_subgrid(data, name=name, write=False)

    def read_region_subgrid(self) -> None:
        """Reads all region subgrid data arrays from disk based on the file library."""
        # first read and set the mask. This is required.
        region_subgrid_files: dict[str, Path] = self.files["region_subgrid"]
        if len(region_subgrid_files) == 0:
            return
        mask: xr.DataArray = read_zarr(Path(self.root) / region_subgrid_files["mask"])
        self.set_region_subgrid(mask, name="mask", write=False)
        for name, fn in self.files["region_subgrid"].items():
            if name == "mask":  # mask already read
                continue
            data: xr.DataArray = read_zarr(Path(self.root) / fn)
            self.set_region_subgrid(data, name=name, write=False)

    def read_other(self) -> None:
        """Reads all "other" data arrays from disk based on the file library."""
        for name, fn in self.files["other"].items():
            self.other[name] = Path(self.root, fn)

    def read(self) -> None:
        """Reads all data from disk based on the file library.

        Useful when continuing a build from an existing folder.
        """
        self.files = self.read_or_create_file_library()
        with suppress_logging_warning(self.logger):
            self.read_geom()
            self.read_array()
            self.read_table()
            self.read_params()

            self.read_subgrid()
            self.read_grid()
            self.read_region_subgrid()

            self.read_other()

    def set_other(
        self,
        da: xr.DataArray,
        name: str,
        write: bool = True,
        x_chunksize: int = XY_CHUNKSIZE,
        y_chunksize: int = XY_CHUNKSIZE,
        time_chunksize: int = 1,
        **kwargs: Any,
    ) -> xr.DataArray:
        """Set a data array that does not fit into other categories.

        The data array is saved to disk as a zarr file. Unlike the other gridded
        methods, this category does not assume any specific coordinate system
        or dimension order, and does not combine multiple layers into a single dataset.

        Args:
            da: The data array to set.
            name: The name of the data array.
            write: If True, write the data array to disk. Defaults to True.
            x_chunksize: The chunk size in the x dimension for writing to zarr.
                Defaults to XY_CHUNKSIZE.
            y_chunksize: The chunk size in the y dimension for writing to zarr.
                Defaults to XY_CHUNKSIZE.
            time_chunksize: The chunk size in the time dimension for writing to zarr.
                Defaults to 1.
            **kwargs: Additional keyword arguments to pass to write_zarr.

        Returns:
            The data array that was set. If write=True, this is the data array read from
            disk, so it is not the same object as the input data array.
        """
        assert isinstance(da, xr.DataArray)

        if write:
            self.logger.info(f"Write {name}")

            fp: Path = Path("other") / (name + ".zarr")
            self.files["other"][name] = fp

            fp_with_root: Path = Path(self.root, fp)

            da: xr.DataArray = write_zarr(
                da,
                fp_with_root,
                x_chunksize=x_chunksize,
                y_chunksize=y_chunksize,
                time_chunksize=time_chunksize,
                crs=da.rio.crs,
                **kwargs,
            )
            self.other[name] = fp_with_root
        return self.other[name]

    def _set_grid(
        self,
        grid_name: str,
        grid: xr.Dataset,
        data: xr.DataArray,
        name: str,
        write: bool,
        x_chunksize: int = XY_CHUNKSIZE,
        y_chunksize: int = XY_CHUNKSIZE,
    ) -> xr.Dataset:
        """Add data to grid dataset.

        All layers of grid must have identical spatial coordinates.

        Args:
            grid_name: name of the grid, e.g. "grid", "subgrid", "region_subgrid"
            grid: the gridded dataset itself
            data: the data to add to the grid
            write: if True, write the data to disk
            name: the name of the layer that will be added to the grid.
            x_chunksize: the chunk size in the x dimension for writing to zarr
            y_chunksize: the chunk size in the y dimension for writing to zarr

        Returns:
            grid: the updated grid with the new layer addedå
        """
        assert isinstance(data, xr.DataArray)

        if name in grid:
            self.logger.warning(f"Replacing grid map: {name}")
            grid = grid.drop_vars(name)

        if len(grid) == 0:
            assert name == "mask", "First grid layer must be mask"
            assert data.dtype == bool
        else:
            assert np.array_equal(data.x.values, grid.x.values)
            assert np.array_equal(data.y.values, grid.y.values)

            # when updating, it is possible that the mask already exists.
            if name != "mask":
                # if the mask exists, mask the data, saving some valuable space on disk
                data_ = xr.where(
                    ~grid["mask"],
                    data,
                    data.attrs["_FillValue"] if data.dtype != bool else False,
                    keep_attrs=True,
                )
                # depending on the order of the dimensions, xr.where may change the dimension order
                # so we change it back
                data = data_.transpose(*data.dims)

        if write:
            fn = Path(grid_name) / (name + ".zarr")
            self.logger.info(f"Writing file {fn}")
            data = write_zarr(
                data,
                path=self.root / fn,
                x_chunksize=x_chunksize,
                y_chunksize=y_chunksize,
                crs=4326,
            )
            self.files[grid_name][name] = Path(grid_name) / (name + ".zarr")

        grid[name] = data
        return grid

    def set_grid(
        self, data: xr.DataArray, name: str, write: bool = True
    ) -> xr.DataArray:
        """Set a new grid layer.

        When the first layer is added to the grid, it must be the mask layer.
        This layer is used to define the spatial extent of the grid and set
        the active cells.

        Args:
            data: The data to add to the grid. Must have the same spatial coordinates
            name: The name of the layer to add to the grid.
            write: If True, write the data to disk. Defaults to True.

        Returns:
            The added grid layer. The returned layer is read from disk if write=True, so
            it is not the same object as the input data.
        """
        self._set_grid("grid", self.grid, data, write=write, name=name)
        return self.grid[name]

    def set_subgrid(
        self, data: xr.DataArray, name: str, write: bool = True
    ) -> xr.DataArray:
        """Set a new subgrid layer.

        When the first layer is added to the subgrid, it must be the mask layer.
        This layer is used to define the spatial extent of the subgrid and set
        the active cells.

        Args:
            data: The data to add to the subgrid. Must have the same spatial coordinates
                as the existing subgrid.
            name: The name of the layer to add to the subgrid.
            write: If True, write the data to disk. Defaults to True.

        Returns:
            The added subgrid layer. The returned layer is read from disk if write=True, so
            it is not the same object as the input data.
        """
        self.subgrid = self._set_grid(
            "subgrid", self.subgrid, data, write=write, name=name
        )
        return self.subgrid[name]

    def set_region_subgrid(
        self, data: xr.DataArray, name: str, write: bool = True
    ) -> xr.DataArray:
        """Set a new region subgrid layer.

        When the first layer is added to the region subgrid, it must be the mask layer.
        This layer is used to define the spatial extent of the region subgrid and set
        the active cells.

        Args:
            data: The data to add to the region subgrid. Must have the same spatial coordinates
                as the existing region subgrid.
            name: The name of the layer to add to the region subgrid.
            write: If True, write the data to disk. Defaults to True.

        Returns:
            The added region subgrid layer. The returned layer is read from disk if write=True, so
            it is not the same object as the input data.
        """
        self.region_subgrid = self._set_grid(
            "region_subgrid",
            self.region_subgrid,
            data,
            write=write,
            name=name,
        )
        return self.region_subgrid[name]

    @property
    def subgrid_factor(self) -> int:
        """The factor by which the subgrid is smaller than the original grid.

        Returns:
            The subgrid factor as an integer.
        """
        subgrid_factor: int = self.subgrid.sizes["x"] // self.grid.sizes["x"]
        assert subgrid_factor == self.subgrid.sizes["y"] // self.grid.sizes["y"]
        return subgrid_factor

    @property
    def ldd_scale_factor(self) -> int:
        """The factor by which the original D8 flow directions are higher resolution than the model grid.

        Returns:
            The scale factor as an integer.
        """
        scale_factor: int = (
            self.other["drainage/original_d8_flow_directions"].shape[0]
            // self.grid["mask"].shape[0]
        )
        assert scale_factor == (
            self.other["drainage/original_d8_flow_directions"].shape[1]
            // self.grid["mask"].shape[1]
        )
        return scale_factor

    @property
    def region(self) -> gpd.GeoDataFrame:
        """Get the region geometry.

        Returns:
            The region geometry as a GeoDataFrame.
        """
        return self.geom["mask"]

    @property
    def bounds(self) -> tuple[float, float, float, float]:
        """Get the bounds of the model grid.

        Returns:
            The bounds of the model grid as a tuple (minx, miny, maxx, maxy).
        """
        return self.grid.rio.bounds(recalc=True)

    @property
    def root(self) -> Path:
        """Get the root directory for the model.

        Returns:
            The root directory path.
        """
        return self._root

    @root.setter
    def root(self, root: Path) -> None:
        """Set the root directory for the model.

        Makes the path absolute.

        Args:
            root: The root directory path.
        """
        self._root = Path(root).absolute()
        self._root.mkdir(parents=True, exist_ok=True)

    @property
    def report_dir(self) -> Path:
        """Directory to save reports to for logging and checking.

        Returns:
            Directory path
        """
        path: Path = Path(self.root).parent / "output" / "build"
        path.mkdir(parents=True, exist_ok=True)
        return path

    def full_like(
        self,
        data: xr.DataArray,
        fill_value: int | float | bool,
        nodata: int | float | bool | None,
        attrs: dict | None = None,
        **kwargs: Any,
    ) -> xr.DataArray:
        """Create a DataArray full of a specified value, with the same shape and coordinates as another DataArray.

        Args:
            data: The DataArray to use as a template.
            fill_value: The value to fill the new DataArray with.
            nodata: The value to use for no data in the new DataArray.
            attrs: Optional dictionary of attributes to set on the new DataArray.
            *args: Additional positional arguments to pass to xr.full_like.
            **kwargs: Additional keyword arguments to pass to xr.full_like.

        Returns:
            A new DataArray with the same shape and coordinates as `data`, filled with `fill_value` and with `nodata` set.
        """
        return full_like(
            data=data,
            fill_value=fill_value,
            nodata=nodata,
            attrs=attrs,
            **kwargs,
        )

    def run_method(self, method: str, *args: Any, **kwargs: Any) -> None:
        """Log method parameters before running a method.

        Args:
            method: The name of the method to run.
            *args: Positional arguments to pass to the method.
            **kwargs: Keyword arguments to pass to the method.

        """
        func: Callable = getattr(self, method)
        signature = inspect.signature(func)
        # combine user and default options
        params: dict[str, Any] = {}
        for i, (k, v) in enumerate(signature.parameters.items()):
            if k in ["args", "kwargs"]:
                if k == "args":
                    params[k] = args[i:]
                else:
                    params.update(**kwargs)
            else:
                v: Any = kwargs.get(k, v.default)
                if len(args) > i:
                    v: Any = args[i]

                params[k] = v

        # call the method
        func(*args, **kwargs)

    def run_methods(
        self,
        methods: dict[str, Any],
        validate_order: bool = True,
        record_progress: bool = False,
        continue_: bool = False,
    ) -> None:
        """Run methods in the order specified in the methods dictionary.

        Args:
            methods: A dictionary with method names as keys and their parameters as values.
            validate_order: If True, validate the order of methods using the build_method decorator.
            record_progress: If True, record progress after each method.
            continue_: Continue previous build if it was interrupted or failed.

        Raises:
            ValueError: If continuing a build and completed methods are not in the methods to run
                or if the order is incorrect.
        """
        # then loop over other methods
        # TODO: Allow validate order for custom models
        build_method.validate_methods(methods, validate_order=validate_order)
        self.files = self.read_or_create_file_library()

        completed_methods: list[str] = (
            build_method.read_progress(self.progress_path) if continue_ else []
        )

        # check if all completed methods are in the methods to run and if order is correct
        if continue_:
            methods_to_run = list(methods.keys())
            for i, completed_method in enumerate(completed_methods):
                if completed_method not in methods_to_run:
                    raise ValueError(
                        f"Cannot continue build: completed method {completed_method} not in methods to run. Restore the method or start a new build."
                    )
                if completed_method != methods_to_run[i]:
                    raise ValueError(
                        f"Cannot continue build: completed method {completed_method} is out of order. Restore the method order or start a new build."
                    )

        for method in methods:
            if method in completed_methods:
                self.logger.info(f"Skipping already completed method: {method}")
                continue

            kwargs = {} if methods[method] is None else methods[method]
            self.run_method(method, **kwargs)
            self.write_file_library()

            if record_progress:
                build_method.record_progress(
                    self.progress_path,
                    method,
                )

        self.logger.info("Finished!")

        build_method.log_time_taken()

    def build(
        self, region: dict, methods: dict[str, dict[str, Any]], continue_: bool
    ) -> None:
        """Build the model with the specified region and methods.

        Args:
            region: A dictionary defining the region to build the model for.
            methods: A dictionary with method names as keys and their parameters as values.
            continue_: Continue previous build if it was interrupted or failed.

        Raises:
            ValueError: If "setup_region" is not in methods when building a new model.
            ValueError: If continue is requested but the code version does not match the build version.
        """
        methods: dict[str, dict[str, Any]] = methods or {}
        if "setup_region" not in methods:
            raise ValueError(
                '"setup_region" must be present in methods when building a new model.'
            )
        methods["setup_region"].update(region=region)

        build_method.check_required_methods(methods.keys())

        # if not continuing, remove existing files path
        if continue_:
            if not self.version_is_current():
                raise ValueError(
                    "Cannot continue build: version mismatch. The version of the existing build is different from the current version of the code. This likely means that the code was updated since the last build. To continue, either restore the old version of the code or start a new build."
                )

            self.read()
        else:
            # for new build, remove existing files path and progress file
            self.files_path.unlink(missing_ok=True)
            self.progress_path.unlink(missing_ok=True)

            # Fresh build so remove the version file and set
            # to current version.
            self.version_path.unlink(missing_ok=True)
            self.set_current_version()

        self.run_methods(
            methods,
            validate_order=True and type(self) is GEBModel,
            record_progress=True,
            continue_=continue_,
        )

        self.build_complete_path.write_text("Build complete")

    def update(
        self,
        methods: dict,
    ) -> None:
        """This function updates an already existing model by running the specified methods.

        Args:
            methods: A dictionary with method names as keys and their parameters as values.

        Raises:
            ValueError: If "setup_region" is in methods, as this can only be called when
                building a new model.
        """
        self.read()

        methods = methods or {}

        if "setup_region" in methods:
            raise ValueError(
                '"setup_region" can only be called when starting a new model.'
            )

        self.run_methods(methods, validate_order=False, record_progress=False)
