surface_distance.py

"""Surface distance along 3D terrain via multi-source Dijkstra.

Computes the minimum accumulated distance along the terrain surface
from each pixel to the nearest target pixel, accounting for elevation
relief.  A steep hillside has more surface distance than its flat
map projection.

Algorithm
---------
Multi-source Dijkstra with edge costs derived from elevation::

    edge_cost = sqrt(horizontal_dist² + (elev_v − elev_u)²)

NaN elevation marks impassable pixels (barriers).

Dask strategy
-------------
For finite ``max_distance``, the maximum pixel radius any path can
reach is ``max_distance / min_cellsize`` (since surface distance >=
horizontal distance).  This becomes the ``depth`` parameter to
``dask.array.map_overlap``, giving exact results.

If ``max_distance`` is infinite or the implied radius exceeds chunk
dimensions, an iterative boundary-only Dijkstra is used (same
tile-sweep pattern as ``cost_distance``).
"""

from __future__ import annotations

import math as _math
import warnings
from math import sqrt

import numpy as np
import xarray as xr

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

from numba import cuda

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

from xrspatial.cost_distance import _heap_push, _heap_pop
from xrspatial.proximity import _vectorized_calc_direction
from xrspatial.utils import (
    _validate_raster,
    cuda_args, get_dataarray_resolution, ngjit,
    has_cuda_and_cupy, is_cupy_array, is_dask_cupy,
)
from xrspatial.dataset_support import supports_dataset

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

DISTANCE = 0
ALLOCATION = 1
DIRECTION = 2

# ---------------------------------------------------------------------------
# Numba kernels
# ---------------------------------------------------------------------------


@ngjit
def _seed_sources(source_data, elev_data, target_values,
                  dist, alloc, src_row, src_col):
    """Seed source pixels into pre-allocated output arrays.

    Source pixels: dist=0, alloc=pixel_value, src_row/col=pixel position.
    Only seeds if elevation is finite (passable).
    """
    height, width = source_data.shape
    n_values = len(target_values)

    for r in range(height):
        for c in range(width):
            val = source_data[r, c]
            is_target = False
            if n_values == 0:
                if val != 0.0 and np.isfinite(val):
                    is_target = True
            else:
                for k in range(n_values):
                    if val == target_values[k]:
                        is_target = True
                        break
            if is_target:
                if np.isfinite(elev_data[r, c]):
                    dist[r, c] = 0.0
                    alloc[r, c] = val
                    src_row[r, c] = r
                    src_col[r, c] = c


@ngjit
def _dijkstra(elev_data, height, width, max_distance,
              dy, dx, dd, dist, alloc, src_row, src_col):
    """Planar multi-source Dijkstra.  Modifies arrays in-place.

    Pre-seeded pixels (dist < inf) are added to the heap.
    Edge cost = sqrt(dd[i]^2 + dz^2).
    """
    n_neighbors = len(dy)

    max_heap = height * width
    h_keys = np.empty(max_heap, dtype=np.float64)
    h_rows = np.empty(max_heap, dtype=np.int64)
    h_cols = np.empty(max_heap, dtype=np.int64)
    h_size = 0

    visited = np.zeros((height, width), dtype=np.int8)

    # Add all pre-seeded pixels to the heap
    for r in range(height):
        for c in range(width):
            if dist[r, c] < np.inf:
                h_size = _heap_push(h_keys, h_rows, h_cols, h_size,
                                    dist[r, c], r, c)

    # Dijkstra main loop
    while h_size > 0:
        cost_u, ur, uc, h_size = _heap_pop(h_keys, h_rows, h_cols, h_size)

        if visited[ur, uc]:
            continue
        visited[ur, uc] = 1

        if cost_u > max_distance:
            break

        elev_u = elev_data[ur, uc]

        for i in range(n_neighbors):
            vr = ur + dy[i]
            vc = uc + dx[i]
            if vr < 0 or vr >= height or vc < 0 or vc >= width:
                continue
            if visited[vr, vc]:
                continue

            elev_v = elev_data[vr, vc]
            if not np.isfinite(elev_v):
                continue

            dz = elev_v - elev_u
            edge_cost = np.sqrt(dd[i] * dd[i] + dz * dz)
            new_cost = cost_u + edge_cost

            if new_cost < dist[vr, vc]:
                dist[vr, vc] = new_cost
                alloc[vr, vc] = alloc[ur, uc]
                src_row[vr, vc] = src_row[ur, uc]
                src_col[vr, vc] = src_col[ur, uc]
                h_size = _heap_push(h_keys, h_rows, h_cols, h_size,
                                    new_cost, vr, vc)


@ngjit
def _dijkstra_geodesic(elev_data, height, width, max_distance,
                       dy, dx, dd_grid, dist, alloc, src_row, src_col):
    """Geodesic variant with per-pixel horizontal distances.

    dd_grid[i, r, c] = great-circle horizontal distance to neighbour i
    from pixel (r, c).
    """
    n_neighbors = len(dy)

    max_heap = height * width
    h_keys = np.empty(max_heap, dtype=np.float64)
    h_rows = np.empty(max_heap, dtype=np.int64)
    h_cols = np.empty(max_heap, dtype=np.int64)
    h_size = 0

    visited = np.zeros((height, width), dtype=np.int8)

    for r in range(height):
        for c in range(width):
            if dist[r, c] < np.inf:
                h_size = _heap_push(h_keys, h_rows, h_cols, h_size,
                                    dist[r, c], r, c)

    while h_size > 0:
        cost_u, ur, uc, h_size = _heap_pop(h_keys, h_rows, h_cols, h_size)

        if visited[ur, uc]:
            continue
        visited[ur, uc] = 1

        if cost_u > max_distance:
            break

        elev_u = elev_data[ur, uc]

        for i in range(n_neighbors):
            vr = ur + dy[i]
            vc = uc + dx[i]
            if vr < 0 or vr >= height or vc < 0 or vc >= width:
                continue
            if visited[vr, vc]:
                continue

            elev_v = elev_data[vr, vc]
            if not np.isfinite(elev_v):
                continue

            hdist = dd_grid[i, ur, uc]
            dz = elev_v - elev_u
            edge_cost = np.sqrt(hdist * hdist + dz * dz)
            new_cost = cost_u + edge_cost

            if new_cost < dist[vr, vc]:
                dist[vr, vc] = new_cost
                alloc[vr, vc] = alloc[ur, uc]
                src_row[vr, vc] = src_row[ur, uc]
                src_col[vr, vc] = src_col[ur, uc]
                h_size = _heap_push(h_keys, h_rows, h_cols, h_size,
                                    new_cost, vr, vc)


