flow_accumulation_mfd.py

"""Flow accumulation for Multiple Flow Direction (MFD) grids.

Takes the (8, H, W) fractional flow direction output from
``flow_direction_mfd`` and accumulates upstream area through all
downslope paths simultaneously.

Algorithm
---------
CPU : Kahn's BFS topological sort -- O(N).
GPU : iterative frontier peeling with pull-based kernels.
Dask: iterative tile sweep with boundary propagation (one tile in
      RAM at a time), following the ``flow_accumulation.py`` pattern.
"""

from __future__ import annotations

import numpy as np
import xarray as xr
from numba import cuda

try:
    import cupy
except ImportError:
    class cupy:  # type: ignore[no-redef]
        ndarray = False

try:
    import dask.array as da
except ImportError:
    da = None

from xrspatial.utils import (
    _dask_task_name_kwargs,
    _validate_mfd_fractions,
    _validate_raster,
    cuda_args,
    has_cuda_and_cupy,
    is_cupy_array,
    is_dask_cupy,
    ngjit,
)
from xrspatial.hydro._boundary_store import BoundaryStore
from xrspatial.dataset_support import supports_dataset


# =====================================================================
# Memory guards
# =====================================================================
#
# CPU peak working set per pixel for ``_flow_accum_mfd_cpu``:
#   accum    : float64 -> 8
#   in_degree: int32   -> 4
#   valid    : int8    -> 1
#   queue_r  : int64   -> 8
#   queue_c  : int64   -> 8
# Total ~29 bytes/pixel.  The caller-provided ``flow_dir_mfd`` array
# already lives in RAM before the kernel runs and is not double-counted
# here -- this matches the convention in ``flow_accumulation_d8``.
_BYTES_PER_PIXEL = 29

# GPU peak working set per pixel for ``_flow_accum_mfd_cupy``:
#   accum     : float64 -> 8
#   in_degree : int32   -> 4
#   state     : int32   -> 4
# Total ~16 bytes/pixel.  ``fractions_data`` already lives on the device
# before the kernel runs and is not double-counted here.
_GPU_BYTES_PER_PIXEL = 16


def _available_memory_bytes():
    """Best-effort estimate of available host memory in bytes."""
    try:
        with open('/proc/meminfo', 'r') as f:
            for line in f:
                if line.startswith('MemAvailable:'):
                    return int(line.split()[1]) * 1024  # kB -> bytes
    except (OSError, ValueError, IndexError):
        pass
    try:
        import psutil
        return psutil.virtual_memory().available
    except (ImportError, AttributeError):
        pass
    return 2 * 1024 ** 3


def _available_gpu_memory_bytes():
    """Best-effort estimate of free GPU memory in bytes.

    Returns 0 if CuPy / CUDA is unavailable or the query fails -- callers
    use that as a sentinel meaning "no GPU info, skip the guard".
    """
    try:
        import cupy as _cp
        free, _total = _cp.cuda.runtime.memGetInfo()
        return int(free)
    except Exception:
        return 0


def _check_memory(height, width):
    """Raise MemoryError if the BFS kernel would exceed 50% of RAM."""
    required = int(height) * int(width) * _BYTES_PER_PIXEL
    available = _available_memory_bytes()
    if required > 0.5 * available:
        raise MemoryError(
            f"flow_accumulation_mfd on a {height}x{width} grid requires "
            f"~{required / 1e9:.1f} GB of working memory but only "
            f"~{available / 1e9:.1f} GB is available.  Use a "
            f"dask-backed DataArray for out-of-core processing."
        )


def _check_gpu_memory(height, width):
    """Raise MemoryError if the CuPy kernel would exceed 50% of free GPU RAM.

    Skips the check (returns silently) when ``_available_gpu_memory_bytes``
    cannot determine the free memory -- e.g. on hosts without CUDA, where
    the kernel will fail at the cupy.asarray boundary anyway.
    """
    available = _available_gpu_memory_bytes()
    if available <= 0:
        return
    required = int(height) * int(width) * _GPU_BYTES_PER_PIXEL
    if required > 0.5 * available:
        raise MemoryError(
            f"flow_accumulation_mfd on a {height}x{width} grid requires "
            f"~{required / 1e9:.1f} GB of GPU working memory but only "
            f"~{available / 1e9:.1f} GB is free on the active device.  "
            f"Use a dask+cupy DataArray for out-of-core processing."
        )


