flow_accumulation_d8.py

"""Flow accumulation: count of upstream cells draining through each cell.

Supports D8 (integer codes) flow direction grids.  For D-infinity
(continuous angles), see ``flow_accumulation_dinf``.

For D8 input, each cell drains to exactly one downstream neighbor.

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 ``cost_distance.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 (
    _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


def _to_numpy_f64(arr):
    """Convert *arr* to a contiguous numpy float64 array.

    Handles CuPy arrays transparently via ``.get()``.
    """
    if hasattr(arr, 'get'):
        arr = arr.get()
    return np.asarray(arr, dtype=np.float64)


# =====================================================================
# Direction helpers
# =====================================================================

@ngjit
def _code_to_offset(code):
    """Return (dy, dx) row/col offset for a D8 direction code."""
    c = int(code)
    if c == 1:
        return 0, 1
    elif c == 2:
        return 1, 1
    elif c == 4:
        return 1, 0
    elif c == 8:
        return 1, -1
    elif c == 16:
        return 0, -1
    elif c == 32:
        return -1, -1
    elif c == 64:
        return -1, 0
    elif c == 128:
        return -1, 1
    return 0, 0


def _code_to_offset_py(code):
    """Pure-Python version for non-numba contexts."""
    c = int(code)
    _map = {1: (0, 1), 2: (1, 1), 4: (1, 0), 8: (1, -1),
            16: (0, -1), 32: (-1, -1), 64: (-1, 0), 128: (-1, 1)}
    return _map.get(c, (0, 0))


# =====================================================================
# Flow type detection (kept for backward compat)
# =====================================================================

@ngjit
def _classify_flow_block(flow_dir, height, width):
    """Classify a block of flow direction values.

    Returns 1 for definite Dinf, -1 for definite D8, 0 for ambiguous
    (all NaN or only zeros, which are valid in both conventions).
    """
    for r in range(height):
        for c in range(width):
            v = flow_dir[r, c]
            if v != v:  # NaN
                continue
            iv = int(v)
            if float(iv) != v:
                return 1  # non-integer -> Dinf
            if iv == 0:
                continue  # ambiguous: D8 pit or Dinf east
            if (iv == 1 or iv == 2 or iv == 4 or iv == 8
                    or iv == 16 or iv == 32 or iv == 64 or iv == 128):
                return -1  # definite D8
            return 1  # integer but not a D8 code (e.g. -1) -> Dinf
    return 0  # ambiguous


def _detect_flow_type(data):
    """Return 'd8' or 'dinf' based on flow direction values."""
    if da is not None and isinstance(data, da.Array):
        for iy in range(len(data.chunks[0])):
            for ix in range(len(data.chunks[1])):
                block = data.blocks[iy, ix].compute()
                sample = np.asarray(
                    block.get() if hasattr(block, 'get') else block)
                result = _classify_flow_block(
                    sample, sample.shape[0], sample.shape[1])
                if result == 1:
                    return "dinf"
                elif result == -1:
                    return "d8"
        return "d8"  # all ambiguous -> default D8
    elif hasattr(data, 'get'):  # cupy
        sample = np.asarray(data.get())
    elif isinstance(data, np.ndarray):
        sample = data
    else:
        return "d8"
    result = _classify_flow_block(sample, sample.shape[0], sample.shape[1])
    return "dinf" if result == 1 else "d8"


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

@ngjit
def _flow_accum_cpu(flow_dir, height, width):
    """Kahn's BFS topological sort for flow accumulation."""
    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
    for r in range(height):
        for c in range(width):
            v = flow_dir[r, c]
            if v == v:  # not NaN
                valid[r, c] = 1
                accum[r, c] = 1.0
            else:
                accum[r, c] = np.nan

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

    # BFS queue (flat arrays with head/tail pointers)
    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

        dy, dx = _code_to_offset(flow_dir[r, c])
        if dy == 0 and dx == 0:
            continue
        nr = r + dy
        nc = c + dx
        if 0 <= nr < height and 0 <= nc < width and valid[nr, nc] == 1:
            accum[nr, nc] += accum[r, c]
            in_degree[nr, nc] -= 1
            if in_degree[nr, nc] == 0:
                queue_r[tail] = nr
                queue_c[tail] = nc
                tail += 1

    return accum


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

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

    v = flow_dir[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

    # Decode direction (inline -- can't call @ngjit from @cuda.jit)
    code = int(v)
    dy = 0
    dx = 0
    if code == 1:
        dy, dx = 0, 1
    elif code == 2:
        dy, dx = 1, 1
    elif code == 4:
        dy, dx = 1, 0
    elif code == 8:
        dy, dx = 1, -1
    elif code == 16:
        dy, dx = 0, -1
    elif code == 32:
        dy, dx = -1, -1
    elif code == 64:
        dy, dx = -1, 0
    elif code == 128:
        dy, dx = -1, 1

    if dy == 0 and dx == 0:
        return  # pit

    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(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(flow_dir, accum, in_degree, state, H, W):
    """Active cells pull accumulation from frontier neighbours."""
    i, j = cuda.grid(2)
    if i >= H or j >= W:
        return

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

    # Check all 8 neighbours
    for k in range(8):
        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 ni < 0 or ni >= H or nj < 0 or nj >= W:
            continue

        if state[ni, nj] != 2:
            continue

        # Check if neighbour's flow_dir points to me
        nv = flow_dir[ni, nj]
        ncode = int(nv)
        ndy = 0
        ndx = 0
        if ncode == 1:
            ndy, ndx = 0, 1
        elif ncode == 2:
            ndy, ndx = 1, 1
        elif ncode == 4:
            ndy, ndx = 1, 0
        elif ncode == 8:
            ndy, ndx = 1, -1
        elif ncode == 16:
            ndy, ndx = 0, -1
        elif ncode == 32:
            ndy, ndx = -1, -1
        elif ncode == 64:
            ndy, ndx = -1, 0
        elif ncode == 128:
            ndy, ndx = -1, 1

        if ni + ndy == i and nj + ndx == j:
            accum[i, j] += accum[ni, nj]
            in_degree[i, j] -= 1


def _flow_accum_cupy(flow_dir_data):
    """GPU driver: iterative frontier peeling."""
    import cupy as cp

    H, W = flow_dir_data.shape
    flow_dir_f64 = flow_dir_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[griddim, blockdim](
        flow_dir_f64, accum, in_degree, state, H, W)

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

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

        _pull_from_frontier[griddim, blockdim](
            flow_dir_f64, accum, in_degree, state, H, W)

    # Convert invalid cells to NaN
    accum = cp.where(state == 0, cp.nan, accum)
    return accum


def _flow_accum_tile_cupy(flow_dir_data,
                           seed_top, seed_bottom, seed_left, seed_right,
                           seed_tl, seed_tr, seed_bl, seed_br):
    """GPU seeded flow accumulation for a single tile.

    Same algorithm as ``_flow_accum_cupy`` but injects external seed
    values at boundary cells before frontier peeling.  Seeds are
    NumPy arrays; they are transferred to GPU inside this function.
    """
    import cupy as cp

    H, W = flow_dir_data.shape
    flow_dir_f64 = flow_dir_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[griddim, blockdim](
        flow_dir_f64, accum, in_degree, state, H, W)

    # Inject seeds at boundary cells.  Invalid cells (state==0) are
    # masked to NaN at the end and never enter frontier peeling, so
    # adding seeds to them is harmless.
    accum[0, :] += cp.asarray(seed_top)
    accum[H - 1, :] += cp.asarray(seed_bottom)
    accum[:, 0] += cp.asarray(seed_left)
    accum[:, W - 1] += cp.asarray(seed_right)
    accum[0, 0] += seed_tl
    accum[0, W - 1] += seed_tr
    accum[H - 1, 0] += seed_bl
    accum[H - 1, W - 1] += seed_br

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

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

        _pull_from_frontier[griddim, blockdim](
            flow_dir_f64, accum, in_degree, state, H, W)

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


# =====================================================================
# Tile kernel for dask iterative path
# =====================================================================

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

    Same as ``_flow_accum_cpu`` but adds external seeds to boundary
    cells before BFS.
    """
    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
    for r in range(h):
        for c in range(w):
            v = flow_dir[r, c]
            if v == v:
                valid[r, c] = 1
                accum[r, c] = 1.0
            else:
                accum[r, c] = np.nan

    # Add external seeds to boundary cells
    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]

    # Corner seeds
    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
            dy, dx = _code_to_offset(flow_dir[r, c])
            if dy == 0 and dx == 0:
                continue
            nr = r + dy
            nc = c + dx
            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

        dy, dx = _code_to_offset(flow_dir[r, c])
        if dy == 0 and dx == 0:
            continue
        nr = r + dy
        nc = c + dx
        if 0 <= nr < h and 0 <= nc < w and valid[nr, nc] == 1:
            accum[nr, nc] += accum[r, c]
            in_degree[nr, nc] -= 1
            if in_degree[nr, nc] == 0:
                queue_r[tail] = nr
                queue_c[tail] = nc
                tail += 1

    return accum


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

def _preprocess_tiles(flow_dir_da, chunks_y, chunks_x):
    """Extract boundary flow-direction strips into a BoundaryStore."""
    n_tile_y = len(chunks_y)
    n_tile_x = len(chunks_x)

    flow_bdry = BoundaryStore(chunks_y, chunks_x, fill_value=np.nan)

    for iy in range(n_tile_y):
        for ix in range(n_tile_x):
            chunk = flow_dir_da.blocks[iy, ix].compute()
            flow_bdry.set('top', iy, ix, _to_numpy_f64(chunk[0, :]))
            flow_bdry.set('bottom', iy, ix, _to_numpy_f64(chunk[-1, :]))
            flow_bdry.set('left', iy, ix, _to_numpy_f64(chunk[:, 0]))
            flow_bdry.set('right', iy, ix, _to_numpy_f64(chunk[:, -1]))

    return flow_bdry


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

    For each boundary cell of the current tile, checks whether cells
    in adjacent tiles flow INTO this tile.  If so, adds the adjacent
    cell's boundary accum value as a seed.

    Returns (seed_top, seed_bottom, seed_left, seed_right,
             seed_tl, seed_tr, seed_bl, seed_br).
    """
    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

    # --- Top edge: bottom of tile above ---
    if iy > 0:
        nb_fdir = flow_bdry.get('bottom', iy - 1, ix)
        nb_accum = boundaries.get('bottom', iy - 1, ix)
        w = len(nb_fdir)
        # Cardinal S (code 4): nb cell j -> our (0, j)
        mask = (nb_fdir == 4)
        seed_top += np.where(mask, nb_accum, 0.0)
        if w > 1:
            # Diagonal SE (code 2): nb cell j -> our (0, j+1)
            mask = (nb_fdir[:-1] == 2)
            seed_top[1:] += np.where(mask, nb_accum[:-1], 0.0)
            # Diagonal SW (code 8): nb cell j -> our (0, j-1)
            mask = (nb_fdir[1:] == 8)
            seed_top[:-1] += np.where(mask, nb_accum[1:], 0.0)

    # --- Bottom edge: top of tile below ---
    if iy < n_tile_y - 1:
        nb_fdir = flow_bdry.get('top', iy + 1, ix)
        nb_accum = boundaries.get('top', iy + 1, ix)
        w = len(nb_fdir)
        # Cardinal N (code 64)
        mask = (nb_fdir == 64)
        seed_bottom += np.where(mask, nb_accum, 0.0)
        if w > 1:
            # Diagonal NE (code 128): nb cell j -> our (h-1, j+1)
            mask = (nb_fdir[:-1] == 128)
            seed_bottom[1:] += np.where(mask, nb_accum[:-1], 0.0)
            # Diagonal NW (code 32): nb cell j -> our (h-1, j-1)
            mask = (nb_fdir[1:] == 32)
            seed_bottom[:-1] += np.where(mask, nb_accum[1:], 0.0)

    # --- Left edge: right column of tile to the left ---
    if ix > 0:
        nb_fdir = flow_bdry.get('right', iy, ix - 1)
        nb_accum = boundaries.get('right', iy, ix - 1)
        h = len(nb_fdir)
        # Cardinal E (code 1)
        mask = (nb_fdir == 1)
        seed_left += np.where(mask, nb_accum, 0.0)
        if h > 1:
            # Diagonal SE (code 2): nb cell i -> our (i+1, 0)
            mask = (nb_fdir[:-1] == 2)
            seed_left[1:] += np.where(mask, nb_accum[:-1], 0.0)
            # Diagonal NE (code 128): nb cell i -> our (i-1, 0)
            mask = (nb_fdir[1:] == 128)
            seed_left[:-1] += np.where(mask, nb_accum[1:], 0.0)

    # --- Right edge: left column of tile to the right ---
    if ix < n_tile_x - 1:
        nb_fdir = flow_bdry.get('left', iy, ix + 1)
        nb_accum = boundaries.get('left', iy, ix + 1)
        h = len(nb_fdir)
        # Cardinal W (code 16)
        mask = (nb_fdir == 16)
        seed_right += np.where(mask, nb_accum, 0.0)
        if h > 1:
            # Diagonal SW (code 8): nb cell i -> our (i+1, w-1)
            mask = (nb_fdir[:-1] == 8)
            seed_right[1:] += np.where(mask, nb_accum[:-1], 0.0)
            # Diagonal NW (code 32): nb cell i -> our (i-1, w-1)
            mask = (nb_fdir[1:] == 32)
            seed_right[:-1] += np.where(mask, nb_accum[1:], 0.0)

    # --- Diagonal corner seeds ---
    # TL: bottom-right of (iy-1, ix-1) flows SE (code 2)
    if iy > 0 and ix > 0:
        fdir = flow_bdry.get('bottom', iy - 1, ix - 1)[-1]
        if fdir == 2:
            seed_tl = float(boundaries.get('bottom', iy - 1, ix - 1)[-1])
    # TR: bottom-left of (iy-1, ix+1) flows SW (code 8)
    if iy > 0 and ix < n_tile_x - 1:
        fdir = flow_bdry.get('bottom', iy - 1, ix + 1)[0]
        if fdir == 8:
            seed_tr = float(boundaries.get('bottom', iy - 1, ix + 1)[0])
    # BL: top-right of (iy+1, ix-1) flows NE (code 128)
    if iy < n_tile_y - 1 and ix > 0:
        fdir = flow_bdry.get('top', iy + 1, ix - 1)[-1]
        if fdir == 128:
            seed_bl = float(boundaries.get('top', iy + 1, ix - 1)[-1])
    # BR: top-left of (iy+1, ix+1) flows NW (code 32)
    if iy < n_tile_y - 1 and ix < n_tile_x - 1:
        fdir = flow_bdry.get('top', iy + 1, ix + 1)[0]
        if fdir == 32:
            seed_br = float(boundaries.get('top', iy + 1, ix + 1)[0])

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


def _process_tile(iy, ix, flow_dir_da, boundaries, flow_bdry,
                  chunks_y, chunks_x, n_tile_y, n_tile_x):
    """Run seeded BFS on one tile; update boundaries in-place.

    Returns the maximum absolute boundary change (float).
    """
    chunk = np.asarray(
        flow_dir_da.blocks[iy, ix].compute(), dtype=np.float64)
    h, w = chunk.shape

    seeds = _compute_seeds(
        iy, ix, boundaries, flow_bdry,
        chunks_y, chunks_x, n_tile_y, n_tile_x)

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

    # Extract new boundary strips
    new_top = accum[0, :].copy()
    new_bottom = accum[-1, :].copy()
    new_left = accum[:, 0].copy()
    new_right = accum[:, -1].copy()

    # Compute max absolute change
    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

    # Store updated boundaries
    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_dask_iterative(flow_dir_da):
    """Iterative boundary-propagation for arbitrarily large dask arrays.

    Memory usage is O(tile_size + boundary_strips) per iteration.
    """
    chunks_y = flow_dir_da.chunks[0]
    chunks_x = flow_dir_da.chunks[1]
    n_tile_y = len(chunks_y)
    n_tile_x = len(chunks_x)

    # Phase 0: extract boundary flow dirs
    flow_bdry = _preprocess_tiles(flow_dir_da, chunks_y, chunks_x)
    flow_bdry = flow_bdry.snapshot()  # read-only from here; release temp files

    # 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

        # Forward sweep (top-left -> bottom-right)
        for iy in range(n_tile_y):
            for ix in range(n_tile_x):
                c = _process_tile(iy, ix, flow_dir_da, boundaries,
                                  flow_bdry, chunks_y, chunks_x,
                                  n_tile_y, n_tile_x)
                if c > max_change:
                    max_change = c

        # Backward sweep (bottom-right -> top-left)
        for iy in reversed(range(n_tile_y)):
            for ix in reversed(range(n_tile_x)):
                c = _process_tile(iy, ix, flow_dir_da, boundaries,
                                  flow_bdry, chunks_y, chunks_x,
                                  n_tile_y, n_tile_x)
                if c > max_change:
                    max_change = c

        if max_change == 0.0:
            break

    # Snapshot converged boundaries before assembly (releases temp files)
    boundaries = boundaries.snapshot()

    # Phase 3: lazy assembly via da.map_blocks
    return _assemble_result(flow_dir_da, boundaries, flow_bdry,
                            chunks_y, chunks_x, n_tile_y, n_tile_x)


def _assemble_result(flow_dir_da, boundaries, flow_bdry,
                     chunks_y, chunks_x, n_tile_y, n_tile_x):
    """Build a lazy dask array by re-running each tile with converged seeds."""

    def _tile_fn(flow_dir_block, block_info=None):
        if block_info is None or 0 not in block_info:
            return np.full(flow_dir_block.shape, np.nan, dtype=np.float64)
        iy, ix = block_info[0]['chunk-location']
        h, w = flow_dir_block.shape
        seeds = _compute_seeds(
            iy, ix, boundaries, flow_bdry,
            chunks_y, chunks_x, n_tile_y, n_tile_x)
        return _flow_accum_tile_kernel(
            np.asarray(flow_dir_block, dtype=np.float64), h, w, *seeds)

    return da.map_blocks(
        _tile_fn,
        flow_dir_da,
        dtype=np.float64,
        meta=np.array((), dtype=np.float64),
    )


def _process_tile_cupy(iy, ix, flow_dir_da, boundaries, flow_bdry,
                        chunks_y, chunks_x, n_tile_y, n_tile_x):
    """Run seeded GPU flow accumulation on one tile; update boundaries."""
    import cupy as cp

    chunk = cp.asarray(
        flow_dir_da.blocks[iy, ix].compute(), dtype=cp.float64)

    seeds = _compute_seeds(
        iy, ix, boundaries, flow_bdry,
        chunks_y, chunks_x, n_tile_y, n_tile_x)

    accum = _flow_accum_tile_cupy(chunk, *seeds)

    # Extract boundaries to CPU (small 1-D strips)
    new_top = accum[0, :].get().copy()
    new_bottom = accum[-1, :].get().copy()
    new_left = accum[:, 0].get().copy()
    new_right = accum[:, -1].get().copy()

    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 _assemble_result_cupy(flow_dir_da, boundaries, flow_bdry,
                           chunks_y, chunks_x, n_tile_y, n_tile_x):
    """Build a lazy dask+cupy array using GPU tile kernel."""
    import cupy as cp

    def _tile_fn(flow_dir_block, block_info=None):
        if block_info is None or 0 not in block_info:
            return cp.full(flow_dir_block.shape, cp.nan, dtype=cp.float64)
        iy, ix = block_info[0]['chunk-location']
        seeds = _compute_seeds(
            iy, ix, boundaries, flow_bdry,
            chunks_y, chunks_x, n_tile_y, n_tile_x)
        return _flow_accum_tile_cupy(
            cp.asarray(flow_dir_block, dtype=cp.float64), *seeds)

    return da.map_blocks(
        _tile_fn,
        flow_dir_da,
        dtype=np.float64,
        meta=cp.array((), dtype=cp.float64),
    )


def _flow_accum_dask_cupy(flow_dir_da):
    """Dask+CuPy D8: native GPU processing per tile."""
    chunks_y = flow_dir_da.chunks[0]
    chunks_x = flow_dir_da.chunks[1]
    n_tile_y = len(chunks_y)
    n_tile_x = len(chunks_x)

    flow_bdry = _preprocess_tiles(flow_dir_da, chunks_y, chunks_x)
    flow_bdry = flow_bdry.snapshot()

    boundaries = BoundaryStore(chunks_y, chunks_x, fill_value=0.0)

    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_cupy(iy, ix, flow_dir_da, boundaries,
                                        flow_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_cupy(iy, ix, flow_dir_da, boundaries,
                                        flow_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_cupy(flow_dir_da, boundaries, flow_bdry,
                                  chunks_y, chunks_x, n_tile_y, n_tile_x)


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

@supports_dataset
def flow_accumulation_d8(flow_dir: xr.DataArray,
                      name: str = 'flow_accumulation') -> xr.DataArray:
    """Compute flow accumulation from a D8 flow direction grid.

    Each cell drains to exactly one downstream neighbor based on
    integer D8 direction codes.  For D-infinity (continuous angle)
    grids, use ``flow_accumulation_dinf`` instead.

    Parameters
    ----------
    flow_dir : xarray.DataArray or xr.Dataset
        2D flow direction grid with D8 codes:
        0/1/2/4/8/16/32/64/128 (NaN for nodata).
        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'
        Name of output DataArray.

    Returns
    -------
    xarray.DataArray or xr.Dataset
        2D float64 array of flow accumulation values.  Each cell
        contains the count of upstream cells (including itself) that
        drain through it.  Cells with NaN flow direction produce NaN.

    References
    ----------
    Jenson, S.K. and Domingue, J.O. (1988). Extracting Topographic
    Structure from Digital Elevation Data for Geographic Information
    System Analysis. Photogrammetric Engineering and Remote Sensing,
    54(11), 1593-1600.
    """
    _validate_raster(flow_dir, func_name='flow_accumulation', name='flow_dir')

    data = flow_dir.data

    if isinstance(data, np.ndarray):
        out = _flow_accum_cpu(data.astype(np.float64), *data.shape)
    elif has_cuda_and_cupy() and is_cupy_array(data):
        out = _flow_accum_cupy(data)
    elif has_cuda_and_cupy() and is_dask_cupy(flow_dir):
        out = _flow_accum_dask_cupy(data)
    elif da is not None and isinstance(data, da.Array):
        out = _flow_accum_dask_iterative(data)
    else:
        raise TypeError(f"Unsupported array type: {type(data)}")

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