# ---------------------------------------------------------------------------
# Post-processing helpers
# ---------------------------------------------------------------------------


def _init_arrays(H, W):
    """Create and initialize output arrays for the Dijkstra kernel."""
    dist = np.full((H, W), np.inf, dtype=np.float64)
    alloc = np.full((H, W), np.nan, dtype=np.float64)
    src_row = np.full((H, W), -1, dtype=np.int64)
    src_col = np.full((H, W), -1, dtype=np.int64)
    return dist, alloc, src_row, src_col


def _finalize_dist(dist, max_distance):
    """Convert float64 dist to float32; inf / over-budget -> NaN."""
    out = np.where(
        np.isinf(dist) | (dist > max_distance), np.nan, dist,
    ).astype(np.float32)
    return out


def _finalize_alloc(alloc, dist, max_distance):
    """Float32 allocation; NaN where unreachable."""
    out = np.where(
        np.isinf(dist) | (dist > max_distance), np.nan, alloc,
    ).astype(np.float32)
    return out


def _finalize_direction(src_row, src_col, dist, cellsize_x, cellsize_y,
                        max_distance):
    """Compute compass bearing from each pixel to its allocated source.

    Uses pixel index differences scaled by cell size.
    """
    H, W = dist.shape
    row_idx, col_idx = np.meshgrid(np.arange(H), np.arange(W), indexing='ij')

    # Coordinate differences (source - pixel)
    dx = (src_col.astype(np.float64) - col_idx) * cellsize_x
    dy = (src_row.astype(np.float64) - row_idx) * cellsize_y

    result = _vectorized_calc_direction(
        np.zeros((H, W), dtype=np.float64), dx,
        np.zeros((H, W), dtype=np.float64), dy,
    )

    # Mask unreachable and no-source pixels
    mask = np.isinf(dist) | (dist > max_distance) | (src_row < 0)
    result[mask] = np.nan
    return result


def _extract_output(dist, alloc, src_row, src_col,
                    cellsize_x, cellsize_y, max_distance, mode):
    """Select and finalize the requested output from raw Dijkstra arrays."""
    if mode == DISTANCE:
        return _finalize_dist(dist, max_distance)
    elif mode == ALLOCATION:
        return _finalize_alloc(alloc, dist, max_distance)
    else:
        return _finalize_direction(src_row, src_col, dist,
                                   cellsize_x, cellsize_y, max_distance)


# ---------------------------------------------------------------------------
# Geodesic dd_grid precomputation
# ---------------------------------------------------------------------------

EARTH_RADIUS = 6378137.0  # meters


def _precompute_dd_grid(lat_2d, lon_2d, dy, dx):
    """Precompute per-pixel great-circle horizontal distances.

    Returns dd_grid[n_neighbors, H, W] in meters.
    """
    H, W = lat_2d.shape
    n = len(dy)
    # Memory guard: dd_grid is (n_neighbors, H, W) float64
    estimated = n * H * W * 8
    try:
        from xrspatial.zonal import _available_memory_bytes
        avail = _available_memory_bytes()
    except ImportError:
        avail = 2 * 1024**3
    if estimated > 0.8 * avail:
        raise MemoryError(
            f"Geodesic dd_grid needs ~{estimated / 1e9:.1f} GB "
            f"({n} neighbors x {H}x{W} x 8 bytes) but only "
            f"~{avail / 1e9:.1f} GB available.  Use planar mode "
            f"or downsample the raster."
        )
    dd_grid = np.zeros((n, H, W), dtype=np.float64)

    for i in range(n):
        dr, dc = int(dy[i]), int(dx[i])
        # Source region
        r0 = max(0, -dr)
        r1 = H - max(0, dr)
        c0 = max(0, -dc)
        c1 = W - max(0, dc)
        # Neighbour region
        nr0 = max(0, dr)
        nr1 = H - max(0, -dr)
        nc0 = max(0, dc)
        nc1 = W - max(0, -dc)

        lat1 = np.radians(lat_2d[r0:r1, c0:c1])
        lon1 = np.radians(lon_2d[r0:r1, c0:c1])
        lat2 = np.radians(lat_2d[nr0:nr1, nc0:nc1])
        lon2 = np.radians(lon_2d[nr0:nr1, nc0:nc1])

        dlat = lat2 - lat1
        dlon = lon2 - lon1
        a = (np.sin(dlat / 2) ** 2
             + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2) ** 2)
        dd_grid[i, r0:r1, c0:c1] = (
            EARTH_RADIUS * 2 * np.arctan2(np.sqrt(a), np.sqrt(1 - a))
        )

    return dd_grid


# ---------------------------------------------------------------------------
# NumPy wrapper
# ---------------------------------------------------------------------------


def _surface_distance_numpy(source_data, elev_data, cellsize_x, cellsize_y,
                            max_distance, target_values, dy, dx, dd,
                            dd_grid, use_geodesic, mode):
    """NumPy backend: run Dijkstra and extract requested output."""
    H, W = source_data.shape
    dist, alloc, src_row, src_col = _init_arrays(H, W)

    _seed_sources(source_data, elev_data, target_values,
                  dist, alloc, src_row, src_col)

    if use_geodesic:
        _dijkstra_geodesic(elev_data, H, W, max_distance,
                           dy, dx, dd_grid, dist, alloc, src_row, src_col)
    else:
        _dijkstra(elev_data, H, W, max_distance,
                  dy, dx, dd, dist, alloc, src_row, src_col)

    return _extract_output(dist, alloc, src_row, src_col,
                           cellsize_x, cellsize_y, max_distance, mode)


# ---------------------------------------------------------------------------
# CuPy GPU backend — iterative parallel relaxation
# ---------------------------------------------------------------------------


@cuda.jit
def _sd_relax_kernel(elev, dist, alloc, src_row, src_col, changed,
                     height, width, dy, dx, dd, n_neighbors,
                     max_distance):
    """One relaxation pass for surface distance.

    Each pixel checks all neighbours for shorter 3D paths.
    Iterate until *changed* stays 0.
    """
    iy, ix = cuda.grid(2)
    if iy >= height or ix >= width:
        return

    elev_u = elev[iy, ix]
    if not _math.isfinite(elev_u):
        return

    current = dist[iy, ix]
    best = current
    best_alloc = alloc[iy, ix]
    best_srow = src_row[iy, ix]
    best_scol = src_col[iy, ix]

    for k in range(n_neighbors):
        vy = iy + dy[k]
        vx = ix + dx[k]
        if vy < 0 or vy >= height or vx < 0 or vx >= width:
            continue

        d_v = dist[vy, vx]
        if d_v >= best:
            continue

        elev_v = elev[vy, vx]
        if not _math.isfinite(elev_v):
            continue

        dz = elev_u - elev_v
        edge_cost = _math.sqrt(dd[k] * dd[k] + dz * dz)
        new_cost = d_v + edge_cost

        if new_cost < best:
            best = new_cost
            best_alloc = alloc[vy, vx]
            best_srow = src_row[vy, vx]
            best_scol = src_col[vy, vx]

    if best < current and best <= max_distance:
        dist[iy, ix] = best
        alloc[iy, ix] = best_alloc
        src_row[iy, ix] = best_srow
        src_col[iy, ix] = best_scol
        changed[0] = 1