# Neighbor offsets: E, SE, S, SW, W, NW, N, NE
_DY = np.array([0, 1, 1, 1, 0, -1, -1, -1], dtype=np.int64)
_DX = np.array([1, 1, 0, -1, -1, -1, 0, 1], dtype=np.int64)

# Opposite neighbor index (who points back at me?)
# E(0)->W(4), SE(1)->NW(5), S(2)->N(6), SW(3)->NE(7), ...
_OPPOSITE = np.array([4, 5, 6, 7, 0, 1, 2, 3], dtype=np.int64)


# =====================================================================
# CPU kernel
# =====================================================================

@ngjit
def _flow_accum_mfd_cpu(fractions, height, width):
    """Kahn's BFS topological sort for MFD flow accumulation.

    Parameters
    ----------
    fractions : (8, H, W) float64 array of flow fractions
    height, width : int

    Returns
    -------
    accum : (H, W) float64 array
    """
    dy = np.array([0, 1, 1, 1, 0, -1, -1, -1], dtype=np.int64)
    dx = np.array([1, 1, 0, -1, -1, -1, 0, 1], dtype=np.int64)

    accum = np.empty((height, width), dtype=np.float64)
    in_degree = np.zeros((height, width), dtype=np.int32)
    valid = np.zeros((height, width), dtype=np.int8)

    # Pass 1: initialise
    n_valid = 0
    for r in range(height):
        for c in range(width):
            v = fractions[0, r, c]
            if v != v:  # NaN
                accum[r, c] = np.nan
            else:
                valid[r, c] = 1
                accum[r, c] = 1.0
                n_valid += 1

    # Pass 2: compute in-degrees
    for r in range(height):
        for c in range(width):
            if valid[r, c] == 0:
                continue
            for k in range(8):
                if fractions[k, r, c] > 0.0:
                    nr = r + dy[k]
                    nc = c + dx[k]
                    if 0 <= nr < height and 0 <= nc < width:
                        if valid[nr, nc] == 1:
                            in_degree[nr, nc] += 1

    # BFS queue
    queue_r = np.empty(height * width, dtype=np.int64)
    queue_c = np.empty(height * width, dtype=np.int64)
    head = np.int64(0)
    tail = np.int64(0)

    for r in range(height):
        for c in range(width):
            if valid[r, c] == 1 and in_degree[r, c] == 0:
                queue_r[tail] = r
                queue_c[tail] = c
                tail += 1

    while head < tail:
        r = queue_r[head]
        c = queue_c[head]
        head += 1

        for k in range(8):
            frac = fractions[k, r, c]
            if frac > 0.0:
                nr = r + dy[k]
                nc = c + dx[k]
                if 0 <= nr < height and 0 <= nc < width:
                    if valid[nr, nc] == 1:
                        accum[nr, nc] += accum[r, c] * frac
                        in_degree[nr, nc] -= 1
                        if in_degree[nr, nc] == 0:
                            queue_r[tail] = nr
                            queue_c[tail] = nc
                            tail += 1

    # If a cycle remains, some valid cells never reached in_degree 0 and
    # were never dequeued.  head counts the cells processed by the BFS.
    if head < n_valid:
        raise ValueError(
            "flow_accumulation_mfd: the MFD fraction grid contains a cycle "
            "(some cells never reach zero in-degree).  The input must be a "
            "directed acyclic graph, as produced by flow_direction_mfd."
        )

    return accum


# =====================================================================
# GPU kernels
# =====================================================================

@cuda.jit
def _init_accum_indegree_mfd(fractions, accum, in_degree, state, H, W):
    """Initialise accum, in_degree and state for MFD on GPU."""
    i, j = cuda.grid(2)
    if i >= H or j >= W:
        return

    v = fractions[0, i, j]
    if v != v:  # NaN
        state[i, j] = 0
        accum[i, j] = 0.0
        return

    state[i, j] = 1
    accum[i, j] = 1.0

    # Neighbor offsets: E, SE, S, SW, W, NW, N, NE
    for k in range(8):
        frac = fractions[k, i, j]
        if frac <= 0.0:
            continue

        if k == 0:
            dy, dx = 0, 1
        elif k == 1:
            dy, dx = 1, 1
        elif k == 2:
            dy, dx = 1, 0
        elif k == 3:
            dy, dx = 1, -1
        elif k == 4:
            dy, dx = 0, -1
        elif k == 5:
            dy, dx = -1, -1
        elif k == 6:
            dy, dx = -1, 0
        else:
            dy, dx = -1, 1

        ni = i + dy
        nj = j + dx
        if 0 <= ni < H and 0 <= nj < W:
            cuda.atomic.add(in_degree, (ni, nj), 1)


