"""This module contains classes and functions to build and run SFINCS models for flood hazard assessment.

The main class is `SFINCSRootModel`, which is used to create and manage SFINCS models.
It provides methods to build the model, set up simulations with different forcing methods,
and read simulation results.

"""

from __future__ import annotations

import logging
import math
import shutil
from datetime import datetime
from pathlib import Path
from typing import Any, Literal

import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import numpy.typing as npt
import pandas as pd
import pyflwdir
import rasterio
import rasterio.features
import xarray as xr
from hydromt_sfincs import SfincsModel
from hydromt_sfincs.workflows import burn_river_rect
from pyflwdir import FlwdirRaster
from pyflwdir.dem import fill_depressions
from pyproj import CRS
from rioxarray.merge import merge_arrays
from scipy.ndimage import binary_dilation, value_indices
from shapely import line_locate_point
from shapely.geometry import GeometryCollection, LineString, MultiPoint, Point
from shapely.ops import nearest_points
from tqdm import tqdm

from geb.geb_types import (
    ArrayInt64,
    TwoDArrayBool,
    TwoDArrayFloat32,
    TwoDArrayFloat64,
    TwoDArrayInt32,
)
from geb.hydrology.routing import get_river_width
from geb.workflows.extreme_value_analysis import ReturnPeriodModel
from geb.workflows.io import (
    create_hash_from_parameters,
    get_window,
    read_geom,
    read_hash,
    read_zarr,
    write_geom,
    write_hash,
    write_zarr,
)
from geb.workflows.methods import get_utm_zone
from geb.workflows.raster import (
    calculate_cell_area_m2,
    clip_region,
    coord_to_pixel,
    pad_to_grid_alignment,
    pad_xy,
    rasterize_like,
)

from .workflows import get_river_depth, get_river_manning
from .workflows.outflow import create_outflow_in_mask
from .workflows.utils import (
    create_hourly_hydrograph,
    export_rivers,
    get_discharge_and_river_parameters_by_river,
    get_representative_river_points,
    get_start_point,
    make_relative_paths,
    read_flood_depth,
    run_sfincs_simulation,
    select_most_upstream_point,
    to_sfincs_datetime,
)

SFINCS_WATER_LEVEL_BOUNDARY = 2