def _surface_distance_cupy(source_data, elev_data, cellsize_x, cellsize_y,
                           max_distance, target_values, dy, dx, dd, mode):
    """GPU surface distance via iterative parallel relaxation."""
    import cupy as cp

    H, W = source_data.shape
    src = source_data.astype(cp.float64)
    elev = elev_data.astype(cp.float64)

    dist = cp.full((H, W), cp.inf, dtype=cp.float64)
    alloc_arr = cp.full((H, W), cp.nan, dtype=cp.float64)
    srow = cp.full((H, W), -1, dtype=cp.int64)
    scol = cp.full((H, W), -1, dtype=cp.int64)

    # Seed sources
    if len(target_values) == 0:
        mask = cp.isfinite(src) & (src != 0) & cp.isfinite(elev)
    else:
        tv = cp.asarray(target_values, dtype=cp.float64)
        mask = cp.isin(src, tv) & cp.isfinite(elev)

    dist[mask] = 0.0
    alloc_arr[mask] = src[mask]
    rows_g, cols_g = cp.meshgrid(cp.arange(H, dtype=cp.int64),
                                 cp.arange(W, dtype=cp.int64),
                                 indexing='ij')
    srow[mask] = rows_g[mask]
    scol[mask] = cols_g[mask]

    if not cp.any(mask):
        out = cp.full((H, W), cp.nan, dtype=cp.float32)
        return out

    dy_d = cp.asarray(dy, dtype=cp.int64)
    dx_d = cp.asarray(dx, dtype=cp.int64)
    dd_d = cp.asarray(dd, dtype=cp.float64)
    n_neighbors = len(dy)

    changed = cp.zeros(1, dtype=cp.int32)
    griddim, blockdim = cuda_args((H, W))

    max_iterations = H + W
    for _ in range(max_iterations):
        changed[0] = 0
        _sd_relax_kernel[griddim, blockdim](
            elev, dist, alloc_arr, srow, scol, changed,
            H, W,
            dy_d, dx_d, dd_d, n_neighbors,
            np.float64(max_distance),
        )
        if int(changed[0]) == 0:
            break

    # Extract output
    if mode == DISTANCE:
        out = cp.where(cp.isinf(dist) | (dist > max_distance),
                       cp.nan, dist).astype(cp.float32)
    elif mode == ALLOCATION:
        out = cp.where(cp.isinf(dist) | (dist > max_distance),
                       cp.nan, alloc_arr).astype(cp.float32)
    else:  # DIRECTION
        row_idx, col_idx = cp.meshgrid(cp.arange(H, dtype=cp.float64),
                                       cp.arange(W, dtype=cp.float64),
                                       indexing='ij')
        dx_coord = (scol.astype(cp.float64) - col_idx) * cellsize_x
        dy_coord = (srow.astype(cp.float64) - row_idx) * cellsize_y
        # Compute direction on CPU (uses numpy-based vectorized function)
        dx_np = cp.asnumpy(dx_coord)
        dy_np = cp.asnumpy(dy_coord)
        zeros = np.zeros((H, W), dtype=np.float64)
        result = _vectorized_calc_direction(zeros, dx_np, zeros, dy_np)
        mask_np = cp.asnumpy(cp.isinf(dist) | (dist > max_distance)
                             | (srow < 0))
        result[mask_np] = np.nan
        out = cp.asarray(result)

    return out


# ---------------------------------------------------------------------------
# Dask bounded — map_overlap
# ---------------------------------------------------------------------------


def _make_sd_chunk_func(cellsize_x, cellsize_y, max_distance,
                        target_values, dy, dx, dd, mode):
    """Return a function for ``da.map_overlap`` over source + elev."""

    def _chunk(source_block, elev_block):
        return _surface_distance_numpy(
            source_block, elev_block,
            cellsize_x, cellsize_y, max_distance,
            target_values, dy, dx, dd,
            None, False, mode,
        )

    return _chunk


def _surface_distance_dask_bounded(source_da, elev_da,
                                   cellsize_x, cellsize_y,
                                   max_distance, target_values,
                                   dy, dx, dd, mode):
    """Dask bounded path via map_overlap."""
    min_cellsize = min(abs(cellsize_x), abs(cellsize_y))
    pad = int(max_distance / min_cellsize) + 1

    chunk_func = _make_sd_chunk_func(
        cellsize_x, cellsize_y, max_distance,
        target_values, dy, dx, dd, mode,
    )

    return da.map_overlap(
        chunk_func,
        source_da, elev_da,
        depth=(pad, pad),
        boundary=np.nan,
        dtype=np.float32,
        meta=np.array((), dtype=np.float32),
    )


# ---------------------------------------------------------------------------
# Iterative boundary-only Dijkstra for dask arrays
# ---------------------------------------------------------------------------


def _preprocess_tiles_sd(source_da, elev_da, chunks_y, chunks_x,
                         target_values):
    """Extract elevation boundary strips and identify source tiles."""
    n_tile_y = len(chunks_y)
    n_tile_x = len(chunks_x)
    n_values = len(target_values)

    elev_bdry = {
        side: [[None] * n_tile_x for _ in range(n_tile_y)]
        for side in ('top', 'bottom', 'left', 'right')
    }
    has_source = [[False] * n_tile_x for _ in range(n_tile_y)]

    for iy in range(n_tile_y):
        for ix in range(n_tile_x):
            echunk = elev_da.blocks[iy, ix].compute()
            elev_bdry['top'][iy][ix] = echunk[0, :].astype(np.float64)
            elev_bdry['bottom'][iy][ix] = echunk[-1, :].astype(np.float64)
            elev_bdry['left'][iy][ix] = echunk[:, 0].astype(np.float64)
            elev_bdry['right'][iy][ix] = echunk[:, -1].astype(np.float64)

            schunk = source_da.blocks[iy, ix].compute()
            if n_values == 0:
                has_source[iy][ix] = bool(
                    np.any((schunk != 0) & np.isfinite(schunk)
                           & np.isfinite(echunk))
                )
            else:
                for tv in target_values:
                    if np.any((schunk == tv) & np.isfinite(echunk)):
                        has_source[iy][ix] = True
                        break

    return elev_bdry, has_source