@cuda.jit
def _find_ready_and_finalize_mfd(in_degree, state, changed, H, W):
    """Finalize previous frontier (2->3), mark new frontier (1->2)."""
    i, j = cuda.grid(2)
    if i >= H or j >= W:
        return

    if state[i, j] == 2:
        state[i, j] = 3

    if state[i, j] == 1 and in_degree[i, j] == 0:
        state[i, j] = 2
        cuda.atomic.add(changed, 0, 1)


@cuda.jit
def _pull_from_frontier_mfd(fractions, accum, in_degree, state, H, W):
    """Active MFD cells pull accumulation from frontier neighbours."""
    i, j = cuda.grid(2)
    if i >= H or j >= W:
        return

    if state[i, j] != 1:
        return

    # Opposite direction index: if neighbor k sent flow in direction k,
    # I am the opposite direction from them.
    # E(0)->W(4), SE(1)->NW(5), S(2)->N(6), SW(3)->NE(7), etc.
    for nbr in range(8):
        if nbr == 0:
            dy, dx = 0, 1
        elif nbr == 1:
            dy, dx = 1, 1
        elif nbr == 2:
            dy, dx = 1, 0
        elif nbr == 3:
            dy, dx = 1, -1
        elif nbr == 4:
            dy, dx = 0, -1
        elif nbr == 5:
            dy, dx = -1, -1
        elif nbr == 6:
            dy, dx = -1, 0
        else:
            dy, dx = -1, 1

        ni = i + dy
        nj = j + dx
        if ni < 0 or ni >= H or nj < 0 or nj >= W:
            continue
        if state[ni, nj] != 2:
            continue

        # Opposite of nbr: the direction from neighbor back to me
        if nbr == 0:
            opp = 4
        elif nbr == 1:
            opp = 5
        elif nbr == 2:
            opp = 6
        elif nbr == 3:
            opp = 7
        elif nbr == 4:
            opp = 0
        elif nbr == 5:
            opp = 1
        elif nbr == 6:
            opp = 2
        else:
            opp = 3

        frac = fractions[opp, ni, nj]
        if frac > 0.0:
            accum[i, j] += accum[ni, nj] * frac
            in_degree[i, j] -= 1


def _flow_accum_mfd_cupy(fractions_data):
    """GPU driver: iterative frontier peeling for MFD."""
    import cupy as cp

    _, H, W = fractions_data.shape
    fractions_f64 = fractions_data.astype(cp.float64)

    accum = cp.zeros((H, W), dtype=cp.float64)
    in_degree = cp.zeros((H, W), dtype=cp.int32)
    state = cp.zeros((H, W), dtype=cp.int32)
    changed = cp.zeros(1, dtype=cp.int32)

    griddim, blockdim = cuda_args((H, W))

    _init_accum_indegree_mfd[griddim, blockdim](
        fractions_f64, accum, in_degree, state, H, W)

    max_iter = H * W
    for _ in range(max_iter):
        changed[0] = 0
        _find_ready_and_finalize_mfd[griddim, blockdim](
            in_degree, state, changed, H, W)

        if int(changed[0]) == 0:
            break

        _pull_from_frontier_mfd[griddim, blockdim](
            fractions_f64, accum, in_degree, state, H, W)

    accum = cp.where(state == 0, cp.nan, accum)
    return accum


# =====================================================================
# Dask tile kernel
# =====================================================================