class SFINCSRootModel:
    """Builds and updates SFINCS model files for flood hazard modeling."""

    def __init__(self, root: Path, name: str) -> None:
        """Initializes the SFINCSRootModel with a GEBModel and event name.

        Sets up the constant parts of the model (grid, mask, rivers, etc.),
        and has methods to create simulations with actual forcing.

        Args:
            root: The root directory for all SFINCS models.
            name: A string representing the name of the event (e.g., "flood_event_2023").
                Also used to create the path to write the file to disk.
        """
        self._name: str = name
        self._root = root
        self._was_build = None

    @property
    def name(self) -> str:
        """Gets the name of the SFINCS model.

        Returns:
            The name of the SFINCS model.
        """
        return self._name

    @property
    def path(self) -> Path:
        """Gets the root directory for the SFINCS model files.

        Returns:
            The path to the SFINCS model root directory.
        """
        folder: Path = self._root / self.name
        folder.mkdir(parents=True, exist_ok=True)
        return folder

    @property
    def figures_path(self) -> Path:
        """Gets the directory for the SFINCS model diagnostic figures.

        Returns:
            The path to the SFINCS model figures directory.
        """
        folder: Path = self.path / "figures"
        folder.mkdir(parents=True, exist_ok=True)
        return folder

    def exists(self) -> bool:
        """Checks if the SFINCS model already exists in the model root directory.

        Returns:
            True if the SFINCS model exists, False otherwise.
        """
        return Path(self.path / "sfincs.inp").is_file()

    def read(self) -> SFINCSRootModel:
        """Reads an existing SFINCS model from the model root directory.

        Returns:
            The SFINCSRootModel instance with the read model.

        Raises:
            FileNotFoundError: if the SFINCS model does not exist in the specified path.
        """
        if not self.exists():
            raise FileNotFoundError(f"SFINCS model not found in {self.path}")
        self.sfincs_model = SfincsModel(root=str(self.path), mode="r")
        self.sfincs_model.read()

        self.rivers: gpd.GeoDataFrame = read_geom(self.path / "rivers.geoparquet")
        self.subbasins: gpd.GeoDataFrame = read_geom(self.path / "subbasins.geoparquet")

        self.area_of_interest = read_zarr(self.path / "area_of_interest.zarr")

        self.is_build_from_scratch = False
        return self

    def build(
        self,
        DEMs: list[dict[str, str | Path | xr.DataArray | xr.Dataset]],
        subbasins: gpd.GeoDataFrame,
        rivers: gpd.GeoDataFrame,
        discharge: xr.DataArray,
        river_width_alpha: npt.NDArray[np.float32],
        river_width_beta: npt.NDArray[np.float32],
        mannings: xr.DataArray,
        subgrid: bool,
        grid_size_multiplier: int,
        depth_calculation_method: str,
        depth_calculation_parameters: dict[str, float | int] | None = None,
        coastal: bool = False,
        low_elevation_coastal_zone_mask: gpd.GeoDataFrame | None = None,
        coastal_boundary_exclude_mask: gpd.GeoDataFrame | None = None,
        setup_river_outflow_boundary: bool = True,
        outflow_boundary_width_m: float = 500.0,
        initial_water_level: float | None = 0.0,
        custom_rivers_to_burn: gpd.GeoDataFrame | None = None,
        overwrite: bool | Literal["auto"] = True,
        p_value_threshold: float = 0.05,
        selection_strategy: str = "first_significant",
        fixed_shape: float | None = 0.0,
        write_figures: bool = False,
    ) -> SFINCSRootModel:
        """Build a SFINCS model.

        Args:
            DEMs: List of DEM datasets to use for the model. Should be a list of dictionaries with 'path' and 'name' keys.
            subbasins: A GeoDataFrame defining the subbasins of interest.
            rivers: A GeoDataFrame containing river segments.
            discharge: An xarray DataArray containing discharge values for the rivers in m^3/s.
            river_width_alpha: An numpy array of river width alpha parameters. Used for calculating river width.
            river_width_beta: An numpy array of river width beta parameters. Used for calculating river width
            mannings: A xarray DataArray of Manning's n values for the rivers.
            grid_size_multiplier: The number of grid cells to combine when setting up the model grid.
                if 1, no combining is done.
                if 2, every 2x2 grid cells are combined into one cell, etc.
            subgrid: Whether to set up subgrid pixels for the model.
            depth_calculation_method: The method to use for calculating river depth. Can be 'manning' or 'power_law'.
            depth_calculation_parameters: A dictionary of parameters for the depth calculation method. Only used if
                depth_calculation_method is 'power_law', in which case it should contain 'c' and 'd' keys.
            coastal: Whether to set up coastal boundary conditions. Defaults to False.
            low_elevation_coastal_zone_mask: A GeoDataFrame defining the low elevation coastal zone to set as active cells.
            coastal_boundary_exclude_mask: A GeoDataFrame defining areas to exclude from the coastal boundary condition cells.
            setup_river_outflow_boundary: Whether to set up an outflow boundary condition. Defaults to True. Mostly used for testing purposes.
            outflow_boundary_width_m: The width of the outflow boundary condition in meters. Defaults to 500 m.
            initial_water_level: The initial water level to initiate the model. SFINCS fills all cells below this level with water.
            custom_rivers_to_burn: A GeoDataFrame of custom rivers to burn into the model grid. If None, uses the provided rivers GeoDataFrame.
                dataframe must contain 'width' and 'depth' columns.
            overwrite: Whether to overwrite the existing model if it exists. If 'auto', the model is only rebuilt if the input parameters or code have changed.
            p_value_threshold: Anderson-Darling p-value threshold for threshold selection. Defaults to 0.05.
            selection_strategy: Strategy for selecting the best threshold.
                'first_significant': Selects the first threshold (ordered high-to-low) with p_ad > p_value_threshold.
                'best_fit': Evaluates all thresholds and selects the one with the highest p-value.
            fixed_shape: Value to fix the shape parameter (xi) of the GPD. Set to 0.0 to force an Exponential (Gumbel) tail, or null to allow it to be fitted. Defaults to 0.0.
            write_figures: Whether to generate and save diagnostic figures.

        Returns:
            The SFINCSRootModel instance with the built model.

        Raises:
            ValueError: if depth_calculation_method is not 'manning' or 'power_law',
            ValueError: if grid_size_multiplier is not a positive integer.
            ValueError: if resolution of DEM is not square pixels.
        """
        # if overwrite is True, always rebuild the model
        if overwrite is True:
            print("Building new SFINCS model...")
            self.cleanup()
        # determine if we need to rebuild the model based on hash of parameters and
        # code
        elif overwrite == "auto":
            parameters: dict[str, Any] = locals()
            parameters.pop("self")
            hash_file: Path = self.path / "model.hash"
            current_hash: str = create_hash_from_parameters(
                parameters, code_path=Path(__file__)
            )
            if hash_file.exists() and (self.path / "sfincs.inp").exists():
                previous_hash = read_hash(hash_file)
                if previous_hash == current_hash:
                    print("SFINCS model and code unchanged, reading existing model...")
                    return self.read()
                else:
                    print("SFINCS model or code changed, rebuilding model...")
            else:
                print("No existing hash file, building model...")

            self.cleanup()
            self.path.mkdir(parents=True, exist_ok=True)
            write_hash(hash_file, current_hash)

        # if overwrite is False and model exists, read existing model
        elif overwrite is False and self.exists():
            print("Overwrite is False and model exists, reading existing model...")
            return self.read()

        if not isinstance(grid_size_multiplier, int) or grid_size_multiplier <= 0:
            raise ValueError("grid_size_multiplier must be a positive integer")
        if grid_size_multiplier == 1 and subgrid:
            raise ValueError(
                "Cannot use subgrid pixels when grid_size_multiplier is 1 (no aggregation)"
            )

        assert depth_calculation_method in [
            "manning",
            "power_law",
        ], "Method should be 'manning' or 'power_law'"

        self.outflow_boundary_width_m = outflow_boundary_width_m

        mask: xr.DataArray = self.create_mask(
            DEMs=DEMs,
            subbasins=subbasins,
            grid_size_multiplier=grid_size_multiplier,
        )

        subbasins: gpd.GeoDataFrame = subbasins.to_crs(mask.rio.crs)
        rivers: gpd.GeoDataFrame = rivers.to_crs(mask.rio.crs)

        print("Starting SFINCS model build...")
        # build base model
        sf: SfincsModel = SfincsModel(root=str(self.path), mode="w+", write_gis=True)
        self.sfincs_model = sf

        # Configure the model's logger and all handlers to DEBUG level
        # This must be done AFTER creating SfincsModel when handlers are initialized
        sf.logger.setLevel(logging.DEBUG)
        for handler in sf.logger.handlers:
            handler.setLevel(logging.DEBUG)

        # we need this later to exclude these subbasins from the area of interest
        subbasins_of_interest = subbasins[
            (~subbasins["is_downstream_outflow"].astype(bool))
        ].index.tolist()

        self.subbasins, subbasins_burned = self.burn_and_align_subbasins(
            subbasins=subbasins, mask=mask
        )
        self.subbasins.to_parquet(self.path / "subbasins_processed.geoparquet")

        # delete mask and subbasins to ensure we don't accidentally use them later
        del mask
        del subbasins

        resolution: tuple[float, float] = subbasins_burned.rio.resolution()
        if abs(abs(resolution[0]) - abs(resolution[1])) > 1e-8:
            raise ValueError("DEM resolution must be square pixels")

        geotransformation = subbasins_burned.rio.transform(recalc=True)
        crs = subbasins_burned.rio.crs
        if not crs.is_epsg_code:
            raise ValueError("CRS must be an EPSG code")

        sf.setup_grid(
            x0=geotransformation.c,
            y0=geotransformation.f,
            dx=geotransformation.a,
            dy=geotransformation.e,
            mmax=subbasins_burned.rio.width,
            nmax=subbasins_burned.rio.height,
            rotation=0.0,
            epsg=crs.to_epsg(),
        )
        if crs.is_geographic:
            sf.setup_config(crsgeo=1)

        # Copy DEMs to avoid modifying input
        DEMs: list[dict[str, str | Path | xr.DataArray | xr.Dataset]] = [
            DEM.copy() for DEM in DEMs
        ]

        # Remove unnecessary keys (to avoid warnings) and set reproj_method to bilinear for all DEMs
        for DEM in DEMs:
            DEM.pop("name", None)
            DEM.pop("path", None)
            DEM.pop("fill_depressions", None)
            DEM["reproj_method"] = "bilinear"

        sf.setup_dep(datasets_dep=DEMs)

        # Remove rivers that are not represented in the grid and have no upstream rivers
        # TODO: Make an upstream flag in preprocessing for upstream rivers that is more
        # general than the MERIT-hydro specific 'maxup' attribute
        self.rivers: gpd.GeoDataFrame = rivers[
            (rivers["maxup"] > 0) | (rivers["represented_in_grid"])
        ].to_crs(self.crs)
        del rivers

        flood_plain: gpd.GeoDataFrame = self.get_flood_plain()
        sf.setup_mask_active(flood_plain, reset_mask=True)

        self.rivers["outflow_elevation"] = np.nan
        self.rivers["outflow_point_xy"] = None
        self.rivers["is_outflow_boundary"] = False

        if coastal:
            # set zsini based on the minimum elevation
            assert isinstance(sf.config, dict)
            sf.config["zsini"] = initial_water_level  # ty: ignore[invalid-assignment]

            if not isinstance(coastal_boundary_exclude_mask, gpd.GeoDataFrame):
                raise ValueError(
                    "coastal_boundary_exclude_mask must be provided when coastal is True"
                )

            # activate low elevation coastal zone
            if not isinstance(low_elevation_coastal_zone_mask, gpd.GeoDataFrame):
                raise ValueError(
                    "low_elevation_coastal_zone_mask must be provided when coastal is True"
                )
            sf.setup_mask_active(
                include_mask=low_elevation_coastal_zone_mask,
                mask_buffer=5000,
                reset_mask=False,
            )

            # setup the coastal boundary conditions
            sf.setup_mask_bounds(
                btype="waterlevel",
                zmax=3,  # maximum elevation for valid boundary cells
                exclude_mask=coastal_boundary_exclude_mask,
                all_touched=True,
            )
        else:
            self.calculate_outflow_conditions(area=flood_plain)
            if setup_river_outflow_boundary:
                outflow_points = self.rivers[self.rivers["is_outflow_boundary"]][
                    "outflow_point_xy"
                ]
                outflow_points = gpd.GeoDataFrame(
                    geometry=gpd.points_from_xy(
                        outflow_points.apply(lambda xy: xy[0]),
                        outflow_points.apply(lambda xy: xy[1]),
                    ),
                    crs=self.crs,
                )
                outflow_points = outflow_points.to_crs(self.utm_zone)
                outflow_points.geometry = outflow_points.buffer(
                    (self.outflow_boundary_width_m / 2) * 1.1
                ).to_crs(self.crs)

                sf.setup_mask_active(
                    include_mask=outflow_points,
                    all_touched=False,
                    reset_mask=False,
                )

                # create polygon from mask
                area = list(
                    rasterio.features.shapes(
                        self.mask,
                        transform=self.mask.rio.transform(),
                        connectivity=8,
                        mask=self.mask,
                    )
                )
                area = gpd.GeoDataFrame.from_features(
                    [
                        {"geometry": geom, "properties": {"value": int(value)}}
                        for geom, value in area
                    ],
                    crs=self.mask.rio.crs,
                )
                if len(area) > 1:
                    # remove isolated cells by keeping only the largest polygon
                    area["area"] = area.geometry.area
                    area = area.sort_values("area", ascending=False).iloc[0:1]
                    area = area.drop(columns=["area"])
                    # raise ValueError(
                    #     "Calculated outflow area is not a single polygon, cannot calculate outflow conditions"
                    # )

                self.calculate_outflow_conditions(area=area)

                # must be performed BEFORE burning rivers.
                self.setup_river_outflow_boundary()

        # We are only interested in flooding within the subbasins of interest.
        # This area of interest mask is used later to filter the flood maps.
        self.area_of_interest: xr.DataArray = self.mask.astype(bool)
        self.area_of_interest.values[:] = False
        self.area_of_interest.values[
            np.isin(subbasins_burned, subbasins_of_interest)
        ] = True

        # In many cases due to the coarse subbasin resolution compared to the DEM resolution,
        # the area of interest does not align perfectly with the shape of the subbasins at the
        # DEM resolution, ultimately leading to small gaps in the flood maps. We now temporarily
        # fix this by dilating the area of interest mask. A better solution should be found.
        # TODO: Find a better solution for the area of interest alignment issue.
        self.area_of_interest.values = binary_dilation(
            self.area_of_interest.values, iterations=10
        )
        self.area_of_interest.attrs["_FillValue"] = None

        self.plot_rivers()

        # do the river burning only if there are active rivers in the model
        # or if custom rivers to burn are provided
        if not self.active_rivers.empty or (
            custom_rivers_to_burn is not None and not custom_rivers_to_burn.empty
        ):
            river_representative_points = []
            for ID in self.active_rivers.index:
                river_representative_points.append(
                    get_representative_river_points(ID, self.rivers)
                )

            discharge_by_river, river_parameters = (
                get_discharge_and_river_parameters_by_river(
                    self.active_rivers.index.tolist(),
                    river_representative_points,
                    discharge=discharge,
                    river_width_alpha=river_width_alpha,
                    river_width_beta=river_width_beta,
                )
            )

            # resample to daily frequency because this is the frequency for the river
            # width and depth estimation formulas.
            discharge_by_river = discharge_by_river.resample("D", label="left").mean()

            if custom_rivers_to_burn is not None:
                rivers_to_burn = custom_rivers_to_burn.to_crs(sf.crs)
                if "width" not in rivers_to_burn.columns:
                    raise ValueError(
                        "Custom rivers to burn must have a 'width' column when using custom rivers"
                    )
                if "depth" not in rivers_to_burn.columns:
                    raise ValueError(
                        "Custom rivers to burn must have a 'depth' column when using custom rivers"
                    )
            else:
                rivers_to_burn = self.assign_return_periods(
                    self.active_rivers,
                    discharge_by_river,
                    return_periods=[2],
                    p_value_threshold=p_value_threshold,
                    fixed_shape=fixed_shape,
                    selection_strategy=selection_strategy,
                    write_figures=write_figures,
                    output_directory=self.figures_path / "bankfull_estimation",
                )

                river_width_unknown_mask = rivers_to_burn["width"].isnull()

                rivers_to_burn.loc[river_width_unknown_mask, "width"] = get_river_width(
                    river_parameters["river_width_alpha"][river_width_unknown_mask],
                    river_parameters["river_width_beta"][river_width_unknown_mask],
                    rivers_to_burn.loc[river_width_unknown_mask, "Q_2"],
                ).astype(np.float64)

                rivers_to_burn["depth"] = get_river_depth(
                    rivers_to_burn,
                    method=depth_calculation_method,
                    parameters=depth_calculation_parameters,
                    bankfull_column="Q_2",
                )

            rivers_to_burn["manning"] = get_river_manning(rivers_to_burn)

            rivers_to_burn.to_parquet(
                self.path / "rivers_with_widths_and_depths.geoparquet"
            )

            # Because hydromt-sfincs does a lot of filling default values when data
            # is missing, we need to be extra sure that the required columns are
            # present and contain valid data.
            assert rivers_to_burn["width"].notnull().all(), "River width cannot be null"
            assert rivers_to_burn["depth"].notnull().all(), "River depth cannot be null"
            assert rivers_to_burn["manning"].notnull().all(), (
                "River Manning's n cannot be null"
            )

            # only burn rivers that are wider than the grid size
            rivers_to_burn: gpd.GeoDataFrame = rivers_to_burn[
                rivers_to_burn["width"] > self.estimated_cell_size_m
            ]

            rivers_to_burn.to_parquet(self.path / "rivers_to_burn.geoparquet")
        else:
            rivers_to_burn = None

        # if sfincs is run with subgrid, we set up the subgrid, with burned in rivers and mannings
        # roughness within the subgrid. If not, we burn the rivers directly into the main grid,
        # including mannings roughness.
        if subgrid:
            print(
                f"Setting up SFINCS subgrid with {grid_size_multiplier} subgrid pixels..."
            )
            # only burn rivers that are wider than the subgrid pixel size
            sf.setup_subgrid(
                datasets_dep=DEMs,
                datasets_rgh=[
                    {
                        "manning": mannings.to_dataset(name="manning"),
                    }
                ],
                datasets_riv=[
                    {
                        "centerlines": rivers_to_burn.rename(
                            columns={"width": "rivwth", "depth": "rivdph"}
                        )
                    }
                ]
                if rivers_to_burn is not None
                else [],
                write_dep_tif=True,
                write_man_tif=True,
                nr_subgrid_pixels=grid_size_multiplier,
                nlevels=20,
                nrmax=500,
            )

            sf.write_subgrid()
        else:
            print(
                "Setting up SFINCS without subgrid - burning rivers into main grid..."
            )
            # first set up the mannings roughness with the default method
            # (we already have the DEM set up)
            sf.setup_manning_roughness(
                datasets_rgh=[
                    {
                        "manning": mannings.to_dataset(name="manning"),
                    }
                ]
            )

            if rivers_to_burn is not None:
                # retrieve the elevation and mannings grids
                # burn the rivers into these grids
                elevation, mannings = burn_river_rect(
                    da_elv=sf.grid.dep,
                    gdf_riv=rivers_to_burn,
                    da_man=sf.grid.manning,
                    rivwth_name="width",
                    rivdph_name="depth",
                    manning_name="manning",
                    segment_length=90.0,
                )
                # set the modified grids back to the model
                sf.set_grid(elevation, name="dep")
                sf.set_grid(mannings, name="manning")

        # write all components, except forcing which must be done after the model building
        sf.write_grid()
        sf.write_subgrid()
        sf.write_geoms()
        sf.write_states()
        sf.write_config()

        self.subbasins.to_parquet(self.path / "subbasins.geoparquet")
        self.rivers.to_parquet(self.path / "rivers.geoparquet")

        write_zarr(
            self.area_of_interest,
            self.path / "area_of_interest.zarr",
            crs=self.mask.rio.crs,
        )

        sf.plot_basemap(
            fn_out=str(self.figures_path / "basemap.png"),
            vmin=math.floor(self.elevation.min()),
            vmax=max(
                math.ceil(self.elevation.max()), 1
            ),  # vmax is required until bug in hydromt-sfincs fixed, see: https://github.com/Deltares/hydromt_sfincs/issues/324
        )

        self.is_build_from_scratch = True

        return self

    def create_mask(
        self,
        DEMs: list[dict[str, str | Path | xr.DataArray | xr.Dataset]],
        subbasins: gpd.GeoDataFrame,
        grid_size_multiplier: int,
    ) -> xr.DataArray:
        """Creates the model mask based on the provided DEMs and subbasins.

        The mask is created by taking the first DEM dataset and setting all cells
        within the subbasins to True. If grid_size_multiplier is greater than 1,
        the mask is aggregated accordingly.

        It is automatically aligned to DEM grid, and padded in such a way that the
        grid cells the corners are multiples of grid_size_multiplier and the resolution.

        Args:
            DEMs: List of DEM datasets to use for the model. Should be a list of dictionaries with 'path' and 'name' keys.
            subbasins: A GeoDataFrame defining the subbasins of interest.
            grid_size_multiplier: The number of grid cells to combine when setting up the model grid.
                if 1, no combining is done.

        Returns:
            The created model mask as an xarray DataArray.
        """
        mask_ds = DEMs[0]["elevtn"]
        assert isinstance(mask_ds, xr.Dataset)

        mask: xr.DataArray = mask_ds["elevtn"].astype(bool)
        mask.attrs["_FillValue"] = None  # ensure no fill value
        mask.name = "mask"

        # for in case the first DEM does not cover the entire area, we first pad the mask to the subbasins bounds
        # added a small buffer to ensure we don't have issues with rounding errors in the bounds and the mask alignment
        minx, miny, maxx, maxy = (
            subbasins.buffer(0.01).to_crs(mask.rio.crs).total_bounds
        )
        mask = pad_xy(
            mask,
            minx=minx,
            miny=miny,
            maxx=maxx,
            maxy=maxy,
            constant_values=True,
            return_slice=False,
        )

        # then, for when the DEM is (at places) larger than the subbasins, we clip to the subbasins bounds
        mask: xr.DataArray = mask.isel(
            get_window(
                mask.x,
                mask.y,
                subbasins.to_crs(mask.rio.crs).total_bounds,
                buffer=100,
                raise_on_buffer_out_of_bounds=False,
            )
        )

        # set all values to True
        mask.values[:] = True

        # if the grid size multiplier is greater than 1, we need to aggregate the mask
        # accordingly. We first pad the mask to aligning it to the grid size multiplier,
        # resolution and origin, then coarsen it by the grid size multiplier
        if grid_size_multiplier > 1:
            mask = pad_to_grid_alignment(
                mask,
                grid_size_multiplier=grid_size_multiplier,
                constant_values=False,
            )
            mask = mask.coarsen(
                x=grid_size_multiplier,
                y=grid_size_multiplier,
                boundary="exact",
                coord_func="mean",
            ).reduce(np.max)
            mask.rio.write_transform(inplace=True)  # force recalculation of transform

        return mask

    def burn_and_align_subbasins(
        self, subbasins: gpd.GeoDataFrame, mask: xr.DataArray
    ) -> tuple[gpd.GeoDataFrame, xr.DataArray]:
        """Burns the subbasins into a raster aligned with the model mask.

        To ensure alignment with the model grid, the subbasins are first
        rasterized like the model mask, then vectorized back to polygons.

        Args:
            subbasins: A GeoDataFrame defining the subbasins of interest.
            mask: The model mask as an xarray DataArray.

        Returns:
            A tuple containing:
            - The burned subbasins as a GeoDataFrame.
            - The burned subbasins raster as an xarray DataArray aligned with the model mask.
        """
        subbasins.to_parquet(self.path / "subbasins_input.geoparquet")
        subbasins: gpd.GeoDataFrame = subbasins.to_crs(mask.rio.crs)
        subbasins_burned: xr.DataArray = rasterize_like(
            gdf=subbasins.reset_index(),  # ty:ignore[invalid-argument-type]
            column="COMID",
            raster=mask,
            dtype=np.int32,
            nodata=-1,
            all_touched=True,
        )
        assert isinstance(mask, xr.DataArray)

        subbasin_shapes = list(
            rasterio.features.shapes(
                subbasins_burned.values,
                transform=mask.rio.transform(),
                connectivity=8,
                mask=(subbasins_burned.values != -1),
            )
        )

        subbasin_features: list[dict[str, Any]] = [
            {"geometry": geom, "properties": {"COMID": int(value)}}
            for geom, value in subbasin_shapes
        ]
        subbasins: gpd.GeoDataFrame = gpd.GeoDataFrame.from_features(
            subbasin_features,
            crs=mask.rio.crs,
        ).set_index("COMID")  # ty:ignore[invalid-assignment]

        # combine subbasins with the same COMID
        subbasins = subbasins.dissolve(by=subbasins.index)  # ty:ignore[invalid-argument-type]

        subbasins_burned: xr.DataArray = clip_region(subbasins_burned, align=None)[0]

        # if y axis is descending (usually for geographical grids), flip it
        if subbasins_burned.y[-1] < subbasins_burned.y[0]:
            subbasins_burned: xr.DataArray = subbasins_burned.isel(
                y=slice(None, None, -1)
            )

        return subbasins, subbasins_burned

    def plot_rivers(self) -> None:
        """Plots the rivers and subbasins boundary and saves to file."""
        fig, ax = plt.subplots(figsize=(10, 10))
        self.subbasins.boundary.plot(ax=ax, color="black")

        # in coastal areas, there may be no active rivers
        if not self.active_rivers.empty:
            self.active_rivers.plot(ax=ax, color="blue")
        plt.savefig(self.figures_path / "rivers.png")

    def calculate_outflow_conditions(self, area: gpd.GeoDataFrame) -> None:
        """Calculates outflow elevation and coordinates for all rivers.

        Determines the outflow point for each river (intersection with subbasins boundary
        for downstream-most rivers, or downstream endpoint for internal rivers)
        and retrieves the elevation at that point from the DEM.

        Args:
            area: A GeoDataFrame defining the area to consider for outflow point calculation.
                Typically the subbasins or model area boundary.

        Raises:
            ValueError: if the calculated outflow point is not a single point.
            ValueError: if the calculated outflow point is outside of the model grid.
        """

        def export_diagnostics(
            outflow_point: Point | MultiPoint | GeometryCollection,
        ) -> None:
            write_zarr(
                self.mask,
                self.path / "debug_outflow_mask.zarr",
                crs=self.mask.rio.crs,
            )
            write_geom(self.rivers, self.path / "debug_rivers.geoparquet")
            write_geom(self.subbasins, self.path / "debug_region.geoparquet")

            outflow_gdf = gpd.GeoDataFrame({"geometry": [outflow_point]}, crs=self.crs)
            write_geom(outflow_gdf, self.path / "debug_outflow_point.geoparquet")

        boundary = area.union_all().boundary

        for river_idx, river in self.active_rivers[
            self.active_rivers["is_downstream_outflow"]  # any outflow river
            | (
                self.active_rivers["downstream_ID"] == -1
            )  # or river that flows into ocean
        ].iterrows():
            downstream_river_idx = river["downstream_ID"]
            if downstream_river_idx == -1:
                # continue  # skip rivers that flow into ocean, as they don't have an outflow point on the grid
                river = river.geometry

                # Last point of river
                end_point = Point(river.coords[-1])

                # Compute nearest points between end_point and boundary
                _, nearest_boundary_point = nearest_points(end_point, boundary)

                # Extend river directly to that location
                river = LineString(
                    list(river.coords) + [nearest_boundary_point.coords[0]]
                )

                outflow_point = nearest_boundary_point

                if not isinstance(outflow_point, Point):
                    export_diagnostics(outflow_point)
                    raise ValueError(
                        "Calculated outflow point is not a single point. Please check the river geometries and subbasins boundary."
                    )

                outflow_col, outflow_row = coord_to_pixel(
                    (outflow_point.x, outflow_point.y),
                    self.mask.rio.transform().to_gdal(),
                )

                # due to floating point precision, the intersection point
                # may be just outside the model grid. We therefore check if the
                # point is outside the grid, and if so, move it 1 m upstream along the river
                if not self.mask.values[outflow_row, outflow_col]:
                    # move outflow point 1 m upstream. 0.000008983 degrees is approximately 1 m
                    outflow_point: Point | MultiPoint | GeometryCollection = (
                        river.interpolate(
                            line_locate_point(river, outflow_point) - 0.000008983
                            if self.is_geographic
                            else 1.0
                        )
                    )
                    if not isinstance(outflow_point, Point):
                        # if the intersection is not a single point, select the most upstream point
                        outflow_point: Point = select_most_upstream_point(
                            river, outflow_point
                        )

                    outflow_col, outflow_row = coord_to_pixel(
                        (outflow_point.x, outflow_point.y),
                        self.mask.rio.transform().to_gdal(),
                    )

                    # if still outside the grid, raise error for boundary rivers
                    if not self.mask.values[outflow_row, outflow_col]:
                        export_diagnostics(outflow_point)
                        raise ValueError(
                            "Calculated outflow point is outside of the model grid. Please check the river geometries and subbasins boundary."
                        )

                self.rivers.at[river_idx, "outflow_elevation"] = 0
                self.rivers.at[river_idx, "outflow_point_xy"] = (
                    outflow_point.x,
                    outflow_point.y,
                )
                self.rivers.at[river_idx, "is_outflow_boundary"] = True

            else:
                river = river.geometry
                assert isinstance(river, LineString)
                # extend the river geometry to include all downstream segments
                # up to coastal outflow
                while downstream_river_idx != -1:
                    downstream_river = self.rivers.loc[downstream_river_idx]
                    # combine the geometries to ensure intersection is found
                    river = LineString(
                        river.coords[:] + downstream_river.geometry.coords[:]
                    )
                    assert isinstance(river, LineString)
                    downstream_river_idx = downstream_river["downstream_ID"]

                outflow_point: Point | MultiPoint | GeometryCollection = (
                    river.intersection(boundary)
                )

                # while there is no intersection point, we elongate the river geometry to ensure it intersects with the boundary. \
                # Since there is no downstream river segment, we extend the last segment of the river
                if outflow_point.is_empty and downstream_river_idx == -1:
                    # Last point of river
                    end_point = Point(river.coords[-1])

                    # Compute nearest points between end_point and boundary
                    _, nearest_boundary_point = nearest_points(end_point, boundary)

                    # Extend river directly to that location
                    river = LineString(
                        list(river.coords) + [nearest_boundary_point.coords[0]]
                    )

                    outflow_point = nearest_boundary_point

                if not isinstance(outflow_point, Point):
                    export_diagnostics(outflow_point)
                    # if the intersection is not a single point, select the most upstream point
                    outflow_point: Point = select_most_upstream_point(
                        river, outflow_point
                    )

                outflow_col, outflow_row = coord_to_pixel(
                    (outflow_point.x, outflow_point.y),
                    self.mask.rio.transform().to_gdal(),
                )

                # due to floating point precision, the intersection point
                # may be just outside the model grid. We therefore check if the
                # point is outside the grid, and if so, move it 1 m upstream along the river
                if not self.mask.values[outflow_row, outflow_col]:
                    # move outflow point 1 m upstream. 0.000008983 degrees is approximately 1 m
                    outflow_point: Point | MultiPoint | GeometryCollection = (
                        river.interpolate(
                            line_locate_point(river, outflow_point) - 0.000008983
                            if self.is_geographic
                            else 1.0
                        )
                    )
                    if not isinstance(outflow_point, Point):
                        # if the intersection is not a single point, select the most upstream point
                        outflow_point: Point = select_most_upstream_point(
                            river, outflow_point
                        )

                    outflow_col, outflow_row = coord_to_pixel(
                        (outflow_point.x, outflow_point.y),
                        self.mask.rio.transform().to_gdal(),
                    )

                    # if still outside the grid, raise error for boundary rivers
                    if not self.mask.values[outflow_row, outflow_col]:
                        export_diagnostics(outflow_point)
                        raise ValueError(
                            "Calculated outflow point is outside of the model grid. Please check the river geometries and subbasins boundary."
                        )

                outflow_elevation: float = self.elevation[
                    outflow_row, outflow_col
                ].item()
                self.rivers.at[river_idx, "outflow_elevation"] = outflow_elevation
                self.rivers.at[river_idx, "outflow_point_xy"] = (
                    outflow_point.x,
                    outflow_point.y,
                )
                self.rivers.at[river_idx, "is_outflow_boundary"] = True

    def setup_river_outflow_boundary(
        self,
    ) -> None:
        """Sets up river outflow boundary condition for the SFINCS model.

        Uses pre-calculated outflow points to set up boundary conditions for
        downstream-most rivers.
        """
        rivers_with_outflow_boundary = self.rivers[self.rivers["is_outflow_boundary"]]

        for river_idx, river in rivers_with_outflow_boundary.iterrows():
            outflow_point_xy = river["outflow_point_xy"]

            outflow_col, outflow_row = coord_to_pixel(
                outflow_point_xy,
                self.mask.rio.transform().to_gdal(),
            )

            outflow_mask: TwoDArrayBool = create_outflow_in_mask(
                self.mask.values,
                row=outflow_row,
                col=outflow_col,
                width_cells=(
                    math.ceil(
                        (
                            (self.outflow_boundary_width_m / self.estimated_cell_size_m)
                            - 1
                        )
                        / 2
                    )
                    * 2
                    + 1
                ),
            )

            assert self.sfincs_model.grid_type == "regular"
            self.mask.values[outflow_mask] = SFINCS_WATER_LEVEL_BOUNDARY

    def get_flood_plain(self, maximum_hand: float = 30.0) -> gpd.GeoDataFrame:
        """Returns the flood plain grid of the SFINCS model.

        Uses a two-stage approach:
        1. Calculate the Height Above Nearest Drainage (HAND) using the elevation
           grid and the river network and classify all cells with HAND less than
              maximum_hand as flood plain.
        2. Adds a buffer of minimum distance around rivers.

        Args:
            maximum_hand: The maximum Height Above Nearest Drainage (HAND) value to consider as flood plain, in meters. Default is 30.0 m.

        Returns:
            The flood plain as a GeoDataFrame.
        """
        subbasin_raster = rasterize_like(
            gdf=self.subbasins,
            raster=self.elevation,
            dtype=np.uint8,
            nodata=0,
            burn_value=1,
            all_touched=True,
        ).astype(bool)

        # get a raster of all drainage cells (i.e., 0 HAND). We do so
        # by burning the rivers into a raster
        drainage_cells = rasterize_like(
            gdf=self.active_rivers.to_crs(self.elevation.rio.crs),
            raster=self.elevation,
            dtype=np.uint8,
            nodata=0,
            burn_value=1,
            all_touched=True,
        ).astype(bool)
        drainage_cells.attrs["_FillValue"] = None

        # obtain the elevation grid. Some cells on the outside of the subbasins
        # drain to a location outside the subbasins. In the HAND calculation,
        # these cells are given a very low HAND, and thus would also be part of
        # the floodplain. To avoid this, we set the cells outside the subbasins
        # to a very low elevation, so that cells that drain outside the subbasins
        # get an extremly high HAND value, and thus will not be part of the
        # flood plain.
        elevation = self.elevation.values.copy()
        elevation[~subbasin_raster.values] = -10_000

        # Fill depressions to ensure proper flow direction calculation
        elevation, d8 = fill_depressions(elevation, nodata=np.nan)
        flow_raster = pyflwdir.from_array(
            d8,
            transform=self.elevation.rio.transform(),
            latlon=self.is_geographic,
        )

        # also add all cells with large upstream from the DEM to the drainage cells
        upstream_area: TwoDArrayFloat64 = flow_raster.upstream_area(unit="m2")
        drainage_cells: xr.DataArray = drainage_cells | (
            upstream_area > 25_000_000
        )  # 25 km²

        height_above_nearest_drainage = xr.full_like(
            self.elevation, np.nan, dtype=np.float32
        )
        height_above_nearest_drainage.values = flow_raster.hand(
            drain=drainage_cells.values, elevtn=elevation
        )
        height_above_nearest_drainage = xr.where(
            height_above_nearest_drainage != -9999.0,
            height_above_nearest_drainage,
            np.nan,
        )
        height_above_nearest_drainage = xr.where(
            subbasin_raster,
            height_above_nearest_drainage,
            np.nan,
        )

        flood_plain: xr.DataArray = height_above_nearest_drainage <= maximum_hand

        # convert flood plain raster to vector
        flood_plain_geom = list(
            rasterio.features.shapes(
                flood_plain.astype(np.uint8),
                mask=flood_plain,
                connectivity=8,
                transform=flood_plain.rio.transform(recalc=True),
            ),
        )
        assert len(flood_plain_geom) > 0, "No flood plain geometries found"

        flood_plain_geom = gpd.GeoDataFrame.from_features(
            [{"geometry": geom[0], "properties": {}} for geom in flood_plain_geom],
            crs=flood_plain.rio.crs,
        )

        # add buffer around rivers to ensure connectivity
        river_buffer: gpd.GeoDataFrame = self.active_rivers.to_crs(self.utm_zone)
        river_buffer.geometry = river_buffer.geometry.buffer(
            self.outflow_boundary_width_m
        )
        river_buffer = river_buffer.to_crs(flood_plain_geom.crs).dissolve()

        flood_plain_geom: gpd.GeoDataFrame = pd.concat(
            [
                flood_plain_geom,
                river_buffer,
            ],
            ignore_index=True,
        )  # ty:ignore[invalid-assignment]
        return flood_plain_geom

    @property
    def mask(self) -> xr.DataArray:
        """Returns the mask grid of the SFINCS model.

        Returns:
            The mask grid as an xarray DataArray.
        """
        return self.sfincs_model.grid["msk"]

    @property
    def elevation(self) -> xr.DataArray:
        """Returns the elevation grid of the SFINCS model.

        Returns:
            The elevation grid as an xarray DataArray.
        """
        return self.sfincs_model.grid["dep"]

    @property
    def crs(self) -> rasterio.crs.CRS:
        """Returns the coordinate reference system (CRS) of the SFINCS model.

        Returns:
            The CRS of the SFINCS model.
        """
        return self.elevation.rio.crs

    @property
    def shape(self) -> tuple[int, int]:
        """Returns the shape of the SFINCS model grid.

        Returns:
            The shape of the SFINCS model grid as a tuple (nrows, ncols).
        """
        return self.sfincs_model.grid["dep"].shape

    @property
    def is_geographic(self) -> bool:
        """Returns whether the SFINCS model uses a geographic coordinate system.

        Returns:
            True if the SFINCS model uses a geographic coordinate system, False otherwise.

        Raises:
            ValueError: if the CRS is not defined for the SFINCS model.
        """
        if self.sfincs_model.crs is None:
            raise ValueError("CRS is not defined for the SFINCS model.")
        return self.sfincs_model.crs.is_geographic

    @property
    def cell_area(self) -> float | int | xr.DataArray:
        """Returns the area of a single cell in the SFINCS model grid in square meters.

        Returns:
            The area of a single cell in the SFINCS model grid in m².
        """
        if self.is_geographic:
            cell_area: xr.DataArray = xr.full_like(
                self.elevation, np.nan, dtype=np.float32
            )
            height, width = self.shape
            cell_area.values = calculate_cell_area_m2(
                self.elevation.rio.transform(), height, width
            )
            return cell_area
        else:
            return xr.full_like(
                self.elevation,
                (
                    self.sfincs_model.grid["dep"].rio.resolution()[0]
                    * self.sfincs_model.grid["dep"].rio.resolution()[1]
                ),
                dtype=np.float32,
            )

    @property
    def estimated_cell_size_m(self) -> float:
        """Returns the estimated cell size of the SFINCS model grid in meters.

        For geographic coordinate systems, this is an approximate value based on the average
        cell area. For projected coordinate systems, this is the actual cell size.

        Returns:
            The estimated cell size of the SFINCS model grid in meters.
        """
        if self.is_geographic:
            assert isinstance(self.cell_area, xr.DataArray)
            avg_cell_area = self.cell_area.mean().item()
            estimated_cell_size = np.sqrt(avg_cell_area)
            return estimated_cell_size
        else:
            return math.sqrt(
                self.sfincs_model.grid["msk"].rio.resolution()[0]
                * self.sfincs_model.grid["msk"].rio.resolution()[1]
            )

    @staticmethod
    def assign_return_periods(
        rivers: gpd.GeoDataFrame,
        discharge_dataframe: pd.DataFrame,
        return_periods: list[int | float],
        prefix: str = "Q",
        min_exceed: int = 30,
        nboot: int = 2000,
        fixed_shape: float | None = None,
        fixed_scale: float | None = None,
        p_value_threshold: float = 0.05,
        selection_strategy: str = "first_significant",
        write_figures: bool = False,
        output_directory: Path | None = None,
    ) -> gpd.GeoDataFrame:
        """Assign return periods to rivers using GPD-POT analysis.

        Based on https://doi.org/10.1002/2016WR019426

        Uses Generalized Pareto Distribution Peaks-Over-Threshold method with:
            - Daily maxima resampling from input time series
            - GPD-POT exceedance model fitting above threshold
            - Anderson-Darling bootstrap p-value threshold selection
            - Return level estimation for specified return periods

        Args:
            rivers: GeoDataFrame with river IDs as index that must match columns in discharge_dataframe.
            discharge_dataframe: Time series DataFrame with datetime index containing discharge data for all rivers (m³/s).
            return_periods: List of return periods in years to compute return levels for.
            prefix: Column prefix for output return level columns. Defaults to "Q".
            min_exceed: Minimum number of exceedances required for reliable GPD fit. Defaults to 30.
            nboot: Number of bootstrap samples for Anderson-Darling threshold selection. Defaults to 2000.
            fixed_shape: Value to fix the shape parameter (xi) (dimensionless). If None, it is fitted.
                Set to 0 to force a Gumbel (Exponential) distribution.
            fixed_scale: Value to fix the scale parameter (sigma) (dimensionless). If None, it is fitted.
            p_value_threshold: Anderson-Darling p-value threshold for threshold selection. Defaults to 0.05.
            selection_strategy: Strategy for selecting the best threshold.
                'first_significant': Selects the first threshold (ordered high-to-low) with p_ad > p_value_threshold.
                'best_fit': Evaluates all thresholds and selects the one with the highest p-value.
            write_figures: Whether to generate and save diagnostic figures for each river.
            output_directory: Directory to save the generated figures.

        Returns:
            Updated rivers GeoDataFrame with return level columns added (m³/s).

        Raises:
            TypeError: If discharge series does not have a DatetimeIndex.
            ValueError: If return periods contain non-positive values.
        """
        assert isinstance(return_periods, list)
        if not all((isinstance(T, (int, float)) and T > 0) for T in return_periods):
            raise ValueError("All return periods must be positive numbers (years > 0).")
        return_periods_arr = np.asarray(return_periods, float)

        if write_figures and output_directory:
            output_directory.mkdir(parents=True, exist_ok=True)

        for idx in tqdm(rivers.index, total=len(rivers), desc="GPD-POT Return Periods"):
            discharge = discharge_dataframe[idx].dropna()

            # If all values are zero, assign zeros
            if (discharge < 1e-10).all():
                print(f"Discharge all zero for river {idx}, assigning zeros.")
                for T in return_periods:
                    rivers.loc[idx, f"{prefix}_{T}"] = 0.0
                continue

            if not isinstance(discharge.index, pd.DatetimeIndex):
                raise TypeError(
                    f"Discharge series for river {idx} must have a DatetimeIndex."
                )

            model = ReturnPeriodModel(
                series=discharge,
                return_periods=return_periods_arr,
                min_exceed=min_exceed,
                nboot=nboot,
                fixed_shape=fixed_shape,
                fixed_scale=fixed_scale,
                p_value_threshold=p_value_threshold,
                selection_strategy=selection_strategy,
            )
            print(f"GPD-POT analysis completed for river {idx}.")

            if write_figures and output_directory:
                fig = model.plot_diagnostics()
                fig.savefig(output_directory / f"return_period_river_{idx}.svg")
                plt.close(fig)

            for return_period, return_water_level in model.rl_table.set_index(
                "T_years"
            )["GPD_POT_RL"].items():
                rivers.loc[idx, f"{prefix}_{return_period}"] = return_water_level

        return rivers

    def estimate_discharge_for_return_periods(
        self,
        discharge: xr.DataArray,
        rising_limb_hours: int = 72,
        return_periods: list[int | float] = [2, 5, 10, 20, 50, 100, 250, 500, 1000],
        p_value_threshold: float = 0.05,
        selection_strategy: str = "first_significant",
        fixed_shape: float | None = 0.0,
        write_figures: bool = False,
    ) -> None:
        """Estimate discharge for specified return periods and create hydrographs.

        Args:
            discharge: xr.DataArray containing the discharge data
            rising_limb_hours: number of hours for the rising limb of the hydrograph.
            return_periods: list of return periods for which to estimate discharge.
            p_value_threshold: Anderson-Darling p-value threshold for threshold selection. Defaults to 0.05.
            selection_strategy: Strategy for selecting the best threshold.
                'first_significant': Selects the first threshold (ordered high-to-low) with p_ad > p_value_threshold.
                'best_fit': Evaluates all thresholds and selects the one with the highest p-value.
            fixed_shape: Value to fix the shape parameter (xi) of the GPD. Set to 0.0 to force an Exponential (Gumbel) tail, or null to allow it to be fitted. Defaults to 0.0.
            write_figures: Whether to generate and save diagnostic figures.
        """
        recession_limb_hours: int = rising_limb_hours

        # here we only select the rivers that have an upstream forcing point
        rivers_with_return_period: gpd.GeoDataFrame = self.active_rivers[
            ~self.active_rivers["is_downstream_outflow"]
        ]

        river_representative_points: list[list[tuple[int, int]]] = []
        for ID in rivers_with_return_period.index:
            river_representative_points.append(
                get_representative_river_points(ID, self.rivers)
            )

        discharge_by_river, _ = get_discharge_and_river_parameters_by_river(
            rivers_with_return_period.index,
            river_representative_points,
            discharge=discharge,
        )
        rivers_with_return_period: gpd.GeoDataFrame = self.assign_return_periods(
            rivers_with_return_period,
            discharge_by_river,
            return_periods=return_periods,
            p_value_threshold=p_value_threshold,
            fixed_shape=fixed_shape,
            selection_strategy=selection_strategy,
            write_figures=write_figures,
            output_directory=self.figures_path / "return_periods",
        )

        for return_period in return_periods:
            self.rivers[f"hydrograph_{return_period}"] = None

        for river_idx in rivers_with_return_period.index:
            for return_period in return_periods:
                discharge_for_return_period = rivers_with_return_period.at[
                    river_idx, f"Q_{return_period}"
                ]
                hydrograph: pd.DataFrame = create_hourly_hydrograph(
                    discharge_for_return_period,
                    rising_limb_hours,
                    recession_limb_hours,
                )
                hydrograph: dict[str, Any] = {
                    time.isoformat(): Q.item()  # ty: ignore[unresolved-attribute]
                    for time, Q in hydrograph.iterrows()
                }
                self.rivers.at[river_idx, f"hydrograph_{return_period}"] = hydrograph

        export_rivers(self.path, rivers_with_return_period, postfix="_return_periods")

    @property
    def root(self) -> Path:
        """Gets the root directory for all SFINCS models.

        Returns:
            The path to the SFINCS model root directory.
        """
        return self._root

    def create_simulation(
        self,
        *args: Any,
        **kwargs: Any,
    ) -> SFINCSSimulation:
        """Sets forcing for a SFINCS model based on the provided parameters.

        Creates a new simulation directory and creteas a new sfincs model
        in that folder. Variables that do not change between simulations
        (e.g., grid, mask, etc.) are retained in the original model folder
        and inherited by the new simulation using relative paths.

        Args:
            *args: Positional arguments to pass to the SFINCSSimulation constructor.
            **kwargs: Keyword arguments to pass to the SFINCSSimulation constructor.

        Returns:
            An instance of SFINCSSimulation with the configured forcing.
        """
        return SFINCSSimulation(
            sfincs_root_model_path=self.root,
            sfincs_root_model_name=self.name,
            **kwargs,
        )

    def create_coastal_return_period_simulation(
        self,
        return_period: int,
        locations: gpd.GeoDataFrame,
        offset: xr.DataArray,
        sea_level_rise: pd.DataFrame | None = None,
        year: int | None = None,
    ) -> SFINCSSimulation:
        """
        Creates a SFINCS simulation with coastal water level forcing for a specified return period.

        It reads the coastal hydrograph timeseries and locations from pre-defined files, and sets the coastal water level forcing for the simulation.

        Args:
            return_period: The return period for which to create the coastal simulation.
            locations: A GeoDataFrame containing the locations of GTSM forcing stations.
            offset: The offset to apply to the coastal water level forcing based on mean sea level topography.
            sea_level_rise: A DataFrame containing sea level rise data.
            year: The year for which to apply sea level rise adjustments.
        Returns:
            An instance of SFINCSSimulation configured with coastal water level forcing.
        Raises:
            ValueError: If sea level rise data is provided but the station is not found in the data for the specified year.
        """
        # prepare coastal timeseries and locations
        timeseries = pd.read_csv(
            Path(
                f"output/hydrographs/gtsm_spring_tide_hydrograph_rp{return_period:04d}.csv"
            ),
            index_col=0,
        )

        # apply sea level rise adjustment if provided
        if sea_level_rise is not None and year is not None:
            # get sea level rise adjustment for the specified year without mutating
            # the original index (which may be reused elsewhere as a DatetimeIndex)
            year_mask = sea_level_rise.index.year == year  # ty:ignore[unresolved-attribute]
            if not year_mask.any():
                raise ValueError(f"No sea level rise data found for year {year}.")
            # select the first matching row; this results in a Series indexed by station ID
            sea_level_adjustment = sea_level_rise.loc[year_mask].iloc[0]
            # add the adjustment to the timeseries
            for station_id in timeseries.columns:
                if int(station_id) in sea_level_adjustment.index:
                    adjustment_value = sea_level_adjustment.loc[int(station_id)]
                    timeseries[station_id] += adjustment_value
                else:
                    raise ValueError(
                        f"Station ID {station_id} not found in sea level rise data for year {year}."
                    )
        # convert index to int
        # make a copy to avoid overwriting the original locations
        locations_copy = locations.copy()
        locations_copy.index = locations_copy.index.astype(int)

        timeseries.index = pd.to_datetime(timeseries.index, format="%Y-%m-%d %H:%M:%S")
        # convert columns to int
        timeseries.columns = timeseries.columns.astype(int)

        # Align timeseries columns with locations index
        timeseries = timeseries.loc[:, locations_copy.index]

        # now convert to incrementing integers starting from 0
        timeseries.columns = range(len(timeseries.columns))
        locations_copy.index = range(len(locations_copy.index))

        simulation: SFINCSSimulation = self.create_simulation(
            simulation_name=f"rp_{return_period}_coastal",
            start_time=timeseries.index[0],
            end_time=timeseries.index[-1],
        )

        # set coastal forcing model
        simulation.set_coastal_waterlevel_forcing(
            timeseries=timeseries, locations=locations_copy, offset=offset
        )
        return simulation

    @property
    def active_rivers(self) -> gpd.GeoDataFrame:
        """Returns a GeoDataFrame of rivers that are associated with the included subbasins.

        Returns:
            A GeoDataFrame of active rivers.
        """
        active_subbasin_ids = self.subbasins.index.tolist()
        return self.rivers[self.rivers.index.isin(active_subbasin_ids)]

    @property
    def active_cells(self) -> xr.DataArray:
        """Returns a boolean mask of the active cells in the SFINCS model.

        Returns:
            A boolean mask of the active cells in the SFINCS model.
        """
        return self.sfincs_model.grid["msk"] == 1

    @property
    def inflow_rivers(self) -> gpd.GeoDataFrame:
        """Returns a GeoDataFrame of rivers that are inflow points to the model.

        Returns:
            A GeoDataFrame of rivers that are inflow rivers.
        """
        non_headwater_rivers: gpd.GeoDataFrame = self.rivers[self.rivers["maxup"] > 0]
        non_outflow_basins: gpd.GeoDataFrame = non_headwater_rivers[
            (~non_headwater_rivers["is_downstream_outflow"])
            & (~non_headwater_rivers["is_further_downstream_outflow"])
        ]
        upstream_branches_in_domain = np.unique(
            self.rivers["downstream_ID"], return_counts=True
        )

        rivers_with_inflow = []
        for idx, row in non_outflow_basins.iterrows():
            if idx in upstream_branches_in_domain[0] and (
                upstream_branches_in_domain[1][
                    np.where(upstream_branches_in_domain[0] == idx)
                ]
                < row["maxup"]
            ):
                rivers_with_inflow.append(idx)

        return self.rivers[self.rivers.index.isin(rivers_with_inflow)]

    @property
    def non_inflow_rivers(self) -> gpd.GeoDataFrame:
        """Returns a GeoDataFrame of rivers that are not inflow points to the model.

        Returns:
            A GeoDataFrame of rivers that are not inflow rivers.
        """
        inflow_rivers: gpd.GeoDataFrame = self.inflow_rivers
        return self.rivers[~self.rivers.index.isin(inflow_rivers.index)]

    @property
    def headwater_rivers(self) -> gpd.GeoDataFrame:
        """Returns a GeoDataFrame of headwater rivers in the model.

        Returns:
            A GeoDataFrame of headwater rivers.
        """
        return self.rivers[self.rivers["maxup"] == 0]

    @property
    def has_inflow(self) -> bool:
        """Checks if the model has any inflow rivers.

        Returns:
            True if the model has inflow rivers, False otherwise.
        """
        return not self.inflow_rivers.empty

    @property
    def utm_zone(self) -> CRS:
        """Returns the UTM zone of the SFINCS model.

        Returns:
            The UTM zone as an integer.
        """
        centroid = self.subbasins.to_crs("EPSG:4326").unary_union.centroid
        return get_utm_zone(centroid.y, centroid.x)

    @property
    def is_build_from_scratch(self) -> bool:
        """Checks if the SFINCS model was built or read from disk.

        Returns:
            True if the SFINCS model was built, False otherwise.
        """
        return self._is_build_from_scratch

    @is_build_from_scratch.setter
    def is_build_from_scratch(self, value: bool) -> None:
        """Sets the is_build_from_scratch property of the SFINCS model.

        Args:
            value: A boolean indicating if the SFINCS model was built.
        """
        self._is_build_from_scratch = value

    def cleanup(self) -> None:
        """Cleans up the SFINCS model directory."""
        shutil.rmtree(self.path, ignore_errors=True)