def _init_boundaries_sd(chunks_y, chunks_x):
    """Create boundary arrays: dist, alloc, src_row, src_col."""
    n_y = len(chunks_y)
    n_x = len(chunks_x)

    def _make(side_sizes, fill, dtype):
        return [[np.full(side_sizes(iy, ix), fill, dtype=dtype)
                 for ix in range(n_x)]
                for iy in range(n_y)]

    horiz = lambda iy, ix: chunks_x[ix]  # noqa: E731
    vert = lambda iy, ix: chunks_y[iy]  # noqa: E731

    boundaries = {}
    for key, fill, dtype in [
        ('dist', np.inf, np.float64),
        ('alloc', np.nan, np.float64),
        ('src_row', np.inf, np.float64),
        ('src_col', np.inf, np.float64),
    ]:
        boundaries[key] = {
            'top': _make(horiz, fill, dtype),
            'bottom': _make(horiz, fill, dtype),
            'left': _make(vert, fill, dtype),
            'right': _make(vert, fill, dtype),
        }
    return boundaries


def _compute_seeds_sd(iy, ix, boundaries, elev_bdry,
                      cellsize_x, cellsize_y, chunks_y, chunks_x,
                      n_tile_y, n_tile_x, connectivity):
    """Compute seed arrays for tile (iy, ix) from neighbour boundaries.

    Returns dict with keys 'dist', 'alloc', 'src_row', 'src_col',
    each containing (top, bottom, left, right, tl, tr, bl, br).
    Cardinal seeds are 1-D float64 arrays; corner seeds are float64 scalars.
    """
    tile_h = chunks_y[iy]
    tile_w = chunks_x[ix]
    diag_dist = sqrt(cellsize_x ** 2 + cellsize_y ** 2)

    # Initialize seeds
    seeds = {}
    for key in ('dist', 'alloc', 'src_row', 'src_col'):
        fill = np.inf if key == 'dist' else (np.nan if key == 'alloc'
                                             else np.inf)
        seeds[key] = {
            'top': np.full(tile_w, fill),
            'bottom': np.full(tile_w, fill),
            'left': np.full(tile_h, fill),
            'right': np.full(tile_h, fill),
            'tl': fill, 'tr': fill, 'bl': fill, 'br': fill,
        }

    my_top = elev_bdry['top'][iy][ix]
    my_bottom = elev_bdry['bottom'][iy][ix]
    my_left = elev_bdry['left'][iy][ix]
    my_right = elev_bdry['right'][iy][ix]

    def _edge_seeds(nb_dist, nb_elev, my_elev,
                    nb_alloc, nb_srow, nb_scol,
                    cardinal_dist):
        """Compute min-cost seed per boundary pixel, tracking alloc/src."""
        n = len(nb_dist)

        # Cardinal: pixel c <- pixel c in neighbour
        dz = my_elev - nb_elev
        cost = nb_dist + np.sqrt(cardinal_dist ** 2 + dz ** 2)
        valid = (np.isfinite(nb_dist) & np.isfinite(nb_elev)
                 & np.isfinite(my_elev))

        s_dist = np.where(valid, cost, np.inf)
        s_alloc = np.where(valid, nb_alloc, np.nan)
        s_srow = np.where(valid, nb_srow, np.inf)
        s_scol = np.where(valid, nb_scol, np.inf)

        if connectivity == 8 and n > 1:
            # Diagonal left: pixel c <- pixel c-1 in neighbour
            dz_l = my_elev[1:] - nb_elev[:-1]
            cost_l = nb_dist[:-1] + np.sqrt(diag_dist ** 2 + dz_l ** 2)
            valid_l = (np.isfinite(nb_dist[:-1]) & np.isfinite(nb_elev[:-1])
                       & np.isfinite(my_elev[1:]))
            cost_l = np.where(valid_l, cost_l, np.inf)
            better = cost_l < s_dist[1:]
            s_dist[1:] = np.where(better, cost_l, s_dist[1:])
            s_alloc[1:] = np.where(better, nb_alloc[:-1], s_alloc[1:])
            s_srow[1:] = np.where(better, nb_srow[:-1], s_srow[1:])
            s_scol[1:] = np.where(better, nb_scol[:-1], s_scol[1:])

            # Diagonal right: pixel c <- pixel c+1 in neighbour
            dz_r = my_elev[:-1] - nb_elev[1:]
            cost_r = nb_dist[1:] + np.sqrt(diag_dist ** 2 + dz_r ** 2)
            valid_r = (np.isfinite(nb_dist[1:]) & np.isfinite(nb_elev[1:])
                       & np.isfinite(my_elev[:-1]))
            cost_r = np.where(valid_r, cost_r, np.inf)
            better_r = cost_r < s_dist[:-1]
            s_dist[:-1] = np.where(better_r, cost_r, s_dist[:-1])
            s_alloc[:-1] = np.where(better_r, nb_alloc[1:], s_alloc[:-1])
            s_srow[:-1] = np.where(better_r, nb_srow[1:], s_srow[:-1])
            s_scol[:-1] = np.where(better_r, nb_scol[1:], s_scol[:-1])

        return s_dist, s_alloc, s_srow, s_scol

    def _get_bdry(key, side, iy_, ix_):
        return boundaries[key][side][iy_][ix_]

    # Edge neighbours
    if iy > 0:
        result = _edge_seeds(
            _get_bdry('dist', 'bottom', iy - 1, ix),
            elev_bdry['bottom'][iy - 1][ix],
            my_top,
            _get_bdry('alloc', 'bottom', iy - 1, ix),
            _get_bdry('src_row', 'bottom', iy - 1, ix),
            _get_bdry('src_col', 'bottom', iy - 1, ix),
            cellsize_y,
        )
        for i, key in enumerate(('dist', 'alloc', 'src_row', 'src_col')):
            seeds[key]['top'] = result[i]

    if iy < n_tile_y - 1:
        result = _edge_seeds(
            _get_bdry('dist', 'top', iy + 1, ix),
            elev_bdry['top'][iy + 1][ix],
            my_bottom,
            _get_bdry('alloc', 'top', iy + 1, ix),
            _get_bdry('src_row', 'top', iy + 1, ix),
            _get_bdry('src_col', 'top', iy + 1, ix),
            cellsize_y,
        )
        for i, key in enumerate(('dist', 'alloc', 'src_row', 'src_col')):
            seeds[key]['bottom'] = result[i]

    if ix > 0:
        result = _edge_seeds(
            _get_bdry('dist', 'right', iy, ix - 1),
            elev_bdry['right'][iy][ix - 1],
            my_left,
            _get_bdry('alloc', 'right', iy, ix - 1),
            _get_bdry('src_row', 'right', iy, ix - 1),
            _get_bdry('src_col', 'right', iy, ix - 1),
            cellsize_x,
        )
        for i, key in enumerate(('dist', 'alloc', 'src_row', 'src_col')):
            seeds[key]['left'] = result[i]

    if ix < n_tile_x - 1:
        result = _edge_seeds(
            _get_bdry('dist', 'left', iy, ix + 1),
            elev_bdry['left'][iy][ix + 1],
            my_right,
            _get_bdry('alloc', 'left', iy, ix + 1),
            _get_bdry('src_row', 'left', iy, ix + 1),
            _get_bdry('src_col', 'left', iy, ix + 1),
            cellsize_x,
        )
        for i, key in enumerate(('dist', 'alloc', 'src_row', 'src_col')):
            seeds[key]['right'] = result[i]

    # Diagonal corner seeds (8-connectivity only)
    if connectivity == 8:
        def _corner(nb_d, nb_e, my_e, nb_a, nb_sr, nb_sc):
            nb_d = float(nb_d)
            nb_e = float(nb_e)
            my_e = float(my_e)
            if (np.isfinite(nb_d) and np.isfinite(nb_e)
                    and np.isfinite(my_e)):
                dz = my_e - nb_e
                cost = nb_d + sqrt(diag_dist ** 2 + dz ** 2)
                return cost, float(nb_a), float(nb_sr), float(nb_sc)
            return np.inf, np.nan, np.inf, np.inf

        corners = [
            # (tile_offset_y, tile_offset_x, nb_side, nb_col_idx, my_elev_val,
            #  seed_key)
            (iy - 1, ix - 1, 'bottom', -1, my_top[0], 'tl'),
            (iy - 1, ix + 1, 'bottom', 0, my_top[-1], 'tr'),
            (iy + 1, ix - 1, 'top', -1, my_bottom[0], 'bl'),
            (iy + 1, ix + 1, 'top', 0, my_bottom[-1], 'br'),
        ]
        for niy, nix, nb_side, nb_idx, my_e, skey in corners:
            if 0 <= niy < n_tile_y and 0 <= nix < n_tile_x:
                nb_d = boundaries['dist'][nb_side][niy][nix][nb_idx]
                nb_e = elev_bdry[nb_side][niy][nix][nb_idx]
                nb_a = boundaries['alloc'][nb_side][niy][nix][nb_idx]
                nb_sr = boundaries['src_row'][nb_side][niy][nix][nb_idx]
                nb_sc = boundaries['src_col'][nb_side][niy][nix][nb_idx]
                cd, ca, csr, csc = _corner(nb_d, nb_e, my_e,
                                           nb_a, nb_sr, nb_sc)
                seeds['dist'][skey] = cd
                seeds['alloc'][skey] = ca
                seeds['src_row'][skey] = csr
                seeds['src_col'][skey] = csc

    return seeds