@ngjit
def _flow_accum_mfd_tile_kernel(fractions, h, w,
                                 seed_top, seed_bottom,
                                 seed_left, seed_right,
                                 seed_tl, seed_tr, seed_bl, seed_br):
    """Seeded BFS MFD flow accumulation for a single tile.

    Parameters
    ----------
    fractions : (8, h, w) float64 -- MFD flow fractions for this tile
    """
    dy = np.array([0, 1, 1, 1, 0, -1, -1, -1], dtype=np.int64)
    dx = np.array([1, 1, 0, -1, -1, -1, 0, 1], dtype=np.int64)

    accum = np.empty((h, w), dtype=np.float64)
    in_degree = np.zeros((h, w), dtype=np.int32)
    valid = np.zeros((h, w), dtype=np.int8)

    # Initialise
    n_valid = 0
    for r in range(h):
        for c in range(w):
            v = fractions[0, r, c]
            if v == v:  # not NaN
                valid[r, c] = 1
                accum[r, c] = 1.0
                n_valid += 1
            else:
                accum[r, c] = np.nan

    # Add external seeds
    for c in range(w):
        if valid[0, c] == 1:
            accum[0, c] += seed_top[c]
        if valid[h - 1, c] == 1:
            accum[h - 1, c] += seed_bottom[c]
    for r in range(h):
        if valid[r, 0] == 1:
            accum[r, 0] += seed_left[r]
        if valid[r, w - 1] == 1:
            accum[r, w - 1] += seed_right[r]

    if valid[0, 0] == 1:
        accum[0, 0] += seed_tl
    if valid[0, w - 1] == 1:
        accum[0, w - 1] += seed_tr
    if valid[h - 1, 0] == 1:
        accum[h - 1, 0] += seed_bl
    if valid[h - 1, w - 1] == 1:
        accum[h - 1, w - 1] += seed_br

    # Compute in-degrees
    for r in range(h):
        for c in range(w):
            if valid[r, c] == 0:
                continue
            for k in range(8):
                if fractions[k, r, c] > 0.0:
                    nr = r + dy[k]
                    nc = c + dx[k]
                    if 0 <= nr < h and 0 <= nc < w and valid[nr, nc] == 1:
                        in_degree[nr, nc] += 1

    # BFS
    queue_r = np.empty(h * w, dtype=np.int64)
    queue_c = np.empty(h * w, dtype=np.int64)
    head = np.int64(0)
    tail = np.int64(0)

    for r in range(h):
        for c in range(w):
            if valid[r, c] == 1 and in_degree[r, c] == 0:
                queue_r[tail] = r
                queue_c[tail] = c
                tail += 1

    while head < tail:
        r = queue_r[head]
        c = queue_c[head]
        head += 1

        for k in range(8):
            frac = fractions[k, r, c]
            if frac > 0.0:
                nr = r + dy[k]
                nc = c + dx[k]
                if 0 <= nr < h and 0 <= nc < w and valid[nr, nc] == 1:
                    accum[nr, nc] += accum[r, c] * frac
                    in_degree[nr, nc] -= 1
                    if in_degree[nr, nc] == 0:
                        queue_r[tail] = nr
                        queue_c[tail] = nc
                        tail += 1

    # A cycle within this tile leaves some valid cells undequeued.
    if head < n_valid:
        raise ValueError(
            "flow_accumulation_mfd: the MFD fraction grid contains a cycle "
            "(some cells never reach zero in-degree).  The input must be a "
            "directed acyclic graph, as produced by flow_direction_mfd."
        )

    return accum


# =====================================================================
# Dask iterative tile sweep
# =====================================================================

def _preprocess_mfd_tiles(fractions_da, chunks_y, chunks_x):
    """Extract boundary fraction strips into a dict.

    For MFD we need the full 8-band fractions at each boundary cell,
    so we store them as (8, length) arrays.
    """
    n_tile_y = len(chunks_y)
    n_tile_x = len(chunks_x)

    # Store fraction strips keyed by (side, iy, ix)
    frac_bdry = {}

    for iy in range(n_tile_y):
        for ix in range(n_tile_x):
            # fractions_da is (8, H, W) dask array
            # Each tile's fractions: shape (8, tile_h, tile_w)
            chunk = fractions_da[:, sum(chunks_y[:iy]):sum(chunks_y[:iy+1]),
                                    sum(chunks_x[:ix]):sum(chunks_x[:ix+1])].compute()
            chunk = np.asarray(chunk, dtype=np.float64)

            # top row: (8, tile_w)
            frac_bdry[('top', iy, ix)] = chunk[:, 0, :].copy()
            # bottom row: (8, tile_w)
            frac_bdry[('bottom', iy, ix)] = chunk[:, -1, :].copy()
            # left col: (8, tile_h)
            frac_bdry[('left', iy, ix)] = chunk[:, :, 0].copy()
            # right col: (8, tile_h)
            frac_bdry[('right', iy, ix)] = chunk[:, :, -1].copy()

    return frac_bdry