class MultipleSFINCSSimulations:
    """Manages multiple SFINCS simulations as a single entity."""

    def __init__(self, simulations: list[SFINCSSimulation]) -> None:
        """Simulates running multiple SFINCS simulations as one.

        Args:
            simulations: A list of SFINCSSimulation instances to manage together.
        """
        self.simulations = simulations

    def run(
        self,
        ncpus: int | str = "auto",
        gpu: bool | str = "auto",
    ) -> None:
        """Runs all contained SFINCS simulations.

        Args:
            ncpus: Number of CPU cores to use for the simulation. Can be
                an integer or 'auto' to automatically detect available cores.
            gpu: Whether to use GPU acceleration for the simulations. Can be
                True, False, or 'auto' to automatically detect GPU availability.
        """
        for simulation in self.simulations:
            simulation.run(ncpus=ncpus, gpu=gpu)

    def read_max_flood_depth(self, minimum_flood_depth: float | int) -> xr.DataArray:
        """Reads the maximum flood depth map from the simulation output.

        Args:
            minimum_flood_depth: Minimum flood depth to consider in the output.

        Returns:
            An xarray DataArray containing the maximum flood depth.
        """
        flood_depths: list[xr.DataArray] = []
        for simulation in self.simulations:
            simulation_max_flood_depth: xr.DataArray = simulation.read_max_flood_depth(
                minimum_flood_depth
            )
            # merge_arrays expects the origin to be bottom-left, so we flip the y-axis
            simulation_max_flood_depth = simulation_max_flood_depth.isel(
                y=slice(None, None, -1)
            )
            flood_depths.append(simulation_max_flood_depth)

        rp_map: xr.DataArray = merge_arrays(flood_depths, method="max")
        rp_map = rp_map.isel(
            y=slice(None, None, -1)
        )  # flip back to original orientation

        assert rp_map.rio.crs is not None

        return rp_map

    def cleanup(self) -> None:
        """Cleans up all simulation directories."""
        for simulation in self.simulations:
            simulation.cleanup()