def _can_skip_sd(iy, ix, has_source, boundaries,
                 n_tile_y, n_tile_x, connectivity):
    """True when a tile cannot possibly receive any distance information."""
    if has_source[iy][ix]:
        return False
    bdist = boundaries['dist']
    if iy > 0 and np.any(np.isfinite(bdist['bottom'][iy - 1][ix])):
        return False
    if (iy < n_tile_y - 1
            and np.any(np.isfinite(bdist['top'][iy + 1][ix]))):
        return False
    if ix > 0 and np.any(np.isfinite(bdist['right'][iy][ix - 1])):
        return False
    if (ix < n_tile_x - 1
            and np.any(np.isfinite(bdist['left'][iy][ix + 1]))):
        return False
    if connectivity == 8:
        if (iy > 0 and ix > 0
                and np.isfinite(bdist['bottom'][iy - 1][ix - 1][-1])):
            return False
        if (iy > 0 and ix < n_tile_x - 1
                and np.isfinite(bdist['bottom'][iy - 1][ix + 1][0])):
            return False
        if (iy < n_tile_y - 1 and ix > 0
                and np.isfinite(bdist['top'][iy + 1][ix - 1][-1])):
            return False
        if (iy < n_tile_y - 1 and ix < n_tile_x - 1
                and np.isfinite(bdist['top'][iy + 1][ix + 1][0])):
            return False
    return True


def _run_tile(source_chunk, elev_chunk, seeds, target_values,
              max_distance, dy, dx, dd, row_offset, col_offset):
    """Run Dijkstra on one tile with boundary seeds.

    Returns (dist, alloc, src_row, src_col) as raw float64/int64 arrays.
    src_row/src_col are global indices.
    """
    h, w = source_chunk.shape
    dist, alloc, src_row, src_col = _init_arrays(h, w)

    # Seed source pixels (using global coords for src)
    n_values = len(target_values)
    for r in range(h):
        for c in range(w):
            val = source_chunk[r, c]
            is_target = False
            if n_values == 0:
                if val != 0.0 and np.isfinite(val):
                    is_target = True
            else:
                for k in range(n_values):
                    if val == target_values[k]:
                        is_target = True
                        break
            if is_target and np.isfinite(elev_chunk[r, c]):
                dist[r, c] = 0.0
                alloc[r, c] = val
                src_row[r, c] = r + row_offset
                src_col[r, c] = c + col_offset

    # Seed boundary pixels from neighbour tiles
    sd = seeds['dist']
    sa = seeds['alloc']
    ssr = seeds['src_row']
    ssc = seeds['src_col']

    # Top edge
    for c in range(w):
        if sd['top'][c] < dist[0, c] and np.isfinite(elev_chunk[0, c]):
            dist[0, c] = sd['top'][c]
            alloc[0, c] = sa['top'][c]
            src_row[0, c] = int(ssr['top'][c])
            src_col[0, c] = int(ssc['top'][c])
    # Bottom edge
    for c in range(w):
        if sd['bottom'][c] < dist[h - 1, c] and np.isfinite(
                elev_chunk[h - 1, c]):
            dist[h - 1, c] = sd['bottom'][c]
            alloc[h - 1, c] = sa['bottom'][c]
            src_row[h - 1, c] = int(ssr['bottom'][c])
            src_col[h - 1, c] = int(ssc['bottom'][c])
    # Left edge
    for r in range(h):
        if sd['left'][r] < dist[r, 0] and np.isfinite(elev_chunk[r, 0]):
            dist[r, 0] = sd['left'][r]
            alloc[r, 0] = sa['left'][r]
            src_row[r, 0] = int(ssr['left'][r])
            src_col[r, 0] = int(ssc['left'][r])
    # Right edge
    for r in range(h):
        if sd['right'][r] < dist[r, w - 1] and np.isfinite(
                elev_chunk[r, w - 1]):
            dist[r, w - 1] = sd['right'][r]
            alloc[r, w - 1] = sa['right'][r]
            src_row[r, w - 1] = int(ssr['right'][r])
            src_col[r, w - 1] = int(ssc['right'][r])

    # Corner seeds
    _corners = [
        (0, 0, 'tl'),
        (0, w - 1, 'tr'),
        (h - 1, 0, 'bl'),
        (h - 1, w - 1, 'br'),
    ]
    for cr, cc, skey in _corners:
        sv = sd[skey]
        if sv < dist[cr, cc] and np.isfinite(elev_chunk[cr, cc]):
            dist[cr, cc] = sv
            alloc[cr, cc] = sa[skey]
            src_row[cr, cc] = int(ssr[skey])
            src_col[cr, cc] = int(ssc[skey])

    # Run Dijkstra
    _dijkstra(elev_chunk, h, w, max_distance,
              dy, dx, dd, dist, alloc, src_row, src_col)

    return dist, alloc, src_row, src_col