def _compute_seeds_mfd(iy, ix, boundaries, frac_bdry,
                        chunks_y, chunks_x, n_tile_y, n_tile_x):
    """Compute seed arrays for tile (iy, ix) from neighbour boundaries.

    For MFD, a neighbor cell flows into the current tile if its fraction
    for the direction pointing into our tile is > 0.
    """
    # Neighbor offsets: E(0), SE(1), S(2), SW(3), W(4), NW(5), N(6), NE(7)
    # Opposite:         W(4), NW(5), N(6), NE(7), E(0), SE(1), S(2), SE(3)

    tile_h = chunks_y[iy]
    tile_w = chunks_x[ix]

    seed_top = np.zeros(tile_w, dtype=np.float64)
    seed_bottom = np.zeros(tile_w, dtype=np.float64)
    seed_left = np.zeros(tile_h, dtype=np.float64)
    seed_right = np.zeros(tile_h, dtype=np.float64)
    seed_tl = 0.0
    seed_tr = 0.0
    seed_bl = 0.0
    seed_br = 0.0

    dy_arr = np.array([0, 1, 1, 1, 0, -1, -1, -1], dtype=np.int64)
    dx_arr = np.array([1, 1, 0, -1, -1, -1, 0, 1], dtype=np.int64)

    # --- Top edge: bottom row of tile above ---
    if iy > 0:
        nb_frac = frac_bdry[('bottom', iy - 1, ix)]  # (8, tile_w)
        nb_accum = boundaries.get('bottom', iy - 1, ix)
        w = nb_frac.shape[1]
        for c in range(w):
            for k in range(8):
                if not (nb_frac[k, c] > 0.0):
                    continue
                # Direction k from neighbor: dy_arr[k], dx_arr[k]
                # Neighbor is in row above, so dy must be +1 to enter our tile
                ndy = dy_arr[k]
                ndx = dx_arr[k]
                if ndy == 1:  # flows south into our tile
                    tc = c + ndx
                    if 0 <= tc < tile_w:
                        seed_top[tc] += nb_accum[c] * nb_frac[k, c]

    # --- Bottom edge: top row of tile below ---
    if iy < n_tile_y - 1:
        nb_frac = frac_bdry[('top', iy + 1, ix)]  # (8, tile_w)
        nb_accum = boundaries.get('top', iy + 1, ix)
        w = nb_frac.shape[1]
        for c in range(w):
            for k in range(8):
                if not (nb_frac[k, c] > 0.0):
                    continue
                ndy = dy_arr[k]
                ndx = dx_arr[k]
                if ndy == -1:  # flows north into our tile
                    tc = c + ndx
                    if 0 <= tc < tile_w:
                        seed_bottom[tc] += nb_accum[c] * nb_frac[k, c]

    # --- Left edge: right column of tile to the left ---
    if ix > 0:
        nb_frac = frac_bdry[('right', iy, ix - 1)]  # (8, tile_h)
        nb_accum = boundaries.get('right', iy, ix - 1)
        h = nb_frac.shape[1]
        for r in range(h):
            for k in range(8):
                if not (nb_frac[k, r] > 0.0):
                    continue
                ndy = dy_arr[k]
                ndx = dx_arr[k]
                if ndx == 1:  # flows east into our tile
                    tr = r + ndy
                    if 0 <= tr < tile_h:
                        seed_left[tr] += nb_accum[r] * nb_frac[k, r]

    # --- Right edge: left column of tile to the right ---
    if ix < n_tile_x - 1:
        nb_frac = frac_bdry[('left', iy, ix + 1)]  # (8, tile_h)
        nb_accum = boundaries.get('left', iy, ix + 1)
        h = nb_frac.shape[1]
        for r in range(h):
            for k in range(8):
                if not (nb_frac[k, r] > 0.0):
                    continue
                ndy = dy_arr[k]
                ndx = dx_arr[k]
                if ndx == -1:  # flows west into our tile
                    tr = r + ndy
                    if 0 <= tr < tile_h:
                        seed_right[tr] += nb_accum[r] * nb_frac[k, r]

    # --- Diagonal corner seeds ---
    # TL: bottom-right cell of (iy-1, ix-1) flows SE (dy=1, dx=1 -> k=1)
    if iy > 0 and ix > 0:
        nb_frac = frac_bdry[('bottom', iy - 1, ix - 1)]  # (8, w)
        av = float(boundaries.get('bottom', iy - 1, ix - 1)[-1])
        frac_se = nb_frac[1, -1]  # SE direction
        if frac_se > 0.0:
            seed_tl += av * frac_se

    # TR: bottom-left cell of (iy-1, ix+1) flows SW (dy=1, dx=-1 -> k=3)
    if iy > 0 and ix < n_tile_x - 1:
        nb_frac = frac_bdry[('bottom', iy - 1, ix + 1)]  # (8, w)
        av = float(boundaries.get('bottom', iy - 1, ix + 1)[0])
        frac_sw = nb_frac[3, 0]  # SW direction
        if frac_sw > 0.0:
            seed_tr += av * frac_sw

    # BL: top-right cell of (iy+1, ix-1) flows NE (dy=-1, dx=1 -> k=7)
    if iy < n_tile_y - 1 and ix > 0:
        nb_frac = frac_bdry[('top', iy + 1, ix - 1)]  # (8, w)
        av = float(boundaries.get('top', iy + 1, ix - 1)[-1])
        frac_ne = nb_frac[7, -1]  # NE direction
        if frac_ne > 0.0:
            seed_bl += av * frac_ne

    # BR: top-left cell of (iy+1, ix+1) flows NW (dy=-1, dx=-1 -> k=5)
    if iy < n_tile_y - 1 and ix < n_tile_x - 1:
        nb_frac = frac_bdry[('top', iy + 1, ix + 1)]  # (8, w)
        av = float(boundaries.get('top', iy + 1, ix + 1)[0])
        frac_nw = nb_frac[5, 0]  # NW direction
        if frac_nw > 0.0:
            seed_br += av * frac_nw

    return (seed_top, seed_bottom, seed_left, seed_right,
            seed_tl, seed_tr, seed_bl, seed_br)