class SFINCSSimulation:
    """A SFINCS simulation with specific forcing and configuration.

    Created fro m a SFINCSRootModel instance which already contains the constant parts of the model.
    """

    def __init__(
        self,
        sfincs_root_model_path: Path,
        sfincs_root_model_name: str,
        simulation_name: str,
        start_time: datetime,
        end_time: datetime,
        spinup_seconds: int = 86400,
        write_figures: bool = True,
        flood_map_output_interval_seconds: int | None = None,
    ) -> None:
        """Initializes a SFINCSSimulation with specific forcing and configuration.

        Args:
            sfincs_root_model_path: Path to the root SFINCS model directory.
            sfincs_root_model_name: Name of the root SFINCS model.
            simulation_name: A string representing the name of the simulation.
                Also used to create the path to write the file to disk.
            start_time: The start time of the simulation as a datetime object.
            end_time: The end time of the simulation as a datetime object.
            spinup_seconds: The number of seconds to use for model spin-up. Defaults to 86400 (1 day).
            write_figures: Whether to generate and save figures for the model. Defaults to False.
            flood_map_output_interval_seconds: The interval in seconds at which to output flood maps. Defaults to None (15 minutes).
                None means no output.
        """
        self._name = simulation_name
        self.write_figures = write_figures
        self.start_time = start_time
        self.end_time = end_time

        # Read a new independent instance of the SFINCS model to avoid shared state
        self.root_model: SFINCSRootModel = SFINCSRootModel(
            root=sfincs_root_model_path, name=sfincs_root_model_name
        ).read()

        self.cleanup()

        sfincs_model = self.root_model.sfincs_model
        sfincs_model.set_root(str(self.path), mode="w+")

        # Configure the model's logger and all handlers to DEBUG level
        # This must be done AFTER reading SfincsModel when handlers are initialized
        sfincs_model.logger.setLevel(logging.DEBUG)
        for handler in sfincs_model.logger.handlers:
            handler.setLevel(logging.DEBUG)

        sfincs_model.setup_config(
            alpha=0.5,  # alpha is the parameter for the CFL-condition reduction. Decrease for additional numerical stability, minimum value is 0.1 and maximum is 0.75 (0.5 default value)
            # h73table=1,  # use h^(7/3) table for friction calculation. This is slightly less accurate but up to 30% faster, disabled by default. Enabling this may cause instability issues in the GPU runs, so use with caution.
            nc_deflate_level=9,  # compression level for netcdf output files (0-9)
            tspinup=spinup_seconds,  # spinup time in seconds
            dtout=flood_map_output_interval_seconds
            or 999999999,  # output flood maps every n seconds, 0 means no output
            dtmaxout=999999999,  # only report max flood depth at the end of the simulation
            tref=to_sfincs_datetime(start_time),  # reference time for the simulation
            tstart=to_sfincs_datetime(start_time),  # simulation start time
            tstop=to_sfincs_datetime(end_time),  # simulation end time
            **make_relative_paths(sfincs_model.config, self.root_path, self.path),
        )
        sfincs_model.write_config()

        self.sfincs_model: SfincsModel = sfincs_model

        # Track total volumes added via forcings (for water balance debugging)
        self.total_runoff_volume_m3: float = 0.0
        self.total_discharge_volume_m3: float = 0.0

        self.set_river_outflow_boundary_condition()

    def print_forcing_volume(self) -> None:
        """Print all forcing volumes for debugging the water balance."""
        msg: str = (
            f"SFINCS Forcing volumes: runoff={int(self.total_runoff_volume_m3)} m3, "
            f"discharge={int(self.total_discharge_volume_m3)} m3"
        )
        print(msg)

    def set_coastal_waterlevel_forcing(
        self,
        locations: gpd.GeoDataFrame,
        timeseries: pd.DataFrame,
        buffer: int = 100_000,
        offset: xr.DataArray | None = None,
    ) -> None:
        """Sets up coastal water level forcing for the SFINCS model from a timeseries.

        Args:
            locations: A GeoDataFrame containing the locations of the water level forcing points.
            timeseries: A DataFrame containing the water level timeseries for each node.
                The columns should match the index of the locations GeoDataFrame.
            buffer: Buffer distance in meters to extend the model domain for coastal forcing points.
            offset: The offset of water levels based on the m
        """
        # select only locations that are in the model
        self.sfincs_model.read_forcing()
        self.sfincs_model.setup_waterlevel_forcing(
            locations=locations,
            timeseries=timeseries,
            buffer=buffer,
            offset=offset,  # ty: ignore[invalid-argument-type]
        )
        self.sfincs_model.write_forcing()
        self.sfincs_model.write_config()

        if self.write_figures:
            self.sfincs_model.plot_forcing(
                fn_out=str(self.figures_path / "forcing.png")
            )
            self.sfincs_model.plot_basemap(
                fn_out=str(self.figures_path / "basemap.png"),
                vmin=math.floor(self.root_model.elevation.min()),
                vmax=max(
                    math.ceil(self.root_model.elevation.max()), 1
                ),  # vmax is required until bug in hydromt-sfincs fixed, see: https://github.com/Deltares/hydromt_sfincs/issues/324)
            )

    def set_forcing_from_grid(
        self, nodes: gpd.GeoDataFrame, discharge_grid: xr.DataArray
    ) -> None:
        """Sets up discharge forcing for the SFINCS model from a gridded dataset.

        Args:
            nodes: A GeoDataFrame containing the locations of the discharge forcing points.
            discharge_grid: Path to a raster file or an xarray DataArray containing discharge values in m^3/s.
                Usually this is from a hydrological model.
        """
        nodes: gpd.GeoDataFrame = nodes.copy()
        nodes["geometry"] = nodes["geometry"].apply(get_start_point)

        river_representative_points = []
        for ID in nodes.index:
            river_representative_points.append(
                get_representative_river_points(
                    ID,
                    nodes,
                )
            )

        discharge_by_river, _ = get_discharge_and_river_parameters_by_river(
            nodes.index,
            river_representative_points,
            discharge=discharge_grid,
        )

        locations = nodes.to_crs(self.sfincs_model.crs)

        self.set_discharge_forcing_from_nodes(
            nodes=locations,
            timeseries=discharge_by_river,
        )

    def set_headwater_forcing_from_grid(
        self,
        discharge_grid: xr.DataArray,
    ) -> None:
        """Sets up discharge forcing for the SFINCS model from a gridded dataset.

        Args:
            discharge_grid: Path to a raster file or an xarray DataArray containing discharge values in m^3/s.
                Usually this is from a hydrological model.
        """
        # Load rivers from file and filter for headwater rivers
        headwater_rivers: gpd.GeoDataFrame = self.root_model.rivers[
            self.root_model.rivers["maxup"] == 0
        ]
        self.set_forcing_from_grid(
            nodes=headwater_rivers,
            discharge_grid=discharge_grid,
        )

    def set_river_outflow_boundary_condition(self) -> None:
        """Sets up river outflow boundary condition for the SFINCS model.

        We use the pre-calculated outflow elevations from the root model
        to set constant water level boundary conditions at the river outflow points.
        """
        outflow_points: gpd.GeoDataFrame = self.root_model.rivers[
            ~self.root_model.rivers["outflow_point_xy"].isna()
        ].copy()
        if not outflow_points.empty:
            outflow_points["geometry"] = outflow_points["outflow_point_xy"].apply(
                lambda xy: Point(xy[0], xy[1])
            )
            outflow_points.index = range(1, len(outflow_points) + 1)

            # Create DataFrame with constant water level for each outflow point
            elevation_time_series_constant: pd.DataFrame = pd.DataFrame(
                data=outflow_points["outflow_elevation"].to_dict(),
                index=pd.date_range(start=self.start_time, end=self.end_time, freq="h"),
            )

            self.sfincs_model.set_forcing_1d(
                gdf_locs=outflow_points,
                df_ts=elevation_time_series_constant,
            )

    def set_inflow_forcing_from_grid(
        self,
        discharge_grid: xr.DataArray,
    ) -> None:
        """Sets up discharge forcing for the SFINCS model from a gridded dataset.

        Args:
            discharge_grid: Path to a raster file or an xarray DataArray containing discharge values in m^3/s.
                Usually this is from a hydrological model.
        """
        # Replicate the inflow_rivers property logic
        non_headwater_rivers: gpd.GeoDataFrame = self.root_model.rivers[
            self.root_model.rivers["maxup"] > 0
        ]
        non_outflow_basins: gpd.GeoDataFrame = non_headwater_rivers[
            ~non_headwater_rivers["is_downstream_outflow"]
        ]
        upstream_branches_in_domain = np.unique(
            self.root_model.rivers["downstream_ID"], return_counts=True
        )

        rivers_with_inflow = []
        for idx, row in non_outflow_basins.iterrows():
            downstream_ID = row["downstream_ID"]
            if downstream_ID not in self.root_model.rivers.index:
                continue
            upstream_branches_count = upstream_branches_in_domain[1][
                upstream_branches_in_domain[0] == downstream_ID
            ][0]
            if upstream_branches_count > 1:
                rivers_with_inflow.append(idx)

        inflow_rivers: gpd.GeoDataFrame = self.root_model.rivers[
            self.root_model.rivers.index.isin(rivers_with_inflow)
        ]
        self.set_forcing_from_grid(
            nodes=inflow_rivers,
            discharge_grid=discharge_grid,
        )

    def set_river_inflow(
        self, nodes: gpd.GeoDataFrame, timeseries: pd.DataFrame
    ) -> None:
        """
        Sets up river inflow boundary conditions for the SFINCS model.

        Here, we use negative indices to indicate inflow boundaries, so that they are
        separate from the runoff generated in the model domain.

        Args:
            nodes: GeoDataFrame containing the locations of river inflow points.
                The index values are converted to negative values to indicate inflow boundaries.
            timeseries: DataFrame containing discharge time series (in m^3/s) for each node.
                The columns should correspond to the node indices; these are also converted to negative values.
        """
        nodes = nodes.copy()
        timeseries = timeseries.copy()
        nodes.index = [-idx for idx in nodes.index]  # SFINCS negative index for inflow
        timeseries.columns = [-col for col in timeseries.columns]
        self.set_discharge_forcing_from_nodes(
            nodes=nodes,
            timeseries=timeseries,
        )

    def set_discharge_forcing_from_nodes(
        self, nodes: gpd.GeoDataFrame, timeseries: pd.DataFrame
    ) -> None:
        """Sets up discharge forcing for the SFINCS model from specified nodes (locations) and timeseries.

        Args:
            nodes: A GeoDataFrame containing the locations of the discharge forcing points.
                The index should start at 1 and be consecutive (1, 2, 3, ...).
            timeseries: A DataFrame containing the discharge timeseries for each node.
                The columns should match the index of the nodes GeoDataFrame.

        Raises:
            ValueError: If forcing locations are outside the model subbasins.
        """
        assert set(timeseries.columns) == set(nodes.index)

        if "dis" in self.sfincs_model.forcing:
            assert not np.isin(
                nodes.index, self.sfincs_model.forcing["dis"].index
            ).any(), "This forcing would overwrite existing discharge forcing points"

        assert (
            self.sfincs_model.region.union_all()
            .contains(nodes.geometry.to_crs(self.sfincs_model.crs))
            .all()
        ), "All forcing locations must be within the model subbasins"

        reprojected_nodes = nodes.to_crs(self.sfincs_model.grid["msk"].rio.crs)
        x_points: xr.DataArray = xr.DataArray(
            reprojected_nodes.geometry.x,
            dims="points",
        )
        y_points: xr.DataArray = xr.DataArray(
            reprojected_nodes.geometry.y,
            dims="points",
        )

        in_masked_area: xr.DataArray = self.sfincs_model.grid["msk"].sel(
            x=x_points,
            y=y_points,
            method="nearest",
        )
        if not (in_masked_area > 0).all():
            raise ValueError("Some forcing locations are outside the active model area")

        self.sfincs_model.setup_discharge_forcing(
            locations=nodes, timeseries=timeseries, merge=True
        )

        self.sfincs_model.write_forcing()
        self.sfincs_model.write_config()

        assert self.end_time == timeseries.index[-1], (
            "End time of timeseries does not match simulation end time, this will lead to accounting errors"
        )

        # the last timestep will not be used in SFINCS because it is the end time, which is why we
        # discard it
        self.total_discharge_volume_m3 += (
            timeseries[:-1].sum(axis=1)
            * (timeseries.index[1:] - timeseries.index[:-1]).total_seconds()
        ).sum()

        self.print_forcing_volume()

        if self.write_figures:
            self.sfincs_model.plot_basemap(
                fn_out=str(self.figures_path / "src_points_check.png"),
                vmin=math.floor(self.root_model.elevation.min()),
                vmax=max(
                    math.ceil(self.root_model.elevation.max()), 1
                ),  # vmax is required until bug in hydromt-sfincs fixed, see: https://github.com/Deltares/hydromt_sfincs/issues/324
            )
            self.sfincs_model.plot_forcing(
                fn_out=str(self.figures_path / "forcing.png")
            )

    def set_accumulated_runoff_forcing(
        self,
        runoff_m: xr.DataArray,
        river_network: FlwdirRaster,
        river_ids: TwoDArrayInt32,
        basin_ids: TwoDArrayInt32,
        upstream_area: TwoDArrayFloat32,
        cell_area: TwoDArrayFloat32,
    ) -> None:
        """Sets up accumulated runoff forcing for the SFINCS model.

        This function accumulates the runoff from the provided runoff grid to the river network
        and sets it as discharge forcing for the SFINCS model. For each river cell in the low-resolution
        runoff grid, the upstream area is determined and the runoff from all cells in that upstream
        area is accumulated to the river cell. The accumulated discharge is then set as forcing
        for the SFINCS model.

        In some cases, the upstream area of the most upstream low-res river point is larger than the
        upstream area of the most upstream high-res river point. This means that the runoff
        would be added further downstream. To avoid this, we check if this is the case, and if so,
        we scale the upstream area of the low-res river points to match the upstream area of the
        most upstream high-res river point.

        Args:
            runoff_m: xarray DataArray containing runoff values in m per time step.
            river_network: FlwdirRaster representing the river network flow directions.
            river_ids: 2D numpy array of river segment IDs for each cell in the grid.
            basin_ids: 2D numpy array of basin IDs for each cell in the grid.
            upstream_area: 2D numpy array of upstream area values for each cell in the grid.
            cell_area: 2D numpy array of cell area values for each cell in the grid.
        """
        INFLOW_MULTIPLICATION_FACTOR = 100000

        # select only the time range needed
        runoff_m: xr.DataArray = runoff_m.sel(
            time=slice(self.start_time, self.end_time)
        )

        # mask out all basins that are not in the model
        mask: TwoDArrayBool = np.isin(basin_ids, self.root_model.rivers.index)
        runoff_m = xr.where(mask, runoff_m, 0.0)  # set runoff to 0 outside model basins
        river_ids = np.where(
            mask, river_ids, -1
        )  # set river IDs to -1 outside model basins

        # first, we want to create unique IDs for each river cell. To allow us to map
        # it back to the original river segment, we create IDs as follows:
        # inflow_ID = river_segment_ID * INFLOW_MULTIPLICATION_FACTOR + offset
        # where offset is a number from 0 to N-1, with N being the number of cells in the river segment
        river_inflow_IDs = np.full_like(river_ids, -1, dtype=np.int64)
        xy_per_river_segment = value_indices(river_ids, ignore_value=-1)
        for ID, (ys, xs) in xy_per_river_segment.items():
            assert len(ys) < INFLOW_MULTIPLICATION_FACTOR - 2, (
                f"River segment has more than {INFLOW_MULTIPLICATION_FACTOR - 2} cells, which is not supported. "
                "Increase the multiplication factor in the inflow ID calculation."
            )

            river_upstream_area = upstream_area[ys, xs]
            up_to_downstream_ids = np.argsort(river_upstream_area)

            ys_up_to_down: npt.NDArray[np.int64] = ys[up_to_downstream_ids]
            xs_up_to_down: npt.NDArray[np.int64] = xs[up_to_downstream_ids]

            for i in range(len(ys_up_to_down)):
                inflow_ID: np.int64 = np.int64(ID) * np.int64(
                    INFLOW_MULTIPLICATION_FACTOR
                ) + np.int64(i)
                assert inflow_ID < 9_223_372_036_854_775_807, (
                    "Inflow ID exceeds maximum int64 value."
                )
                river_inflow_IDs[ys_up_to_down[i], xs_up_to_down[i]] = inflow_ID

        river_cells: TwoDArrayBool = river_inflow_IDs != -1
        river_ids_mapping: ArrayInt64 = river_inflow_IDs[river_cells]

        # starting from each river cell, create an upstream basin map for which
        # the discharge will be accumulated
        # .basins() must have river IDs starting from 1, so we use IDs from 1 to N + 1
        # then, we subtract 1 to make it zero based again
        subbasins: ArrayInt64 = river_network.basins(
            np.where(river_cells.ravel())[0],
            ids=np.arange(1, river_ids_mapping.size + 1, step=1, dtype=np.int64),
        )[mask]
        assert not (subbasins == 0).any()
        subbasins -= 1  # make zero based

        timestep_size: xr.DataArray = runoff_m.time.diff(dim="time").astype(
            "timedelta64[s]"
        )
        # confirm that timestep size is constant
        assert (timestep_size == timestep_size[0]).all()
        timestep_size_seconds: int = timestep_size[0].item().seconds

        # get the generated discharge in m3/s for each cell
        generated_discharge_m3_per_s: xr.DataArray = (
            runoff_m * cell_area / timestep_size_seconds
        )

        # accumulate generated discharge for each river cell
        accumulated_generated_discharge_m3_per_s: TwoDArrayFloat64 = (
            np.apply_along_axis(
                func1d=lambda x: np.bincount(subbasins, weights=x),
                axis=1,
                arr=generated_discharge_m3_per_s.values[:, mask],
            )
        )
        assert (
            accumulated_generated_discharge_m3_per_s.shape[1] == river_ids_mapping.size
        )
        assert accumulated_generated_discharge_m3_per_s.shape[1] == river_cells.sum()

        # create empty timeseries and nodes
        timeseries: pd.DataFrame = pd.DataFrame(
            {
                "time": generated_discharge_m3_per_s.time,
            }
        ).set_index("time")
        inflow_IDs_gdf: list[int] = []
        nodes: list[Point] = []

        for mapped_idx in range(accumulated_generated_discharge_m3_per_s.shape[1]):
            inflow_idx: int = river_ids_mapping[mapped_idx]

            # find the original river segment and offset
            river_ID = inflow_idx // INFLOW_MULTIPLICATION_FACTOR
            inflow_offset = inflow_idx % INFLOW_MULTIPLICATION_FACTOR

            river: gpd.GeoSeries = self.root_model.rivers.loc[river_ID]

            # confirm we have the correct inflow point
            xy_low_res = river["hydrography_xy"][inflow_offset]
            assert river_inflow_IDs[xy_low_res[1], xy_low_res[0]] == inflow_idx

            # find the high-res location corresponding to the low-res inflow point
            hydrography_upstream_area_m2 = river["hydrography_upstream_area_m2"]
            upstream_area_low_res = hydrography_upstream_area_m2[inflow_offset]
            hydrography_high_res_lons_lats = river["hydrography_high_res_lons_lats"]
            hydrography_high_res_upstream_area_m2 = river[
                "hydrography_high_res_upstream_area_m2"
            ]
            closest_upstream_area_index = np.argmin(
                np.abs(hydrography_high_res_upstream_area_m2 - upstream_area_low_res)
            )

            discharge_m3_per_s = accumulated_generated_discharge_m3_per_s[:, mapped_idx]

            # we want the runoff to start at the headwater of the river. However, sometimes due to
            # the low resolution hydrology, it is inherent that sometimes the inflow point is not at the headwater
            # but somewhere downstream. In that case, we need to add an additional inflow point at the headwater
            # and scale the discharge accordingly, based on the upstream area at the headwater and the upstream area.
            # We then place the normal inflow point at the closest high-res point, but with the discharge reduced accordingly.
            if (
                inflow_offset == 0  # check if this is the most upstream low-res point
                and closest_upstream_area_index
                != 0  # check if the closest high-res point is not the high-res headwater point
                and river["maxup"] == 0  # check if this river is a headwater river
            ):
                # create an additional inflow point at the headwater
                lon_headwater, lat_headwater = hydrography_high_res_lons_lats[0]
                nodes.append(Point(lon_headwater, lat_headwater))

                # for the index, we create a unique index based on the river ID. The headwater point
                # will always have offset INFLOW_MULTIPLICATION_FACTOR - 1 which will not
                # collide with any other inflow point
                headwater_idx = (
                    river_ID * INFLOW_MULTIPLICATION_FACTOR
                    + INFLOW_MULTIPLICATION_FACTOR
                    - 1
                )
                inflow_IDs_gdf.append(headwater_idx)

                # find the upstream areas of both the low-res downstream point
                # and the high-res headwater point
                upstream_area_downstream_point_m2 = (
                    hydrography_high_res_upstream_area_m2[closest_upstream_area_index]
                )
                upstream_area_headwater_point_m2 = (
                    hydrography_high_res_upstream_area_m2[0]
                )

                # scale the discharge based on the upstream areas
                headwater_discharge_m3_per_s = (
                    discharge_m3_per_s
                    * upstream_area_headwater_point_m2
                    / upstream_area_downstream_point_m2
                )
                timeseries[headwater_idx] = headwater_discharge_m3_per_s

                # reduce the discharge at the downstream point accordingly
                discharge_m3_per_s = discharge_m3_per_s - headwater_discharge_m3_per_s

            # get the lon lat of the closest high-res point
            lon, lat = hydrography_high_res_lons_lats[closest_upstream_area_index]

            # add the normal inflow point
            nodes.append(Point(lon, lat))
            inflow_IDs_gdf.append(inflow_idx)

            # add the timeseries
            timeseries[inflow_idx] = discharge_m3_per_s

        # create and sort the final nodes and timeseries
        nodes: gpd.GeoDataFrame = gpd.GeoDataFrame(
            index=inflow_IDs_gdf, geometry=nodes, crs=4326
        )

        nodes: gpd.GeoDataFrame = nodes.sort_index()  # ty: ignore[invalid-assignment]
        timeseries: pd.DataFrame = timeseries.sort_index(axis=1)

        self.set_discharge_forcing_from_nodes(
            nodes=nodes,
            timeseries=timeseries,
        )

    def run(self, ncpus: int | str = "auto", gpu: bool | str = "auto") -> None:
        """Runs the SFINCS simulation.

        Args:
            ncpus: Number of CPU cores to use for the simulation. Can be
                an integer or 'auto' to automatically detect available cores.
            gpu: Whether to use GPU acceleration for the simulation. Can be
                True, False, or 'auto' to automatically detect GPU availability.
        """
        assert gpu in [True, False, "auto"], "gpu must be True, False, or 'auto'"
        assert ncpus == "auto" or (isinstance(ncpus, int) and ncpus > 0), (
            "ncpus must be 'auto' or a positive integer"
        )
        run_sfincs_simulation(
            simulation_root=self.path,
            model_root=self.root_path,
            ncpus=ncpus,
            gpu=gpu,
        )

    def read_max_flood_depth(self, minimum_flood_depth: int | float) -> xr.DataArray:
        """Reads the maximum flood depth map from the simulation output.

        Args:
            minimum_flood_depth: Minimum flood depth to consider [m]. Values below this threshold are set to zero.

        Returns:
            An xarray DataArray containing the maximum flood depth.
        """
        flood_map: xr.DataArray = read_flood_depth(
            model_root=self.root_path,
            simulation_root=self.path,
            method="max",
            minimum_flood_depth=minimum_flood_depth,
            end_time=self.end_time,
            mask=self.root_model.area_of_interest,
        )
        return flood_map

    def read_final_flood_depth(self, minimum_flood_depth: float | int) -> xr.DataArray:
        """Reads the final flood depth map from the simulation output.

        Args:
            minimum_flood_depth: Minimum flood depth to consider [m]. Values below this threshold are set to zero.

        Returns:
            An xarray DataArray containing the final flood depth.
        """
        flood_map: xr.DataArray = read_flood_depth(
            model_root=self.root_path,
            simulation_root=self.path,
            method="final",
            minimum_flood_depth=minimum_flood_depth,
            end_time=self.end_time,
        )
        return flood_map

    def get_flood_volume(self, flood_depth: xr.DataArray) -> float:
        """Compute the total flood volume from the flood depth map.

        Args:
            flood_depth: An xarray DataArray containing the flood depth.

        Returns:
            The total flood volume in cubic meters.

        Raises:
            ValueError: If the model CRS is not defined.
        """
        if hasattr(flood_depth, "compute"):
            flood_depth = flood_depth.compute()

        if self.sfincs_model.crs is None:
            raise ValueError("Model CRS is not defined.")

        # Calculate cell area based on the model's CRS
        if self.sfincs_model.crs.is_geographic:
            cell_area: xr.DataArray = xr.full_like(
                self.sfincs_model.grid["dep"], np.nan, dtype=np.float32
            )
            height, width = self.sfincs_model.grid["dep"].shape
            cell_area.values = calculate_cell_area_m2(
                self.sfincs_model.grid["dep"].rio.transform(), height, width
            )
        else:
            cell_area = xr.full_like(
                self.sfincs_model.grid["dep"],
                (
                    self.sfincs_model.grid["dep"].rio.resolution()[0]
                    * self.sfincs_model.grid["dep"].rio.resolution()[1]
                ),
                dtype=np.float32,
            )
        return (flood_depth * cell_area).sum().item()

    def cleanup(self) -> None:
        """Cleans up the simulation directory by removing temporary files."""
        shutil.rmtree(path=self.path, ignore_errors=True)

    @property
    def root_path(self) -> Path:
        """Returns the root directory for the SFINCS model files."""
        return self.root_model.path

    @property
    def path(self) -> Path:
        """Returns the root directory for the SFINCS simulation files."""
        folder: Path = self.root_path / "simulations" / self.name
        folder.mkdir(parents=True, exist_ok=True)
        return folder

    @property
    def figures_path(self) -> Path:
        """Returns the directory for the SFINCS simulation diagnostic figures."""
        folder: Path = self.path / "figures"
        folder.mkdir(parents=True, exist_ok=True)
        return folder

    @property
    def name(self) -> str:
        """Returns the name of the simulation."""
        return self._name

    def has_outflow_boundary(self) -> bool:
        """Checks if the SFINCS model has an outflow boundary condition.

        Returns:
            True if the SFINCS model has an outflow boundary condition, False otherwise.
        """
        return (self.sfincs_model.grid["msk"] == 2).any().item()

    def get_cumulative_precipitation(self) -> xr.DataArray:
        """Reads the cumulative precipitation from the SFINCS model results.

        Returns:
            An xarray DataArray containing the cumulative precipitation.
        """
        self.sfincs_model.read_results()
        cumulative_precipitation = self.sfincs_model.results["cumprcp"].isel(timemax=-1)
        assert isinstance(cumulative_precipitation, xr.DataArray)
        return cumulative_precipitation

    def setup_thin_dams(self, dam_locations: gpd.GeoDataFrame) -> None:
        """Sets up thin dams in the ModFlow groundwater model.

        Args:
            dam_locations: A GeoDataFrame containing the locations of the dams.
        """
        self.sfincs_model.setup_structures(dam_locations, stype="thd")
        self.sfincs_model.write_geoms(data_vars=["thd"])