def _process_tile_sd(iy, ix, source_da, elev_da,
                     boundaries, elev_bdry,
                     cellsize_x, cellsize_y, max_distance, target_values,
                     dy, dx, dd, chunks_y, chunks_x,
                     n_tile_y, n_tile_x, connectivity,
                     cumul_rows, cumul_cols):
    """Run seeded Dijkstra on one tile; update boundaries in-place.

    Returns the maximum absolute boundary distance change.
    """
    source_chunk = source_da.blocks[iy, ix].compute()
    elev_chunk = elev_da.blocks[iy, ix].compute()
    h, w = source_chunk.shape

    row_offset = cumul_rows[iy]
    col_offset = cumul_cols[ix]

    seeds = _compute_seeds_sd(
        iy, ix, boundaries, elev_bdry,
        cellsize_x, cellsize_y, chunks_y, chunks_x,
        n_tile_y, n_tile_x, connectivity,
    )

    dist, alloc_arr, srow, scol = _run_tile(
        source_chunk, elev_chunk, seeds, target_values,
        max_distance, dy, dx, dd, row_offset, col_offset,
    )

    # Extract new boundary strips
    change = 0.0
    for side, sl in [
        ('top', (0, slice(None))),
        ('bottom', (-1, slice(None))),
        ('left', (slice(None), 0)),
        ('right', (slice(None), -1)),
    ]:
        new_d = dist[sl]
        old_d = boundaries['dist'][side][iy][ix]
        with np.errstate(invalid='ignore'):
            diff = np.abs(new_d - old_d)
        diff = np.where(np.isnan(diff), 0.0, diff)
        m = float(np.max(diff))
        if m > change:
            change = m

        boundaries['dist'][side][iy][ix] = new_d.copy()
        boundaries['alloc'][side][iy][ix] = alloc_arr[sl].copy()
        boundaries['src_row'][side][iy][ix] = srow[sl].astype(
            np.float64).copy()
        boundaries['src_col'][side][iy][ix] = scol[sl].astype(
            np.float64).copy()

    return change


def _sd_dask_iterative(source_da, elev_da,
                       cellsize_x, cellsize_y,
                       max_distance, target_values,
                       dy, dx, dd, mode):
    """Iterative boundary-only Dijkstra for arbitrarily large dask arrays."""
    connectivity = len(dy)
    chunks_y = source_da.chunks[0]
    chunks_x = source_da.chunks[1]
    n_tile_y = len(chunks_y)
    n_tile_x = len(chunks_x)

    # Cumulative row/col offsets for global indexing
    cumul_rows = [0]
    for cy in chunks_y[:-1]:
        cumul_rows.append(cumul_rows[-1] + cy)
    cumul_cols = [0]
    for cx in chunks_x[:-1]:
        cumul_cols.append(cumul_cols[-1] + cx)

    # Phase 0: pre-extract elevation boundaries & source flags
    elev_bdry, has_source = _preprocess_tiles_sd(
        source_da, elev_da, chunks_y, chunks_x, target_values,
    )

    # Phase 1: initialise boundaries
    boundaries = _init_boundaries_sd(chunks_y, chunks_x)

    # Phase 2: iterative sweeps
    max_iterations = max(n_tile_y, n_tile_x) + 10
    args = (source_da, elev_da, boundaries, elev_bdry,
            cellsize_x, cellsize_y, max_distance, target_values,
            dy, dx, dd, chunks_y, chunks_x,
            n_tile_y, n_tile_x, connectivity, cumul_rows, cumul_cols)

    for _iteration in range(max_iterations):
        max_change = 0.0

        # Forward sweep
        for iy in range(n_tile_y):
            for ix in range(n_tile_x):
                if _can_skip_sd(iy, ix, has_source, boundaries,
                                n_tile_y, n_tile_x, connectivity):
                    continue
                c = _process_tile_sd(iy, ix, *args)
                if c > max_change:
                    max_change = c

        # Backward sweep
        for iy in reversed(range(n_tile_y)):
            for ix in reversed(range(n_tile_x)):
                if _can_skip_sd(iy, ix, has_source, boundaries,
                                n_tile_y, n_tile_x, connectivity):
                    continue
                c = _process_tile_sd(iy, ix, *args)
                if c > max_change:
                    max_change = c

        if max_change == 0.0:
            break

    # Phase 3: lazy final assembly
    return _assemble_sd(
        source_da, elev_da, boundaries, elev_bdry,
        cellsize_x, cellsize_y, max_distance, target_values,
        dy, dx, dd, chunks_y, chunks_x,
        n_tile_y, n_tile_x, connectivity,
        cumul_rows, cumul_cols, mode,
    )


def _assemble_sd(source_da, elev_da, boundaries, elev_bdry,
                 cellsize_x, cellsize_y, max_distance, target_values,
                 dy, dx, dd, chunks_y, chunks_x,
                 n_tile_y, n_tile_x, connectivity,
                 cumul_rows, cumul_cols, mode):
    """Build lazy dask array by re-running each tile with converged seeds."""

    def _tile_fn(source_block, elev_block, block_info=None):
        if block_info is None or 0 not in block_info:
            return np.full(source_block.shape, np.nan, dtype=np.float32)
        iy, ix = block_info[0]['chunk-location']
        h, w = source_block.shape

        row_offset = cumul_rows[iy]
        col_offset = cumul_cols[ix]

        seeds = _compute_seeds_sd(
            iy, ix, boundaries, elev_bdry,
            cellsize_x, cellsize_y, chunks_y, chunks_x,
            n_tile_y, n_tile_x, connectivity,
        )

        dist, alloc_arr, srow, scol = _run_tile(
            source_block, elev_block, seeds, target_values,
            max_distance, dy, dx, dd, row_offset, col_offset,
        )

        return _extract_output(dist, alloc_arr, srow, scol,
                               cellsize_x, cellsize_y, max_distance, mode)

    return da.map_blocks(
        _tile_fn,
        source_da, elev_da,
        dtype=np.float32,
        meta=np.array((), dtype=np.float32),
    )