def _process_tile_mfd(iy, ix, fractions_da, boundaries, frac_bdry,
                       chunks_y, chunks_x, n_tile_y, n_tile_x):
    """Run seeded MFD BFS on one tile; update boundaries in-place."""
    # Extract this tile's fractions: (8, tile_h, tile_w)
    y_start = sum(chunks_y[:iy])
    y_end = y_start + chunks_y[iy]
    x_start = sum(chunks_x[:ix])
    x_end = x_start + chunks_x[ix]

    chunk = np.asarray(
        fractions_da[:, y_start:y_end, x_start:x_end].compute(),
        dtype=np.float64)
    _, h, w = chunk.shape

    seeds = _compute_seeds_mfd(
        iy, ix, boundaries, frac_bdry,
        chunks_y, chunks_x, n_tile_y, n_tile_x)

    accum = _flow_accum_mfd_tile_kernel(chunk, h, w, *seeds)

    # NaN cells don't contribute flow; replace with 0 for boundary storage
    new_top = np.where(np.isnan(accum[0, :]), 0.0, accum[0, :])
    new_bottom = np.where(np.isnan(accum[-1, :]), 0.0, accum[-1, :])
    new_left = np.where(np.isnan(accum[:, 0]), 0.0, accum[:, 0])
    new_right = np.where(np.isnan(accum[:, -1]), 0.0, accum[:, -1])

    change = 0.0
    for side, new in (('top', new_top), ('bottom', new_bottom),
                      ('left', new_left), ('right', new_right)):
        old = boundaries.get(side, iy, ix)
        with np.errstate(invalid='ignore'):
            diff = np.abs(new - old)
        diff = np.where(np.isnan(diff), 0.0, diff)
        m = float(np.max(diff))
        if m > change:
            change = m

    boundaries.set('top', iy, ix, new_top)
    boundaries.set('bottom', iy, ix, new_bottom)
    boundaries.set('left', iy, ix, new_left)
    boundaries.set('right', iy, ix, new_right)

    return change