# ---------------------------------------------------------------------------
# Dask wrapper
# ---------------------------------------------------------------------------


def _surface_distance_dask(source_da, elev_da, cellsize_x, cellsize_y,
                           max_distance, target_values, dy, dx, dd, mode):
    """Dask path: use map_overlap for bounded, iterative for unbounded."""
    min_cellsize = min(abs(cellsize_x), abs(cellsize_y))
    height, width = source_da.shape

    use_overlap = False
    if np.isfinite(max_distance):
        pad = int(max_distance / min_cellsize) + 1
        chunks_y, chunks_x = source_da.chunks
        if pad < max(chunks_y) and pad < max(chunks_x):
            use_overlap = True

    if use_overlap:
        return _surface_distance_dask_bounded(
            source_da, elev_da, cellsize_x, cellsize_y,
            max_distance, target_values, dy, dx, dd, mode,
        )

    warnings.warn(
        "surface_distance: max_distance is infinite or the implied radius "
        "exceeds chunk dimensions; using iterative tile Dijkstra. "
        "Setting a finite max_distance enables faster single-pass "
        "processing.",
        UserWarning,
        stacklevel=4,
    )
    return _sd_dask_iterative(
        source_da, elev_da, cellsize_x, cellsize_y,
        max_distance, target_values, dy, dx, dd, mode,
    )


# ---------------------------------------------------------------------------
# Dask + CuPy wrapper
# ---------------------------------------------------------------------------


def _surface_distance_dask_cupy(source_da, elev_da,
                                cellsize_x, cellsize_y,
                                max_distance, target_values,
                                dy, dx, dd, mode):
    """Dask+CuPy surface distance.

    Bounded max_distance: map_overlap with per-chunk GPU relaxation.
    Unbounded: convert to dask+numpy, use CPU iterative path.
    """
    import cupy as cp

    min_cellsize = min(abs(cellsize_x), abs(cellsize_y))
    use_overlap = False
    if np.isfinite(max_distance):
        pad = int(max_distance / min_cellsize) + 1
        chunks_y, chunks_x = source_da.chunks
        if pad < max(chunks_y) and pad < max(chunks_x):
            use_overlap = True

    if use_overlap:
        _dy, _dx, _dd = dy, dx, dd
        _mode = mode
        cx, cy = cellsize_x, cellsize_y
        md, tv = max_distance, target_values

        def _chunk_func(source_block, elev_block):
            return _surface_distance_cupy(
                source_block, elev_block, cx, cy,
                md, tv, _dy, _dx, _dd, _mode,
            )

        return da.map_overlap(
            _chunk_func,
            source_da, elev_da,
            depth=(pad, pad),
            boundary=np.nan,
            dtype=np.float32,
            meta=cp.array((), dtype=cp.float32),
        )

    # Unbounded: convert to dask+numpy, use CPU path
    source_np = source_da.map_blocks(
        lambda b: b.get(), dtype=source_da.dtype,
        meta=np.array((), dtype=source_da.dtype),
    )
    elev_np = elev_da.map_blocks(
        lambda b: b.get(), dtype=elev_da.dtype,
        meta=np.array((), dtype=elev_da.dtype),
    )
    result = _surface_distance_dask(
        source_np, elev_np, cellsize_x, cellsize_y,
        max_distance, target_values, dy, dx, dd, mode,
    )
    return result.map_blocks(
        cp.asarray, dtype=result.dtype,
        meta=cp.array((), dtype=result.dtype),
    )


# ---------------------------------------------------------------------------
# Core dispatcher
# ---------------------------------------------------------------------------


def _compute(raster, elevation, x, y, target_values, max_distance,
             connectivity, method, mode):
    """Core dispatcher for surface_distance / allocation / direction."""
    _validate_raster(raster, func_name='surface_distance', name='raster')
    _validate_raster(elevation, func_name='surface_distance',
                     name='elevation')
    if raster.shape != elevation.shape:
        raise ValueError("raster and elevation must have the same shape")
    if raster.dims != (y, x):
        raise ValueError(
            f"raster.dims should be ({y!r}, {x!r}), got {raster.dims}"
        )
    if connectivity not in (4, 8):
        raise ValueError("connectivity must be 4 or 8")
    if method not in ('planar', 'geodesic'):
        raise ValueError("method must be 'planar' or 'geodesic'")

    cellsize_x, cellsize_y = get_dataarray_resolution(raster)
    cellsize_x = abs(float(cellsize_x))
    cellsize_y = abs(float(cellsize_y))

    target_values = np.asarray(target_values, dtype=np.float64)
    max_distance_f = float(max_distance)

    # Build neighbour offsets
    if connectivity == 8:
        dy_arr = np.array([-1, -1, -1, 0, 0, 1, 1, 1], dtype=np.int64)
        dx_arr = np.array([-1, 0, 1, -1, 1, -1, 0, 1], dtype=np.int64)
        dd_arr = np.array([
            sqrt(cellsize_y ** 2 + cellsize_x ** 2),
            cellsize_y,
            sqrt(cellsize_y ** 2 + cellsize_x ** 2),
            cellsize_x,
            cellsize_x,
            sqrt(cellsize_y ** 2 + cellsize_x ** 2),
            cellsize_y,
            sqrt(cellsize_y ** 2 + cellsize_x ** 2),
        ], dtype=np.float64)
    else:
        dy_arr = np.array([0, -1, 1, 0], dtype=np.int64)
        dx_arr = np.array([-1, 0, 0, 1], dtype=np.int64)
        dd_arr = np.array(
            [cellsize_x, cellsize_y, cellsize_y, cellsize_x],
            dtype=np.float64)

    # Geodesic dd_grid precomputation (numpy only for now)
    use_geodesic = (method == 'geodesic')
    dd_grid = None
    if use_geodesic:
        from xrspatial.utils import _extract_latlon_coords
        lat_2d, lon_2d = _extract_latlon_coords(raster)
        dd_grid = _precompute_dd_grid(lat_2d, lon_2d, dy_arr, dx_arr)

    source_data = raster.data
    elev_data = elevation.data

    _is_dask = da is not None and isinstance(source_data, da.Array)
    _is_cupy = (
        not _is_dask
        and has_cuda_and_cupy()
        and is_cupy_array(source_data)
    )
    _is_dask_cupy_flag = _is_dask and has_cuda_and_cupy() and is_dask_cupy(
        raster)

    if _is_dask:
        # Ensure chunks match
        if isinstance(elev_data, da.Array):
            elev_data = elev_data.rechunk(source_data.chunks)
        else:
            elev_data = da.from_array(elev_data, chunks=source_data.chunks)

    if _is_cupy:
        if use_geodesic:
            raise NotImplementedError(
                "geodesic mode is not yet supported for CuPy arrays")
        result_data = _surface_distance_cupy(
            source_data, elev_data, cellsize_x, cellsize_y,
            max_distance_f, target_values, dy_arr, dx_arr, dd_arr, mode,
        )
    elif _is_dask_cupy_flag:
        if use_geodesic:
            raise NotImplementedError(
                "geodesic mode is not yet supported for Dask+CuPy arrays")
        result_data = _surface_distance_dask_cupy(
            source_data, elev_data, cellsize_x, cellsize_y,
            max_distance_f, target_values, dy_arr, dx_arr, dd_arr, mode,
        )
    elif isinstance(source_data, np.ndarray):
        if isinstance(elev_data, np.ndarray):
            result_data = _surface_distance_numpy(
                source_data, elev_data, cellsize_x, cellsize_y,
                max_distance_f, target_values, dy_arr, dx_arr, dd_arr,
                dd_grid, use_geodesic, mode,
            )
        else:
            elev_np = np.asarray(elev_data)
            result_data = _surface_distance_numpy(
                source_data, elev_np, cellsize_x, cellsize_y,
                max_distance_f, target_values, dy_arr, dx_arr, dd_arr,
                dd_grid, use_geodesic, mode,
            )
    elif _is_dask:
        if use_geodesic:
            raise NotImplementedError(
                "geodesic mode is not yet supported for Dask arrays")
        result_data = _surface_distance_dask(
            source_data, elev_data, cellsize_x, cellsize_y,
            max_distance_f, target_values, dy_arr, dx_arr, dd_arr, mode,
        )
    else:
        raise TypeError(f"Unsupported array type: {type(source_data)}")

    return result_data


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