def _flow_accum_mfd_dask_iterative(fractions_da, chunks_y, chunks_x):
    """Iterative boundary-propagation for MFD dask arrays.

    Parameters
    ----------
    fractions_da : dask array of shape (8, H, W)
    chunks_y, chunks_x : tuples of chunk sizes for the spatial dims
    """
    n_tile_y = len(chunks_y)
    n_tile_x = len(chunks_x)

    # The 8 direction bands must stay in a single chunk: every tile kernel
    # needs all 8 fractions, and the lazy assembly drops axis 0 per block.
    if fractions_da.chunks[0] != (fractions_da.shape[0],):
        fractions_da = fractions_da.rechunk({0: fractions_da.shape[0]})

    # Phase 0: extract boundary fraction strips
    frac_bdry = _preprocess_mfd_tiles(fractions_da, chunks_y, chunks_x)

    # Phase 1: initialise boundary accum to 0
    boundaries = BoundaryStore(chunks_y, chunks_x, fill_value=0.0)

    # Phase 2: iterative forward/backward sweeps
    max_iterations = max(n_tile_y, n_tile_x) + 10

    for _iteration in range(max_iterations):
        max_change = 0.0

        for iy in range(n_tile_y):
            for ix in range(n_tile_x):
                c = _process_tile_mfd(iy, ix, fractions_da, boundaries,
                                       frac_bdry, chunks_y, chunks_x,
                                       n_tile_y, n_tile_x)
                if c > max_change:
                    max_change = c

        for iy in reversed(range(n_tile_y)):
            for ix in reversed(range(n_tile_x)):
                c = _process_tile_mfd(iy, ix, fractions_da, boundaries,
                                       frac_bdry, chunks_y, chunks_x,
                                       n_tile_y, n_tile_x)
                if c > max_change:
                    max_change = c

        if max_change == 0.0:
            break

    boundaries = boundaries.snapshot()

    return _assemble_result_mfd(fractions_da, boundaries, frac_bdry,
                                 chunks_y, chunks_x, n_tile_y, n_tile_x)


def _assemble_result_mfd(fractions_da, boundaries, frac_bdry,
                          chunks_y, chunks_x, n_tile_y, n_tile_x):
    """Build a lazy dask array by re-running each MFD tile with converged seeds.

    fractions_da is (8, H, W) chunked one tile per (chunks_y, chunks_x)
    block.  The converged boundary snapshot and fraction strips are small,
    so we capture them in a closure and let ``map_blocks`` run the per-tile
    kernel at compute time.  Nothing here materializes the full output
    raster during the API call.
    """
    # Cumulative tile-start offsets to map a block's spatial origin to (iy, ix).
    y_starts = np.cumsum((0,) + tuple(chunks_y[:-1]))
    x_starts = np.cumsum((0,) + tuple(chunks_x[:-1]))

    def _tile(chunk, block_info=None):
        # block_info[0]['array-location'] gives ((0, 8), (y0, y1), (x0, x1)).
        loc = block_info[0]['array-location']
        y0 = loc[1][0]
        x0 = loc[2][0]
        iy = int(np.searchsorted(y_starts, y0, side='right')) - 1
        ix = int(np.searchsorted(x_starts, x0, side='right')) - 1

        chunk = np.asarray(chunk, dtype=np.float64)
        _, h, w = chunk.shape
        seeds = _compute_seeds_mfd(
            iy, ix, boundaries, frac_bdry,
            chunks_y, chunks_x, n_tile_y, n_tile_x)
        return _flow_accum_mfd_tile_kernel(chunk, h, w, *seeds)

    return da.map_blocks(
        _tile, fractions_da, drop_axis=0,
        dtype=np.float64, meta=np.array((), dtype=np.float64),
        **_dask_task_name_kwargs('xrspatial.flow_accumulation_mfd'),
    )


def _flow_accum_mfd_dask_cupy(fractions_da, chunks_y, chunks_x):
    """Dask+CuPy MFD: convert to numpy, run iterative, convert back."""
    import cupy as cp

    fractions_np = fractions_da.map_blocks(
        lambda b: b.get(), dtype=fractions_da.dtype,
        meta=np.array((), dtype=fractions_da.dtype),
    )
    result = _flow_accum_mfd_dask_iterative(fractions_np, chunks_y, chunks_x)
    return result.map_blocks(
        cp.asarray, dtype=result.dtype,
        meta=cp.array((), dtype=result.dtype),
    )


# =====================================================================
# Public API
# =====================================================================

@supports_dataset
def flow_accumulation_mfd(flow_dir_mfd: xr.DataArray,
                           name: str = 'flow_accumulation_mfd') -> xr.DataArray:
    """Compute flow accumulation from an MFD flow direction grid.

    Takes the 3-D fractional output of ``flow_direction_mfd`` and
    accumulates upstream contributing area through all downslope
    paths simultaneously.  Each cell starts with a value of 1 (itself)
    and passes fractions of its accumulated value to each downstream
    neighbor.

    Parameters
    ----------
    flow_dir_mfd : xarray.DataArray or xr.Dataset
        3-D MFD flow direction array of shape ``(8, H, W)`` as returned
        by ``flow_direction_mfd``.  Values are flow fractions in
        ``[0, 1]`` that sum to 1.0 at each cell (0.0 at pits/flats,
        NaN at edges or nodata cells).
        Supported backends: NumPy, CuPy, NumPy-backed Dask,
        CuPy-backed Dask.
        If a Dataset is passed, the operation is applied to each
        data variable independently.
    name : str, default='flow_accumulation_mfd'
        Name of output DataArray.

    Returns
    -------
    xarray.DataArray or xr.Dataset
        2-D float64 array of flow accumulation values.  Each cell
        holds the total upstream contributing area (including itself)
        that drains through it, weighted by MFD fractions.
        NaN where the input has NaN.

    References
    ----------
    Qin, C., Zhu, A.X., Pei, T., Li, B., Zhou, C., and Yang, L.
    (2007). An adaptive approach to selecting a flow-partition
    exponent for a multiple-flow-direction algorithm. International
    Journal of Geographical Information Science, 21(4), 443-458.

    Quinn, P., Beven, K., Chevallier, P., and Planchon, O. (1991).
    The prediction of hillslope flow paths for distributed
    hydrological modelling using digital terrain models.
    Hydrological Processes, 5(1), 59-79.
    """
    _validate_raster(flow_dir_mfd, func_name='flow_accumulation_mfd',
                     name='flow_dir_mfd', ndim=3)

    data = flow_dir_mfd.data

    if data.ndim != 3 or data.shape[0] != 8:
        raise ValueError(
            "flow_dir_mfd must be a 3-D array of shape (8, H, W), "
            f"got shape {data.shape}"
        )

    _validate_mfd_fractions(data, func_name='flow_accumulation_mfd',
                            name='flow_dir_mfd')

    if isinstance(data, np.ndarray):
        _check_memory(data.shape[1], data.shape[2])
        out = _flow_accum_mfd_cpu(
            data.astype(np.float64), data.shape[1], data.shape[2])
    elif has_cuda_and_cupy() and is_cupy_array(data):
        _check_gpu_memory(data.shape[1], data.shape[2])
        out = _flow_accum_mfd_cupy(data)
    elif has_cuda_and_cupy() and is_dask_cupy(flow_dir_mfd):
        # Spatial chunk sizes from dims 1 and 2
        chunks_y = data.chunks[1]
        chunks_x = data.chunks[2]
        out = _flow_accum_mfd_dask_cupy(data, chunks_y, chunks_x)
    elif da is not None and isinstance(data, da.Array):
        chunks_y = data.chunks[1]
        chunks_x = data.chunks[2]
        out = _flow_accum_mfd_dask_iterative(data, chunks_y, chunks_x)
    else:
        raise TypeError(f"Unsupported array type: {type(data)}")

    # Build 2-D output coords (drop 'neighbor' dim)
    spatial_dims = flow_dir_mfd.dims[1:]
    coords = {k: v for k, v in flow_dir_mfd.coords.items()
              if k != 'neighbor' and k not in flow_dir_mfd.dims[:1]}
    # Copy spatial coordinate arrays
    for d in spatial_dims:
        if d in flow_dir_mfd.coords:
            coords[d] = flow_dir_mfd.coords[d]

    return xr.DataArray(out,
                        name=name,
                        coords=coords,
                        dims=spatial_dims,
                        attrs=flow_dir_mfd.attrs)