@supports_dataset
def surface_distance(
    raster: xr.DataArray,
    elevation: xr.DataArray,
    x: str = "x",
    y: str = "y",
    target_values: list = [],
    max_distance: float = np.inf,
    connectivity: int = 8,
    method: str = 'planar',
) -> xr.DataArray:
    """Compute surface distance along terrain to nearest target pixel.

    For every pixel, computes the minimum accumulated distance along
    the 3D terrain surface to reach the nearest target pixel.
    Edge cost accounts for both horizontal distance and elevation
    change: ``sqrt(horizontal_dist^2 + dz^2)``.

    Parameters
    ----------
    raster : xr.DataArray or xr.Dataset
        2-D source raster.  Target pixels are identified by non-zero
        finite values (or values in *target_values*).
    elevation : xr.DataArray
        2-D elevation surface.  Must have the same shape as *raster*.
        NaN marks impassable barriers.
    x : str, default='x'
        Name of the x coordinate.
    y : str, default='y'
        Name of the y coordinate.
    target_values : list, optional
        Specific pixel values in *raster* to treat as sources.
        If empty, all non-zero finite pixels are sources.
    max_distance : float, default=np.inf
        Maximum surface distance.  Pixels beyond this budget are NaN.
        A finite value enables efficient Dask parallelisation.
    connectivity : int, default=8
        Pixel connectivity: 4 (cardinal) or 8 (cardinal + diagonal).
    method : str, default='planar'
        ``'planar'`` uses cell sizes in map units.
        ``'geodesic'`` computes great-circle horizontal distances
        from lat/lon coordinates (elevation in meters).

    Returns
    -------
    xr.DataArray or xr.Dataset
        2-D array of surface distance values (float32).
        Source pixels have distance 0.  Unreachable pixels are NaN.
    """
    result_data = _compute(
        raster, elevation, x, y, target_values, max_distance,
        connectivity, method, DISTANCE,
    )
    return xr.DataArray(
        result_data,
        coords=raster.coords,
        dims=raster.dims,
        attrs=raster.attrs,
    )


@supports_dataset
def surface_allocation(
    raster: xr.DataArray,
    elevation: xr.DataArray,
    x: str = "x",
    y: str = "y",
    target_values: list = [],
    max_distance: float = np.inf,
    connectivity: int = 8,
    method: str = 'planar',
) -> xr.DataArray:
    """Compute nearest-target allocation along terrain surface.

    For each pixel, returns the value of the nearest target pixel by
    surface distance through the elevation model.

    Parameters
    ----------
    raster : xr.DataArray or xr.Dataset
        2-D source raster with target pixels.
    elevation : xr.DataArray
        2-D elevation surface (same shape as *raster*).
    x, y, target_values, max_distance, connectivity, method
        See :func:`surface_distance`.

    Returns
    -------
    xr.DataArray or xr.Dataset
        2-D array of allocation values (float32).
    """
    result_data = _compute(
        raster, elevation, x, y, target_values, max_distance,
        connectivity, method, ALLOCATION,
    )
    return xr.DataArray(
        result_data,
        coords=raster.coords,
        dims=raster.dims,
        attrs=raster.attrs,
    )


@supports_dataset
def surface_direction(
    raster: xr.DataArray,
    elevation: xr.DataArray,
    x: str = "x",
    y: str = "y",
    target_values: list = [],
    max_distance: float = np.inf,
    connectivity: int = 8,
    method: str = 'planar',
) -> xr.DataArray:
    """Compute compass direction to nearest target along terrain surface.

    For each pixel, returns the compass direction (in degrees) to the
    nearest target pixel by surface distance.  0 = source pixel,
    90 = east, 180 = south, 270 = west, 360 = north.

    Parameters
    ----------
    raster : xr.DataArray or xr.Dataset
        2-D source raster with target pixels.
    elevation : xr.DataArray
        2-D elevation surface (same shape as *raster*).
    x, y, target_values, max_distance, connectivity, method
        See :func:`surface_distance`.

    Returns
    -------
    xr.DataArray or xr.Dataset
        2-D array of direction values (float32, degrees).
    """
    result_data = _compute(
        raster, elevation, x, y, target_values, max_distance,
        connectivity, method, DIRECTION,
    )
    return xr.DataArray(
        result_data,
        coords=raster.coords,
        dims=raster.dims,
        attrs=raster.attrs,
    )