__init__.py

"""Out-of-core CRS reprojection and multi-raster merge.

Public API
----------
reproject(raster, target_crs, ...)
    Reproject a DataArray to a new coordinate reference system.
merge(rasters, ...)
    Merge multiple DataArrays into a single mosaic.
"""
from __future__ import annotations

import math

import numpy as np
import xarray as xr

from xrspatial.utils import _dask_task_name_kwargs, _validate_raster

from ._crs_utils import _detect_band_nodata, _detect_nodata, _detect_source_crs, _resolve_crs
from ._grid import (_MAX_OUTPUT_PIXELS, _chunk_bounds, _compute_chunk_layout, _compute_output_grid,
                    _make_output_coords, _validate_grid_params)
from ._interpolate import _resample_cupy_native, _resample_numpy, _validate_resampling
from ._itrf import itrf_transform
from ._itrf import list_frames as itrf_frames
from ._merge import _merge_arrays_numpy, _validate_strategy
from ._transform import ApproximateTransform
from ._vertical import (depth_to_ellipsoidal, ellipsoidal_to_depth, ellipsoidal_to_orthometric,
                        geoid_height, geoid_height_raster, orthometric_to_ellipsoidal)

__all__ = [
    'reproject', 'merge',
    'geoid_height', 'geoid_height_raster',
    'ellipsoidal_to_orthometric', 'orthometric_to_ellipsoidal',
    'depth_to_ellipsoidal', 'ellipsoidal_to_depth',
    'itrf_transform', 'itrf_frames',
]


# ---------------------------------------------------------------------------
# Source geometry helpers
# ---------------------------------------------------------------------------

_Y_NAMES = {'y', 'lat', 'latitude', 'Y', 'Lat', 'Latitude'}
_X_NAMES = {'x', 'lon', 'longitude', 'X', 'Lon', 'Longitude'}

# Output byte budget above which merge() auto-promotes an in-memory mosaic
# to the lazy dask path instead of allocating the whole array.
_MERGE_OOM_THRESHOLD = 512 * 1024 * 1024  # 512 MB

# Byte budget above which reproject() auto-promotes an in-memory raster
# (numpy or cupy) to the chunked dask path. Compared against the input
# array size and one float64 output array independently, not against
# total memory: the eager numpy path holds ~7 output-sized float64
# temporaries (coordinate grids, pixel-index grids, result), so
# eager-path peak RSS can reach ~7x this budget. That multiplier is why a
# small input upsampled to a large output exhausted memory long before
# the _MAX_OUTPUT_PIXELS guard tripped (#3267). The same budget gates the
# cupy promotion added in #3281.
_REPROJECT_OOM_THRESHOLD = 512 * 1024 * 1024  # 512 MB

# Map friendly vertical datum tokens to EPSG codes so attrs['vertical_crs']
# from reproject output matches the convention used by xrspatial.geotiff,
# which also writes EPSG ints to attrs['vertical_crs'].
_VERTICAL_DATUM_EPSG = {
    'EGM96': 5773,        # EGM96 height
    'EGM2008': 3855,      # EGM2008 height
    'ellipsoidal': 4979,  # WGS 84 (3D, ellipsoidal height)
}

# Sentinel marking the deprecated ``src_vertical_crs`` / ``tgt_vertical_crs``
# kwargs as "not passed". Distinct from None so we can tell an explicit
# ``src_vertical_crs=None`` apart from the default and only warn when the
# caller actually used the old name.
_DEPRECATED = object()


def _resolve_deprecated_vertical_kwarg(old_name, old_val, new_name, new_val):
    """Map a deprecated vertical-CRS kwarg onto its renamed replacement.

    Emits a ``DeprecationWarning`` when the old name is used and rejects
    passing both the old and new spellings at once.
    """
    if old_val is _DEPRECATED:
        return new_val
    import warnings
    warnings.warn(
        f"reproject(): {old_name!r} is deprecated, use {new_name!r} instead.",
        DeprecationWarning,
        stacklevel=3,
    )
    if new_val is not None:
        raise TypeError(
            f"reproject(): pass either {new_name!r} or the deprecated "
            f"{old_name!r}, not both."
        )
    return old_val


def _find_spatial_dims(raster):
    """Find the y and x dimension names, handling multi-band rasters.

    Returns (ydim, xdim).  Checks dim names first, falls back to
    assuming the last two non-band dims are spatial.
    """
    dims = raster.dims
    ydim = xdim = None
    for d in dims:
        if d in _Y_NAMES:
            ydim = d
        elif d in _X_NAMES:
            xdim = d
    if ydim is not None and xdim is not None:
        return ydim, xdim
    # Fallback: last two dims
    return dims[-2], dims[-1]


# Default tolerance for the regular-spacing check. Coordinates loaded from
# real GeoTIFFs can drift a few ULPs from perfectly uniform after pixel-to-
# world transforms, so 1e-6 (relative) is loose enough to accept those while
# still catching the single-pixel perturbation case in #2184.
_REGULAR_COORD_RTOL = 1e-6


def _validate_regular_axis(coords, axis_name, func_name, rtol=_REGULAR_COORD_RTOL):
    """Validate that 1-D coordinate array is strictly monotonic and regular.

    Pixel-resolution math in `_source_bounds` and the chunk workers assumes
    a uniform grid. Without this check, irregular or non-monotonic coords
    silently produce wrong georeferencing (see #2184).

    Parameters
    ----------
    coords : array-like
        1-D coordinate values along one axis.
    axis_name : str
        Name of the axis ("x" or "y") for the error message.
    func_name : str
        Calling function name for the error prefix.
    rtol : float
        Relative tolerance for spacing regularity.

    Raises
    ------
    ValueError
        If coords contain non-finite values, are not strictly monotonic,
        or have spacing that varies by more than ``rtol`` relative to the
        median step.
    """
    arr = np.asarray(coords)
    if arr.ndim != 1:
        raise ValueError(
            f"{func_name}(): coordinate '{axis_name}' must be 1-D, "
            f"got shape {arr.shape}."
        )
    if arr.size < 2:
        # A single-pixel raster has no spacing to validate; the caller
        # will fall back to res=1.0 in _source_bounds, which is fine.
        return
    if not np.all(np.isfinite(arr)):
        raise ValueError(
            f"{func_name}(): coordinate '{axis_name}' contains non-finite "
            f"values (NaN or inf)."
        )
    # np.asarray skips the copy when arr is already float64; np.diff promotes
    # ints to int64, which is fine but we want float steps for the median /
    # tolerance math below.
    diffs = np.diff(np.asarray(arr, dtype=np.float64))
    # Strict monotonicity: every step has the same sign and is non-zero.
    # `diffs > 0` AND `diffs < 0` are both False for zero steps (repeated
    # coords), so the combined check rejects them. Do NOT replace this with
    # a sign-only test like `np.all(np.sign(diffs) == np.sign(diffs[0]))` --
    # that variant accepts zero steps and lets a repeated coord through.
    if not (np.all(diffs > 0) or np.all(diffs < 0)):
        raise ValueError(
            f"{func_name}(): coordinate '{axis_name}' must be strictly "
            f"monotonic (all ascending or all descending). The reproject "
            f"pipeline assumes a uniformly-spaced grid; see #2184."
        )
    median_step = float(np.median(diffs))
    abs_med = abs(median_step)
    deviation = np.abs(diffs - median_step)
    worst = float(np.max(deviation))
    if worst > rtol * abs_med:
        # Report the index of the worst step in the original coords so the
        # caller can locate the offending sample without re-running diff.
        worst_idx = int(np.argmax(deviation))
        raise ValueError(
            f"{func_name}(): coordinate '{axis_name}' is not regularly "
            f"spaced. Median step is {median_step!r}; worst deviation is "
            f"{worst!r} at index {worst_idx} (between {axis_name}[{worst_idx}]"
            f"={float(arr[worst_idx])!r} and {axis_name}[{worst_idx + 1}]"
            f"={float(arr[worst_idx + 1])!r}). The reproject pipeline "
            f"assumes a uniformly-spaced grid; see #2184."
        )


def _validate_source_coords(raster, func_name):
    """Validate both spatial axes of a raster before any reproject work."""
    ydim, xdim = _find_spatial_dims(raster)
    _validate_regular_axis(raster.coords[ydim].values, 'y', func_name)
    _validate_regular_axis(raster.coords[xdim].values, 'x', func_name)


def _source_bounds(raster):
    """Extract (left, bottom, right, top) from a DataArray's coordinates."""
    ydim, xdim = _find_spatial_dims(raster)
    y = raster.coords[ydim].values
    x = raster.coords[xdim].values
    # Compute pixel-edge bounds from pixel-center coords
    if len(y) > 1:
        res_y = abs(float(y[1] - y[0]))
    else:
        res_y = 1.0
    if len(x) > 1:
        res_x = abs(float(x[1] - x[0]))
    else:
        res_x = 1.0
    x_min, x_max = float(np.min(x)), float(np.max(x))
    y_min, y_max = float(np.min(y)), float(np.max(y))
    left = x_min - res_x / 2
    right = x_max + res_x / 2
    bottom = y_min - res_y / 2
    top = y_max + res_y / 2
    return (left, bottom, right, top)


def _is_y_descending(raster):
    """Check if Y axis goes from top (large) to bottom (small)."""
    ydim, _ = _find_spatial_dims(raster)
    y = raster.coords[ydim].values
    if len(y) < 2:
        return True
    return float(y[0]) > float(y[-1])


def _is_x_descending(raster):
    """Check if X axis goes from right (large) to left (small).

    Mirrors :func:`_is_y_descending` for the horizontal axis. The default
    convention for a single-column raster is ascending x (matching
    :func:`_make_output_coords` which always emits ascending x).
    """
    _, xdim = _find_spatial_dims(raster)
    x = raster.coords[xdim].values
    if len(x) < 2:
        return False
    return float(x[0]) > float(x[-1])


# ---------------------------------------------------------------------------
# Per-chunk coordinate transform
# ---------------------------------------------------------------------------

def _transform_coords(transformer, chunk_bounds, chunk_shape,
                      transform_precision, src_crs=None, tgt_crs=None):
    """Compute source CRS coordinates for every output pixel.

    When *transform_precision* is 0, every pixel is transformed through
    pyproj exactly (same strategy as GDAL/rasterio).  Otherwise an
    approximate bilinear control-grid interpolation is used.

    For common CRS pairs (WGS84/NAD83 <-> UTM, WGS84 <-> Web Mercator),
    a Numba JIT fast path bypasses pyproj entirely for ~30x speedup.

    Returns
    -------
    src_y, src_x : ndarray (height, width)
    """
    # Try Numba fast path for common projections.
    # transform_precision == 0 is the documented escape hatch for exact
    # per-pixel pyproj transforms, so skip the approximate fast path then.
    if (transform_precision != 0
            and src_crs is not None and tgt_crs is not None):
        try:
            from ._projections import try_numba_transform
            result = try_numba_transform(
                src_crs, tgt_crs, chunk_bounds, chunk_shape,
            )
            if result is not None:
                return result
        except (ImportError, ModuleNotFoundError):
            pass  # fall through to pyproj

    height, width = chunk_shape
    left, bottom, right, top = chunk_bounds
    res_x = (right - left) / width
    res_y = (top - bottom) / height

    if transform_precision == 0:
        # Exact per-pixel transform via pyproj bulk API.
        # Process in row strips to keep memory bounded and improve
        # cache locality for large rasters.
        out_x_1d = left + (np.arange(width, dtype=np.float64) + 0.5) * res_x
        src_x_out = np.empty((height, width), dtype=np.float64)
        src_y_out = np.empty((height, width), dtype=np.float64)
        strip = 256
        for r0 in range(0, height, strip):
            r1 = min(r0 + strip, height)
            n_rows = r1 - r0
            out_y_strip = top - (np.arange(r0, r1, dtype=np.float64) + 0.5) * res_y
            # Broadcast to (n_rows, width) without allocating a full copy
            sx, sy = transformer.transform(
                np.tile(out_x_1d, n_rows),
                np.repeat(out_y_strip, width),
            )
            src_x_out[r0:r1] = np.asarray(sx, dtype=np.float64).reshape(n_rows, width)
            src_y_out[r0:r1] = np.asarray(sy, dtype=np.float64).reshape(n_rows, width)
        return src_y_out, src_x_out

    # Approximate: bilinear interpolation on a coarse control grid.
    approx = ApproximateTransform(
        transformer, chunk_bounds, chunk_shape,
        precision=transform_precision,
    )
    row_grid = np.arange(height, dtype=np.float64)[:, np.newaxis]
    col_grid = np.arange(width, dtype=np.float64)[np.newaxis, :]
    row_grid = np.broadcast_to(row_grid, (height, width))
    col_grid = np.broadcast_to(col_grid, (height, width))
    return approx(row_grid, col_grid)


# ---------------------------------------------------------------------------
# Per-chunk worker functions
# ---------------------------------------------------------------------------

def _reproject_chunk_numpy(
    source_data, source_bounds_tuple, source_shape, source_y_desc,
    src_wkt, tgt_wkt,
    chunk_bounds_tuple, chunk_shape,
    resampling, nodata, transform_precision,
    source_x_desc=False,
    band_nodata=None,
):
    """Reproject a single output chunk (numpy backend).

    Called inside ``dask.delayed`` for the dask path, or directly for numpy.
    CRS objects are passed as WKT strings for pickle safety.

    ``source_x_desc`` mirrors ``source_y_desc`` for the horizontal axis:
    when True, source column 0 is at the maximum x and column ``src_w-1``
    is at the minimum x. Defaults to False so older callers keep working.
    """
    from ._crs_utils import _crs_from_wkt

    src_crs = _crs_from_wkt(src_wkt)
    tgt_crs = _crs_from_wkt(tgt_wkt)

    # Try Numba fast path first (avoids creating pyproj Transformer).
    # transform_precision == 0 forces the exact pyproj path, so skip Numba.
    numba_result = None
    if transform_precision != 0:
        try:
            from ._projections import try_numba_transform
            numba_result = try_numba_transform(
                src_crs, tgt_crs, chunk_bounds_tuple, chunk_shape,
            )
        except (ImportError, ModuleNotFoundError):
            pass

    if numba_result is not None:
        src_y, src_x = numba_result
    else:
        # Fallback: create pyproj Transformer (expensive)
        from ._crs_utils import _require_pyproj
        pyproj = _require_pyproj()
        transformer = pyproj.Transformer.from_crs(
            tgt_crs, src_crs, always_xy=True
        )
        # Pass src_crs/tgt_crs as None: the numba fast path was already
        # tried above and returned None, and _transform_coords gates its
        # own try_numba_transform retry on both CRSes being non-None.
        # Re-trying would repeat the CRS param parsing and chunk-sized
        # coordinate allocations for nothing (#3106).
        src_y, src_x = _transform_coords(
            transformer, chunk_bounds_tuple, chunk_shape, transform_precision,
        )

    # Convert source CRS coordinates to source pixel coordinates
    src_left, src_bottom, src_right, src_top = source_bounds_tuple
    src_h, src_w = source_shape
    src_res_x = (src_right - src_left) / src_w
    src_res_y = (src_top - src_bottom) / src_h

    if source_x_desc:
        src_col_px = (src_right - src_x) / src_res_x - 0.5
    else:
        src_col_px = (src_x - src_left) / src_res_x - 0.5
    if source_y_desc:
        src_row_px = (src_top - src_y) / src_res_y - 0.5
    else:
        src_row_px = (src_y - src_bottom) / src_res_y - 0.5

    # Determine source window needed
    r_min = np.nanmin(src_row_px)
    r_max = np.nanmax(src_row_px)
    c_min = np.nanmin(src_col_px)
    c_max = np.nanmax(src_col_px)

    # 3-D source: empty-chunk returns must carry the band axis or the
    # dask map_blocks template (which is 3-D for 3-D sources) sees a
    # shape mismatch (#2027).
    if source_data.ndim == 3:
        _empty_shape = (*chunk_shape, source_data.shape[2])
    else:
        _empty_shape = chunk_shape

    # Empty-chunk fills must match the dtype the data path returns and the
    # dask template advertises (#3096): integer sources round-trip back to
    # their dtype, floats stay float64. Without this, a single no-overlap
    # chunk promoted the whole assembled dask output to float64. The
    # resolved nodata is guaranteed representable for integer rasters
    # (#2185/#2572).
    if np.issubdtype(source_data.dtype, np.integer):
        _empty_dtype = source_data.dtype
    else:
        _empty_dtype = np.float64

    if not np.isfinite(r_min) or not np.isfinite(r_max):
        return np.full(_empty_shape, nodata, dtype=_empty_dtype)
    if not np.isfinite(c_min) or not np.isfinite(c_max):
        return np.full(_empty_shape, nodata, dtype=_empty_dtype)

    r_min = int(np.floor(r_min)) - 2
    r_max = int(np.ceil(r_max)) + 3
    c_min = int(np.floor(c_min)) - 2
    c_max = int(np.ceil(c_max)) + 3

    # Check overlap
    if r_min >= src_h or r_max <= 0 or c_min >= src_w or c_max <= 0:
        return np.full(_empty_shape, nodata, dtype=_empty_dtype)

    # Clip to source bounds
    r_min_clip = max(0, r_min)
    r_max_clip = min(src_h, r_max)
    c_min_clip = max(0, c_min)
    c_max_clip = min(src_w, c_max)

    # Guard: cap source window to prevent OOM if projection maps a small
    # output chunk to a huge source area (e.g. polar stereographic edges).
    _MAX_WINDOW_PIXELS = 64 * 1024 * 1024  # 64 Mpix (~512 MB for float64)
    win_pixels = (r_max_clip - r_min_clip) * (c_max_clip - c_min_clip)
    if win_pixels > _MAX_WINDOW_PIXELS:
        return np.full(_empty_shape, nodata, dtype=_empty_dtype)

    # Extract source window
    window = source_data[r_min_clip:r_max_clip, c_min_clip:c_max_clip]
    if hasattr(window, 'compute'):
        window = window.compute()
    window = np.asarray(window)
    orig_dtype = window.dtype

    # Adjust coordinates relative to window
    local_row = src_row_px - r_min_clip
    local_col = src_col_px - c_min_clip

    # Multi-band: reproject each band separately, share coordinates
    if window.ndim == 3:
        n_bands = window.shape[2]
        bands = []
        for b in range(n_bands):
            band_data = window[:, :, b].astype(np.float64)
            # Mask this band with its own source sentinel when the raster
            # declares per-band nodata; otherwise fall back to the single
            # resolved sentinel (#2647).
            src_nd = band_nodata[b] if band_nodata is not None else nodata
            if not np.isnan(src_nd):
                band_data[band_data == src_nd] = np.nan
            band_result = _resample_numpy(band_data, local_row, local_col,
                                          resampling=resampling, nodata=nodata)
            if np.issubdtype(orig_dtype, np.integer):
                info = np.iinfo(orig_dtype)
                band_result = np.clip(np.round(band_result), info.min, info.max).astype(orig_dtype)
            bands.append(band_result)
        return np.stack(bands, axis=-1)

    # Single-band path
    window = window.astype(np.float64)

    # Convert sentinel nodata to NaN so numba kernels can detect it
    if not np.isnan(nodata):
        window[window == nodata] = np.nan

    result = _resample_numpy(window, local_row, local_col,
                             resampling=resampling, nodata=nodata)

    # Clamp and cast back for integer source dtypes
    if np.issubdtype(orig_dtype, np.integer):
        info = np.iinfo(orig_dtype)
        result = np.clip(np.round(result), info.min, info.max).astype(orig_dtype)

    return result


def _reproject_chunk_cupy(
    source_data, source_bounds_tuple, source_shape, source_y_desc,
    src_wkt, tgt_wkt,
    chunk_bounds_tuple, chunk_shape,
    resampling, nodata, transform_precision,
    source_x_desc=False,
    band_nodata=None,
):
    """CuPy variant of ``_reproject_chunk_numpy``.

    ``source_x_desc`` carries the horizontal direction flag (same meaning
    as in :func:`_reproject_chunk_numpy`).
    """
    import cupy as cp

    from ._crs_utils import _crs_from_wkt

    src_crs = _crs_from_wkt(src_wkt)
    tgt_crs = _crs_from_wkt(tgt_wkt)

    # 3-D source: empty-chunk returns must carry the band axis to match
    # the dask+cupy map_blocks template (#2027).
    if source_data.ndim == 3:
        _empty_shape = (*chunk_shape, source_data.shape[2])
    else:
        _empty_shape = chunk_shape

    # Empty-chunk fills must match the dtype the data path returns (#3096);
    # see the matching block in _reproject_chunk_numpy.
    if np.issubdtype(source_data.dtype, np.integer):
        _empty_dtype = source_data.dtype
    else:
        _empty_dtype = np.float64

    # Try CUDA transform first (keeps coordinates on-device).
    # transform_precision == 0 forces the exact pyproj path, so skip CUDA.
    cuda_result = None
    if (transform_precision != 0
            and src_crs is not None and tgt_crs is not None):
        try:
            from ._projections_cuda import try_cuda_transform
            cuda_result = try_cuda_transform(
                src_crs, tgt_crs, chunk_bounds_tuple, chunk_shape,
            )
        except (ImportError, ModuleNotFoundError):
            pass

    if cuda_result is not None:
        src_y, src_x = cuda_result  # cupy arrays
        src_left, src_bottom, src_right, src_top = source_bounds_tuple
        src_h, src_w = source_shape
        src_res_x = (src_right - src_left) / src_w
        src_res_y = (src_top - src_bottom) / src_h
        # Pixel coordinate math stays on GPU via cupy operators
        if source_x_desc:
            src_col_px = (src_right - src_x) / src_res_x - 0.5
        else:
            src_col_px = (src_x - src_left) / src_res_x - 0.5
        if source_y_desc:
            src_row_px = (src_top - src_y) / src_res_y - 0.5
        else:
            src_row_px = (src_y - src_bottom) / src_res_y - 0.5
        # Need min/max on CPU for window selection.
        # Stack the four reductions and pull them across in one device-to-host
        # transfer to avoid four separate synchronous syncs.
        mins_maxes = cp.stack([
            cp.nanmin(src_row_px), cp.nanmax(src_row_px),
            cp.nanmin(src_col_px), cp.nanmax(src_col_px),
        ])
        r_min_val, r_max_val, c_min_val, c_max_val = (
            float(v) for v in mins_maxes.get()
        )
        if not (np.isfinite(r_min_val) and np.isfinite(r_max_val)
                and np.isfinite(c_min_val) and np.isfinite(c_max_val)):
            return cp.full(_empty_shape, nodata, dtype=_empty_dtype)
        r_min = int(np.floor(r_min_val)) - 2
        r_max = int(np.ceil(r_max_val)) + 3
        c_min = int(np.floor(c_min_val)) - 2
        c_max = int(np.ceil(c_max_val)) + 3
        # Coordinates stay as CuPy arrays for native CUDA resampling
    else:
        # CPU fallback (Numba JIT or pyproj)
        from ._crs_utils import _require_pyproj
        pyproj = _require_pyproj()
        transformer = pyproj.Transformer.from_crs(
            tgt_crs, src_crs, always_xy=True
        )
        src_y, src_x = _transform_coords(
            transformer, chunk_bounds_tuple, chunk_shape, transform_precision,
            src_crs=src_crs, tgt_crs=tgt_crs,
        )

        src_left, src_bottom, src_right, src_top = source_bounds_tuple
        src_h, src_w = source_shape
        src_res_x = (src_right - src_left) / src_w
        src_res_y = (src_top - src_bottom) / src_h

        if source_x_desc:
            src_col_px = (src_right - src_x) / src_res_x - 0.5
        else:
            src_col_px = (src_x - src_left) / src_res_x - 0.5
        if source_y_desc:
            src_row_px = (src_top - src_y) / src_res_y - 0.5
        else:
            src_row_px = (src_y - src_bottom) / src_res_y - 0.5

        r_min = np.nanmin(src_row_px)
        r_max = np.nanmax(src_row_px)
        c_min = np.nanmin(src_col_px)
        c_max = np.nanmax(src_col_px)
        if not np.isfinite(r_min) or not np.isfinite(r_max):
            return cp.full(_empty_shape, nodata, dtype=_empty_dtype)
        if not np.isfinite(c_min) or not np.isfinite(c_max):
            return cp.full(_empty_shape, nodata, dtype=_empty_dtype)
        r_min = int(np.floor(r_min)) - 2
        r_max = int(np.ceil(r_max)) + 3
        c_min = int(np.floor(c_min)) - 2
        c_max = int(np.ceil(c_max)) + 3

    if r_min >= src_h or r_max <= 0 or c_min >= src_w or c_max <= 0:
        return cp.full(_empty_shape, nodata, dtype=_empty_dtype)

    r_min_clip = max(0, r_min)
    r_max_clip = min(src_h, r_max)
    c_min_clip = max(0, c_min)
    c_max_clip = min(src_w, c_max)

    # Guard: cap source window to prevent GPU OOM if projection maps a
    # small output chunk to a huge source area (matches numpy path).
    _MAX_WINDOW_PIXELS = 64 * 1024 * 1024  # 64 Mpix (~512 MB for float64)
    win_pixels = (r_max_clip - r_min_clip) * (c_max_clip - c_min_clip)
    if win_pixels > _MAX_WINDOW_PIXELS:
        return cp.full(_empty_shape, nodata, dtype=_empty_dtype)

    window = source_data[r_min_clip:r_max_clip, c_min_clip:c_max_clip]
    if hasattr(window, 'compute'):
        window = window.compute()
    if not isinstance(window, cp.ndarray):
        window = cp.asarray(window)
    orig_dtype = window.dtype

    # Adjust coordinates relative to window (stays on GPU if CuPy)
    local_row = src_row_px - r_min_clip
    local_col = src_col_px - c_min_clip

    # Multi-band: reproject each band separately, share coordinates.
    # Matches the 3-D branch in _reproject_chunk_numpy so 3-D inputs work
    # on cupy and dask+cupy backends instead of crashing with a CUDA
    # signature mismatch (#2027).
    if window.ndim == 3:
        n_bands = window.shape[2]
        # The coordinate arrays are shared by every band. On the CPU
        # transform fallback they arrive as numpy; convert them to the
        # device once here, otherwise _resample_cupy_native re-uploads
        # the same two chunk-sized arrays on every band iteration (#3268).
        if not isinstance(local_row, cp.ndarray):
            local_row = cp.asarray(local_row)
        if not isinstance(local_col, cp.ndarray):
            local_col = cp.asarray(local_col)
        bands = []
        for b in range(n_bands):
            band_data = window[:, :, b].astype(cp.float64)
            # Mask this band with its own source sentinel when the raster
            # declares per-band nodata; otherwise fall back to the single
            # resolved sentinel (#2647). Pre-converting to NaN here lets
            # each band use a different source sentinel; the native kernel
            # still fills out-of-bounds pixels with the resolved `nodata`.
            src_nd = band_nodata[b] if band_nodata is not None else nodata
            if not np.isnan(src_nd):
                band_data = cp.where(
                    band_data == src_nd, cp.nan, band_data,
                )
            # Always resample through the native CUDA kernels so the cupy
            # backend matches numpy exactly. They accept CPU coordinate
            # arrays (transferring them to the GPU) and do the
            # nodata->NaN conversion internally, so they serve both the
            # on-device coordinate path and the pyproj fallback. Using
            # cupyx.scipy.ndimage.map_coordinates here instead would
            # diverge from numpy: it bleeds the cval=0.0 constant into the
            # half-pixel boundary band rather than renormalizing, and its
            # order=3 path is a B-spline rather than Catmull-Rom (#2620).
            band_result = _resample_cupy_native(
                band_data, local_row, local_col,
                resampling=resampling, nodata=nodata,
            )
            if np.issubdtype(orig_dtype, np.integer):
                info = np.iinfo(orig_dtype)
                band_result = cp.clip(
                    cp.round(band_result), info.min, info.max,
                ).astype(orig_dtype)
            bands.append(band_result)
        return cp.stack(bands, axis=-1)

    window = window.astype(cp.float64)

    # Always resample through the native CUDA kernels for numpy parity.
    # local_row/local_col may be CuPy (on-device transform) or numpy
    # (pyproj fallback); _resample_cupy_native handles both and does the
    # nodata->NaN conversion internally. The previous
    # cupyx.scipy.ndimage.map_coordinates fallback diverged from numpy at
    # chunk edges and for cubic resampling (#2620).
    result = _resample_cupy_native(window, local_row, local_col,
                                   resampling=resampling, nodata=nodata)

    # Clamp and cast back for integer source dtypes (parity with numpy path)
    if np.issubdtype(orig_dtype, np.integer):
        info = np.iinfo(orig_dtype)
        result = cp.clip(cp.round(result), info.min, info.max).astype(orig_dtype)
    return result


# ---------------------------------------------------------------------------
# reproject()
# ---------------------------------------------------------------------------

def reproject(
    raster,
    target_crs,
    *,
    source_crs=None,
    resolution=None,
    bounds=None,
    width=None,
    height=None,
    resampling='bilinear',
    nodata=None,
    transform_precision=16,
    chunk_size=None,
    name=None,
    max_memory=None,
    source_vertical_crs=None,
    target_vertical_crs=None,
    bounds_policy="auto",
    src_vertical_crs=_DEPRECATED,
    tgt_vertical_crs=_DEPRECATED,
):
    """Reproject a raster DataArray to a new coordinate reference system.

    Supports numpy, cupy, dask+numpy, and dask+cupy backends. For dask
    inputs, the computation is fully lazy: each output chunk independently
    reads only the source pixels it needs.

    Numpy inputs whose input or output working set exceeds ~512 MB are
    routed through the same lazy dask path when dask is installed, so the
    result is dask-backed in that case. Without dask, a streaming
    fallback bounds memory via ``max_memory``.

    Parameters
    ----------
    raster : xr.DataArray
        Input raster with y/x coordinates.
    target_crs
        Target CRS in any format accepted by ``pyproj.CRS()``.
    source_crs : optional
        Source CRS. Auto-detected from *raster* if None.
    resolution : float or (float, float) or None
        Output pixel size in target CRS units.
    bounds : (left, bottom, right, top) or None
        Explicit output extent in target CRS.
    width, height : int or None
        Explicit output grid dimensions.
    resampling : str
        One of 'nearest', 'bilinear', 'cubic'.
    nodata : float or None
        Nodata value. Auto-detected if None. For integer input dtypes,
        an explicit value that does not fit the dtype range raises
        ``ValueError`` (e.g. ``nodata=-9999`` with a ``uint8`` raster).
        Attrs/rioxarray-derived out-of-range values emit a
        ``UserWarning`` and fall back to ``dtype.min`` for signed or
        ``dtype.max`` for unsigned so legacy files still load (#2572).
    transform_precision : int
        Control-grid subdivisions for the coordinate transform (default 16).
        Higher values increase accuracy at the cost of more pyproj calls.
        Set to 0 for exact per-pixel transforms matching GDAL/rasterio.
    chunk_size : int or (int, int) or None
        Output chunk size for dask. If None, defaults to 512 for the
        standard dask path and 2048 for the in-memory streaming and
        dask+cupy paths (chosen to amortize kernel launch overhead).
    name : str or None
        Name for the output DataArray.
    max_memory : int or str or None
        Maximum memory budget for the reprojection working set.
        Accepts bytes (int) or human-readable strings like ``'4GB'``,
        ``'512MB'``.  Controls how many output tiles are processed
        in parallel for large-dataset streaming mode.  Default None
        uses 1GB.  Has no effect for small datasets that fit in memory.
    source_vertical_crs : str or None
        Source vertical datum for height values. One of:

        - ``'EGM96'`` -- orthometric heights relative to EGM96 geoid (MSL)
        - ``'EGM2008'`` -- orthometric heights relative to EGM2008 geoid
        - ``'ellipsoidal'`` -- heights relative to the WGS84 ellipsoid
        - ``None`` -- no vertical transformation (default)
    target_vertical_crs : str or None
        Target vertical datum. Same options as *source_vertical_crs*.
        Both must be set to trigger a vertical transformation.
    src_vertical_crs : str or None
        Deprecated alias for *source_vertical_crs*. Passing it emits a
        ``DeprecationWarning``.
    tgt_vertical_crs : str or None
        Deprecated alias for *target_vertical_crs*. Passing it emits a
        ``DeprecationWarning``.
    bounds_policy : {"auto", "raw", "clamp", "percentile"}, default "auto"
        How to derive the output extent from the source extent when
        ``bounds`` is not supplied. Only relevant when projecting near a
        singularity (antimeridian, pole, projection edge):

        - ``"raw"``: use the true projected extent of the source corners
          and edges. No clamp, no percentile, no heuristic. The output
          may be very large if the input straddles a projection
          singularity. Use this when you want a true projection of the
          source extent.
        - ``"clamp"``: trim geographic source bounds inward by 0.01 deg
          from +/-180 longitude and +/-90 latitude before projecting.
          Avoids infinities at singularities. No percentile fallback.
          No-op on projected source CRSes (UTM, Mercator, etc.) since
          the clamp only applies in degrees.
        - ``"percentile"``: project a dense interior grid of the source
          extent and use the 2nd/98th percentiles of the result as the
          output bounds. Rejects projection outliers at the cost of
          trimming valid pixels.
        - ``"auto"`` (default): apply ``"clamp"`` for geographic source
          CRSes and fall back to ``"percentile"`` when the projected
          extent is more than 50x the source extent. Matches the
          historical behaviour.

        When ``"auto"``, ``"clamp"``, or ``"percentile"`` actually alters
        the bounds, a ``UserWarning`` is emitted naming the policy and
        reporting the per-side delta versus the raw projected bounds.
        Filter with ``warnings.filterwarnings`` if the crop is intentional.

    Returns
    -------
    xr.DataArray
        The output ``attrs['crs']`` is in WKT format.
        Whenever *target_vertical_crs* is set, ``attrs['vertical_crs']``
        records the target vertical datum's EPSG code (5773 for EGM96,
        3855 for EGM2008, 4979 for ellipsoidal WGS84) to match the
        convention used by ``xrspatial.geotiff``. The friendly string
        token (``'EGM96'`` etc.) is preserved under ``attrs['vertical_datum']``.
        Both attrs are written even when no shift is applied (e.g. when
        *source_vertical_crs* equals *target_vertical_crs*, or when only the
        target is given), so the output's vertical reference is always
        explicit.

        The output y coordinate is always emitted in descending order
        (top-down, north-up) and the output x coordinate is always
        emitted in ascending order (left-to-right) regardless of the
        input directions. This matches the standard raster convention
        and the output of common GIS libraries. Inputs with descending
        x are detected from the x coordinate values and handled the
        same way as descending y: the pixel-index mapping is mirrored
        so the output values stay correct.

        Non-spatial coords from the input (such as a scalar ``time``
        coord or a non-dimension coord that is not aligned to the
        spatial dims) are carried through to the output. Coords that
        are aligned to the input y or x dims are dropped because their
        values do not apply to the rebuilt grid.

    Examples
    --------
    >>> import xarray as xr
    >>> import numpy as np
    >>> from xrspatial.reproject import reproject
    >>> raster = xr.DataArray(
    ...     np.random.rand(64, 64),
    ...     dims=['y', 'x'],
    ...     coords={'y': np.linspace(50, 40, 64),
    ...             'x': np.linspace(-5, 5, 64)},
    ...     attrs={'crs': 'EPSG:4326'},
    ... )
    >>> result = reproject(raster, 'EPSG:3857')
    >>> result.attrs['crs'].startswith(('PROJCRS', 'PROJCS'))
    True
    """
    # Back-compat shim for the old abbreviated kwarg names. These were
    # renamed to source_vertical_crs / target_vertical_crs to match the
    # source_crs / target_crs spelling used by the rest of the signature.
    source_vertical_crs = _resolve_deprecated_vertical_kwarg(
        'src_vertical_crs', src_vertical_crs,
        'source_vertical_crs', source_vertical_crs)
    target_vertical_crs = _resolve_deprecated_vertical_kwarg(
        'tgt_vertical_crs', tgt_vertical_crs,
        'target_vertical_crs', target_vertical_crs)

    _validate_raster(raster, func_name='reproject', name='raster',
                     ndim=(2, 3))

    # Reject irregular / non-monotonic source coords before any CRS
    # resolution or grid math. _source_bounds() infers pixel size from
    # only the first two coord samples and downstream pixel math assumes
    # uniform spacing, so an unchecked irregular input would silently
    # produce wrong georeferencing (#2184).
    _validate_source_coords(raster, 'reproject')

    _validate_grid_params(
        resolution=resolution,
        bounds=bounds,
        width=width,
        height=height,
        transform_precision=transform_precision,
        func_name='reproject',
    )

    _validate_resampling(resampling)

    from ._grid import _validate_bounds_policy
    _validate_bounds_policy(bounds_policy, func_name='reproject')

    # Reject unknown vertical-datum tokens at the API boundary so we never
    # write None into attrs['vertical_crs'] for typos / unsupported values.
    for _name, _val in (('source_vertical_crs', source_vertical_crs),
                        ('target_vertical_crs', target_vertical_crs)):
        if _val is not None and _val not in _VERTICAL_DATUM_EPSG:
            raise ValueError(
                f"Unknown {_name}={_val!r}; expected one of "
                f"{sorted(_VERTICAL_DATUM_EPSG)} or None."
            )

    # Normalize 3-D inputs to canonical (y, x, band) layout.
    # The per-chunk workers slice the source as ``source_data[r:, c:]`` and
    # assume the band axis is trailing. A rasterio/rioxarray-style
    # ``(band, y, x)`` input would otherwise slice the band/y axes instead
    # of the y/x axes and either crash or return wrong-shape data (#2182).
    # We record the input's original dim order so the output can be
    # transposed back at the end, preserving downstream expectations.
    _input_dims = tuple(raster.dims)
    if raster.ndim == 3:
        _ydim_in, _xdim_in = _find_spatial_dims(raster)
        _band_dims_in = [d for d in _input_dims
                         if d not in (_ydim_in, _xdim_in)]
        _band_dim_in = _band_dims_in[0] if _band_dims_in else None
        if _band_dim_in is not None:
            _canonical = (_ydim_in, _xdim_in, _band_dim_in)
            if _input_dims != _canonical:
                raster = raster.transpose(*_canonical)

    # Resolve CRS
    src_crs = _resolve_crs(source_crs)
    if src_crs is None:
        src_crs = _detect_source_crs(raster)
    if src_crs is None:
        raise ValueError(
            "Could not detect source CRS. Pass source_crs explicitly."
        )
    tgt_crs = _resolve_crs(target_crs)

    # Detect nodata. Pass the raster dtype so integer rasters get an
    # integer-compatible sentinel (dtype min for signed, dtype max for
    # unsigned) instead of NaN. Without this hint, the worker's
    # cast-back step would collapse NaN to 0 and `attrs['nodata']`
    # would contradict the array contents (#2185).
    nd = _detect_nodata(raster, nodata, dtype=raster.dtype)

    # Multi-band rasters can declare a distinct source sentinel per band
    # via the rasterio `nodatavals` tuple. `nd` is the single resolved
    # output sentinel; `band_nd` carries the raw per-band source sentinels
    # so each band is masked with its own value before resampling (#2647).
    # `None` means one scalar covers every band -- the workers use `nd`.
    # The raster is in canonical (y, x, band) layout here, so the band
    # axis is trailing.
    _n_bands = raster.shape[2] if raster.ndim == 3 else None
    band_nd = _detect_band_nodata(raster, nodata, _n_bands)

    # Source geometry
    src_bounds = _source_bounds(raster)
    _ydim, _xdim = _find_spatial_dims(raster)
    src_shape = (raster.sizes[_ydim], raster.sizes[_xdim])
    y_desc = _is_y_descending(raster)
    x_desc = _is_x_descending(raster)

    # Detect backend before computing the output grid so the grid's
    # output-size guard can tell a lazy dask output (never materialized
    # in full) from a materializing backend.
    from ..utils import has_dask_array, is_cupy_array

    data = raster.data
    is_dask = False
    if has_dask_array():
        import dask.array as _da
        is_dask = isinstance(data, _da.Array)
    is_cupy = False
    if is_dask:
        # Check underlying type
        try:
            from ..utils import is_cupy_backed
            is_cupy = is_cupy_backed(raster)
        except (ImportError, ModuleNotFoundError):
            pass
    else:
        is_cupy = is_cupy_array(data)

    # Compute the output grid with the size guard disabled: whether the
    # guard applies depends on the final execution path (a lazy dask
    # output never materializes the full grid), and the path decision
    # below needs the output shape. The guard is re-applied after the
    # path is known -- same pattern merge() uses.
    grid = _compute_output_grid(
        src_bounds, src_shape, src_crs, tgt_crs,
        resolution=resolution, bounds=bounds,
        width=width, height=height,
        bounds_policy=bounds_policy,
        lazy_output=True,
    )
    out_bounds = grid['bounds']
    out_shape = grid['shape']

    # For large workloads, wrap in dask for chunked processing. Both the
    # input and the output size count: the eager numpy path holds several
    # output-sized float64 temporaries (coordinate grids, pixel-index
    # grids, result), so a small input upsampled to a large output can
    # exhaust memory well before the _MAX_OUTPUT_PIXELS guard trips
    # (#3267). map_blocks generates an O(1) HighLevelGraph (single
    # blockwise layer) so graph metadata is no longer a concern -- the
    # streaming fallback is only needed when dask itself is unavailable.
    _use_streaming = False
    if not is_dask:
        nbytes = src_shape[0] * src_shape[1] * data.dtype.itemsize
        out_nbytes = out_shape[0] * out_shape[1] * 8  # float64 working set
        if data.ndim == 3:
            nbytes *= data.shape[2]
            out_nbytes *= data.shape[2]
        if (nbytes > _REPROJECT_OOM_THRESHOLD
                or out_nbytes > _REPROJECT_OOM_THRESHOLD):
            cs = chunk_size or 2048
            if isinstance(cs, int):
                cs = (cs, cs)
            if data.ndim == 3:
                # Band axis stays one chunk; the per-chunk workers read
                # all bands of a (y, x) window together.
                cs = (cs[0], cs[1], data.shape[2])
            try:
                import dask.array as _da
                # Wrapping an in-memory cupy array keeps the blocks
                # cupy-backed, so is_cupy stays True and the promoted
                # raster routes through _reproject_dask_cupy -- which has
                # its own GPU VRAM check and chunked map_blocks fallback.
                # Above the threshold an in-memory cupy input therefore
                # yields a dask-of-cupy output, matching merge()'s
                # promotion contract (#3281).
                data = _da.from_array(data, chunks=cs)
                raster = xr.DataArray(
                    data, dims=raster.dims, coords=raster.coords,
                    name=raster.name, attrs=raster.attrs,
                )
                is_dask = True
            except ImportError:
                # dask not available. The streaming fallback is a numpy
                # CPU path, so it only applies to numpy inputs; a cupy
                # input keeps the eager GPU path rather than silently
                # moving its data off the device (#3281).
                if not is_cupy:
                    _use_streaming = True

    # Re-apply the output-size guard for paths that materialize the whole
    # output (in-memory numpy/cupy and streaming). A lazy dask output
    # skips it because peak memory is bounded by the chunk size.
    if not is_dask and out_shape[0] * out_shape[1] > _MAX_OUTPUT_PIXELS:
        raise ValueError(
            f"Computed output grid is too large ({out_shape[1]} x "
            f"{out_shape[0]} = {out_shape[0] * out_shape[1]:,} pixels, "
            f"limit is {_MAX_OUTPUT_PIXELS:,}). Increase the resolution "
            f"parameter or reduce the output extent."
        )

    # Output coordinates
    y_coords, x_coords = _make_output_coords(out_bounds, out_shape)

    # Serialize CRS for pickle safety
    src_wkt = src_crs.to_wkt()
    tgt_wkt = tgt_crs.to_wkt()

    if _use_streaming:
        result_data = _reproject_streaming(
            raster, src_bounds, src_shape, y_desc,
            src_wkt, tgt_wkt,
            out_bounds, out_shape,
            resampling, nd, transform_precision,
            chunk_size or 2048,
            _parse_max_memory(max_memory),
            x_desc=x_desc,
            band_nodata=band_nd,
        )
    elif is_dask and is_cupy:
        result_data = _reproject_dask_cupy(
            raster, src_bounds, src_shape, y_desc,
            src_wkt, tgt_wkt,
            out_bounds, out_shape,
            resampling, nd, transform_precision,
            chunk_size,
            x_desc=x_desc,
            band_nodata=band_nd,
        )
    elif is_dask:
        result_data = _reproject_dask(
            raster, src_bounds, src_shape, y_desc,
            src_wkt, tgt_wkt,
            out_bounds, out_shape,
            resampling, nd, transform_precision,
            chunk_size, False,
            x_desc=x_desc,
            band_nodata=band_nd,
        )
    elif is_cupy:
        result_data = _reproject_inmemory_cupy(
            raster, src_bounds, src_shape, y_desc,
            src_wkt, tgt_wkt,
            out_bounds, out_shape,
            resampling, nd, transform_precision,
            x_desc=x_desc,
            band_nodata=band_nd,
        )
    else:
        result_data = _reproject_inmemory_numpy(
            raster, src_bounds, src_shape, y_desc,
            src_wkt, tgt_wkt,
            out_bounds, out_shape,
            resampling, nd, transform_precision,
            x_desc=x_desc,
            band_nodata=band_nd,
        )

    # Vertical datum transformation (if requested)
    if source_vertical_crs is not None and target_vertical_crs is not None:
        if source_vertical_crs != target_vertical_crs:
            result_data, nd = _apply_vertical_shift(
                result_data, y_coords, x_coords,
                source_vertical_crs, target_vertical_crs, nd,
                tgt_crs_wkt=tgt_wkt,
            )

    ydim, xdim = _find_spatial_dims(raster)
    # Carry input attrs forward so units, long_name, scale_factor, etc.
    # survive the transform. Re-emit `transform` and `res` for the new
    # output grid (rasterio 6-tuple convention). Drop `crs_wkt` since it
    # would duplicate (or contradict) the canonical `crs` we re-emit.
    out_left, _out_bottom, _out_right, out_top = grid['bounds']
    out_res_x = grid['res_x']
    out_res_y = grid['res_y']
    out_attrs = {**raster.attrs}
    out_attrs.pop('crs_wkt', None)
    out_attrs['crs'] = tgt_wkt
    out_attrs['nodata'] = nd
    out_attrs['res'] = (out_res_x, out_res_y)
    out_attrs['transform'] = (
        out_res_x, 0.0, out_left, 0.0, -out_res_y, out_top,
    )
    # If the input used `_FillValue`, propagate the resolved nodata to
    # both keys so downstream consumers reading either find a consistent
    # value. If `_FillValue` was absent, leave it absent.
    if '_FillValue' in raster.attrs:
        out_attrs['_FillValue'] = nd
    # `nodatavals` (rasterio convention) is a tuple of per-band sentinels.
    # Refresh it to the resolved nodata so it doesn't contradict
    # ``out_attrs['nodata']`` after the resample.
    if 'nodatavals' in raster.attrs:
        old_nv = raster.attrs['nodatavals']
        try:
            n_entries = max(1, len(old_nv))
        except TypeError:
            n_entries = 1
        out_attrs['nodatavals'] = tuple(nd for _ in range(n_entries))
    if target_vertical_crs is not None:
        # Align with xrspatial.geotiff: attrs['vertical_crs'] holds the
        # EPSG integer code. The friendly string token is preserved under
        # attrs['vertical_datum'] so the human-readable name is not lost.
        # See GH issue #1570.
        out_attrs['vertical_crs'] = _VERTICAL_DATUM_EPSG.get(target_vertical_crs)
        out_attrs['vertical_datum'] = target_vertical_crs

    # Handle multi-band output (3D result from multi-band source)
    if result_data.ndim == 3:
        # Find the band dimension name from the source
        band_dims = [d for d in raster.dims if d not in (ydim, xdim)]
        band_dim = band_dims[0] if band_dims else 'band'
        out_dims = [ydim, xdim, band_dim]
        out_coords = {ydim: y_coords, xdim: x_coords}
        if band_dim in raster.coords:
            out_coords[band_dim] = raster.coords[band_dim]
        # Carry forward non-spatial coords (e.g. scalar 'time' coord).
        # Skip coords aligned to the rebuilt spatial dims because their
        # values do not apply to the new grid.
        for cname, cval in raster.coords.items():
            if cname in (ydim, xdim, band_dim):
                continue
            if ydim in cval.dims or xdim in cval.dims:
                continue
            out_coords[cname] = cval
    else:
        out_dims = [ydim, xdim]
        out_coords = {ydim: y_coords, xdim: x_coords}
        # Carry forward non-spatial coords (e.g. scalar 'time' coord).
        # Skip coords aligned to the rebuilt spatial dims because their
        # values do not apply to the new grid.
        for cname, cval in raster.coords.items():
            if cname in (ydim, xdim):
                continue
            if ydim in cval.dims or xdim in cval.dims:
                continue
            out_coords[cname] = cval

    result = xr.DataArray(
        result_data,
        dims=out_dims,
        coords=out_coords,
        name=name or raster.name,
        attrs=out_attrs,
    )

    # Preserve the input's dim order so a ``(band, y, x)`` source produces a
    # ``(band, y, x)`` output (#2182). The internal pipeline always builds the
    # array as ``(y, x, band)`` for 3-D rasters; transpose back here.
    if result.ndim == 3 and set(_input_dims) == set(result.dims):
        if tuple(result.dims) != _input_dims:
            result = result.transpose(*_input_dims)

    return result


def _promoted_vertical_dtype(src_dtype):
    """Return the float dtype to use when applying a geoid shift.

    The geoid offset is fractional metres, so any integer input has to
    be promoted before we can add it. Returns ``None`` when the input
    is already a float dtype (including ``float64``) -- we leave its
    precision alone. Returns ``np.float32`` for any non-float input;
    that gives ~0.1 mm resolution in metres at altitudes up to ~10 km,
    which is well below the accuracy of the geoid grids themselves.
    """
    src_dtype = np.dtype(src_dtype)
    if np.issubdtype(src_dtype, np.floating):
        return None
    return np.float32


def _apply_vertical_shift(data, y_coords, x_coords,
                          src_vcrs, tgt_vcrs, nodata,
                          tgt_crs_wkt=None):
    """Apply vertical datum shift to reprojected height values.

    The geoid undulation grid is in geographic (lon/lat) coordinates.
    If the output CRS is projected, coordinates are inverse-projected
    to geographic before the geoid lookup.

    Supported vertical CRS:
    - 'EGM96', 'EGM2008': orthometric heights (above geoid/MSL)
    - 'ellipsoidal': heights above WGS84 ellipsoid

    Backend handling
    ----------------
    The geoid undulation lookup is CPU-only (Numba JIT). To stay
    correct across backends:

    - For ``cupy`` arrays, the input is brought to host, shifted, and
      moved back to GPU.
    - For ``dask`` arrays, the shift is wrapped in ``map_blocks`` so
      each chunk's slab is materialised, shifted, and returned in place.
    - For 3-D ``(y, x, band)`` results, the shift is applied per band
      because the geoid undulation depends only on horizontal position.

    Returns
    -------
    (shifted, out_nodata) : tuple
        ``shifted`` is the height-shifted array, possibly promoted to a
        float dtype if the input was integer (see
        ``_promoted_vertical_dtype``). ``out_nodata`` is the nodata
        sentinel that matches the returned dtype -- ``NaN`` whenever a
        promotion happened so the caller can rely on NaN semantics in
        the float output.
    """
    # Direction (sign convention) is the same regardless of backend.
    geoid_models = []
    signs = []
    if src_vcrs in ('EGM96', 'EGM2008') and tgt_vcrs == 'ellipsoidal':
        geoid_models.append(src_vcrs)
        signs.append(1.0)  # H + N = h
    elif src_vcrs == 'ellipsoidal' and tgt_vcrs in ('EGM96', 'EGM2008'):
        geoid_models.append(tgt_vcrs)
        signs.append(-1.0)  # h - N = H
    elif src_vcrs in ('EGM96', 'EGM2008') and tgt_vcrs in ('EGM96', 'EGM2008'):
        geoid_models.extend([src_vcrs, tgt_vcrs])
        signs.extend([1.0, -1.0])  # H1 + N1 - N2
    else:
        return data, nodata

    x_arr = np.asarray(x_coords, dtype=np.float64)
    y_arr = np.asarray(y_coords, dtype=np.float64)

    # Decide the output dtype once, here, so every backend agrees and
    # the caller can update attrs['nodata'] etc. with the post-shift
    # sentinel.
    promoted = _promoted_vertical_dtype(data.dtype)
    if promoted is None:
        out_dtype = np.dtype(data.dtype)
        out_nodata = nodata
    else:
        out_dtype = np.dtype(promoted)
        # NaN is the natural sentinel in a float output. The numpy path
        # below replaces any cells that matched the input integer
        # ``nodata`` with NaN before applying the shift.
        out_nodata = float('nan')

    # Dask backend: wrap the per-block computation. Each block sees its
    # own row slab of x/y coords, so we route through map_blocks and
    # delegate to the numpy path on every chunk.
    try:
        import dask.array as da
        is_dask = isinstance(data, da.Array)
    except ImportError:
        is_dask = False

    if is_dask:
        shifted = _apply_vertical_shift_dask(
            data, y_arr, x_arr,
            geoid_models, signs, nodata, tgt_crs_wkt,
            out_dtype=out_dtype,
        )
        return shifted, out_nodata

    # CuPy backend: round-trip via host. The geoid lookup is small
    # relative to the reprojection itself, and the CPU JIT path is
    # already well-tested.
    try:
        import cupy as cp
        is_cupy = isinstance(data, cp.ndarray)
    except ImportError:
        is_cupy = False

    if is_cupy:
        host = cp.asnumpy(data)
        shifted = _apply_vertical_shift_numpy(
            host, y_arr, x_arr,
            geoid_models, signs, nodata, tgt_crs_wkt,
            out_dtype=out_dtype,
        )
        return cp.asarray(shifted), out_nodata

    shifted = _apply_vertical_shift_numpy(
        np.asarray(data), y_arr, x_arr,
        geoid_models, signs, nodata, tgt_crs_wkt,
        out_dtype=out_dtype,
    )
    return shifted, out_nodata


def _apply_vertical_shift_numpy(data, y_arr, x_arr,
                                geoid_models, signs, nodata, tgt_crs_wkt,
                                out_dtype=None):
    """Apply geoid shift on a numpy array.

    Handles both 2-D ``(H, W)`` and 3-D ``(H, W, B)`` shapes by looping
    over band slices; the geoid undulation depends only on horizontal
    position, so each band sees the same N(y, x) correction.

    If ``out_dtype`` is given and differs from ``data.dtype`` (the
    integer-DEM case), the input is cast to that float dtype before
    the shift and the integer ``nodata`` sentinel is replaced with
    NaN in the promoted output. When ``out_dtype`` is ``None`` (or
    matches the input), the behaviour is in-place on a copy of the
    input -- same as before this signature was added.
    """
    from ._projections import _PARALLEL_KERNEL_LOCK
    from ._vertical import _interp_geoid_2d, _load_geoid

    # Determine if we need inverse projection (output CRS is projected)
    need_inverse = False
    transformer = None
    if tgt_crs_wkt is not None:
        try:
            from ._crs_utils import _require_pyproj
            pyproj = _require_pyproj()
            tgt_crs = pyproj.CRS.from_wkt(tgt_crs_wkt)
            if not tgt_crs.is_geographic:
                need_inverse = True
                geo_crs = pyproj.CRS.from_epsg(4326)
                transformer = pyproj.Transformer.from_crs(
                    tgt_crs, geo_crs, always_xy=True
                )
        except Exception:
            pass

    assert data.ndim in (2, 3), f"expected 2-D or 3-D, got {data.ndim}-D"
    out_h, out_w = data.shape[:2]
    is_3d = data.ndim == 3

    geoids = [_load_geoid(gm) for gm in geoid_models]

    # Process in row strips to bound memory in the numpy path; dask chunks
    # are usually smaller than one strip so this loop runs once per block.
    # If the caller asked for dtype promotion (integer DEM -> float),
    # build ``result`` in the new dtype and rewrite the input nodata
    # sentinel to NaN so the rest of the loop can treat NaN as the
    # missing-value marker uniformly.
    promoted = (
        out_dtype is not None and np.dtype(out_dtype) != np.dtype(data.dtype)
    )
    if promoted:
        result = data.astype(out_dtype, copy=True)
        # Map the source nodata (e.g. -32768 for int16) onto NaN before
        # shifting so downstream consumers can use NaN semantics.
        if isinstance(nodata, float) and np.isnan(nodata):
            pass  # already NaN
        else:
            result[data == nodata] = np.nan
        is_nan_nodata = True
    else:
        result = data.copy()
        is_nan_nodata = np.isnan(nodata) if isinstance(nodata, float) else False
    strip = 128

    for r0 in range(0, out_h, strip):
        r1 = min(r0 + strip, out_h)
        n_rows = r1 - r0

        # Build strip coordinate grid
        xx_strip = np.tile(x_arr, n_rows).reshape(n_rows, out_w)
        yy_strip = np.repeat(y_arr[r0:r1], out_w).reshape(n_rows, out_w)

        # Inverse project if needed
        if need_inverse and transformer is not None:
            lon_s, lat_s = transformer.transform(xx_strip.ravel(), yy_strip.ravel())
            xx_strip = np.asarray(lon_s, dtype=np.float64).reshape(n_rows, out_w)
            yy_strip = np.asarray(lat_s, dtype=np.float64).reshape(n_rows, out_w)

        # Guard against non-finite output coords (projection singularities,
        # antimeridian, polar regions). Hand NaN to the JIT batch so the
        # longitude wrap loop in _interp_geoid_point does not see inf and
        # spin forever.
        finite_coord = np.isfinite(xx_strip) & np.isfinite(yy_strip)
        if not finite_coord.all():
            xx_strip = np.where(finite_coord, xx_strip, np.nan)
            yy_strip = np.where(finite_coord, yy_strip, np.nan)

        # Compute the total shift N_total for this strip once, regardless
        # of how many geoid models contribute.
        N_total = np.zeros((n_rows, out_w), dtype=np.float64)
        for (grid_data, g_left, g_top, g_rx, g_ry, g_h, g_w), sign in zip(geoids, signs):
            N_strip = np.empty((n_rows, out_w), dtype=np.float64)
            # This runs per-chunk under dask's threaded scheduler (via
            # _apply_vertical_shift_dask), and the kernel is parallel=True,
            # so the launch must be serialized (#3141).
            with _PARALLEL_KERNEL_LOCK:
                _interp_geoid_2d(xx_strip, yy_strip, N_strip,
                                 grid_data, g_left, g_top, g_rx, g_ry, g_h, g_w)
            N_total += sign * N_strip

        # Apply to each band slice. For 2-D this loop runs once.
        n_bands = data.shape[2] if is_3d else 1
        for b in range(n_bands):
            if is_3d:
                strip_data = result[r0:r1, :, b]
            else:
                strip_data = result[r0:r1]
            if is_nan_nodata:
                is_valid = np.isfinite(strip_data)
            else:
                is_valid = strip_data != nodata
            is_valid = is_valid & finite_coord
            strip_data[is_valid] += N_total[is_valid]

    return result


def _apply_vertical_shift_dask(data, y_arr, x_arr,
                               geoid_models, signs, nodata, tgt_crs_wkt,
                               out_dtype=None):
    """Dask-backed geoid shift via ``map_blocks``.

    Each block receives only its row/column slab of the input. The block
    function recomputes per-block y/x slices from ``block_info`` and
    delegates to the numpy path so all backends share one implementation.

    ``out_dtype`` is the (possibly promoted) dtype that
    ``_apply_vertical_shift`` has decided for the whole array, and is
    forwarded into every per-block call so the chunks return the same
    promoted dtype the dask graph metadata advertises.
    """
    import dask.array as da

    if out_dtype is None:
        out_dtype = data.dtype

    # Note: blocks reaching this path are assumed to be numpy-backed.
    # dask-of-cupy is intentionally unreached today because
    # ``_reproject_dask(is_cupy=True)`` collapses to a numpy-backed dask
    # array upstream. If that ever changes, _block would need to detect
    # cupy chunks and host-bounce per chunk.
    def _block(block, block_info):
        info = block_info[0]
        (r0, r1), (c0, c1) = info['array-location'][:2]
        y_slab = y_arr[r0:r1]
        x_slab = x_arr[c0:c1]
        return _apply_vertical_shift_numpy(
            block, y_slab, x_slab,
            geoid_models, signs, nodata, tgt_crs_wkt,
            out_dtype=out_dtype,
        )

    # ``meta`` hardcodes a numpy template to match the assumption above
    # that incoming chunks are numpy-backed. Revisit if dask-of-cupy is
    # ever plumbed through.
    return da.map_blocks(_block, data, dtype=out_dtype, meta=np.array((), dtype=out_dtype),
                         **_dask_task_name_kwargs('xrspatial.reproject_vertical_shift'))


def _reproject_inmemory_numpy(
    raster, src_bounds, src_shape, y_desc,
    src_wkt, tgt_wkt,
    out_bounds, out_shape,
    resampling, nodata, precision,
    x_desc=False,
    band_nodata=None,
):
    """Single-chunk numpy reproject."""
    return _reproject_chunk_numpy(
        raster.values,
        src_bounds, src_shape, y_desc,
        src_wkt, tgt_wkt,
        out_bounds, out_shape,
        resampling, nodata, precision,
        source_x_desc=x_desc,
        band_nodata=band_nodata,
    )


def _reproject_inmemory_cupy(
    raster, src_bounds, src_shape, y_desc,
    src_wkt, tgt_wkt,
    out_bounds, out_shape,
    resampling, nodata, precision,
    x_desc=False,
    band_nodata=None,
):
    """Single-chunk cupy reproject."""
    return _reproject_chunk_cupy(
        raster.data,
        src_bounds, src_shape, y_desc,
        src_wkt, tgt_wkt,
        out_bounds, out_shape,
        resampling, nodata, precision,
        source_x_desc=x_desc,
        band_nodata=band_nodata,
    )


def _parse_max_memory(max_memory):
    """Parse max_memory parameter to bytes.  Accepts int, '4GB', '512MB'."""
    if max_memory is None:
        return 1024 * 1024 * 1024  # 1GB default
    if isinstance(max_memory, (int, float)):
        return int(max_memory)
    s = str(max_memory).strip().upper()
    for suffix, factor in [('TB', 1024**4), ('GB', 1024**3), ('MB', 1024**2), ('KB', 1024)]:
        if s.endswith(suffix):
            return int(float(s[:-len(suffix)]) * factor)
    return int(s)


def _process_tile_batch(batch, source_data, src_bounds, src_shape, y_desc,
                        src_wkt, tgt_wkt, resampling, nodata, precision,
                        max_memory_bytes, tile_mem, x_desc=False,
                        band_nodata=None):
    """Process a batch of tiles within a single worker.

    Uses ThreadPoolExecutor for intra-worker parallelism (Numba
    releases the GIL).  Memory bounded by max_memory_bytes.

    The numba coordinate-transform fast path uses parallel=True kernels,
    which must not run concurrently from multiple host threads (the
    workqueue threading layer aborts the process, #3141).
    try_numba_transform serializes those launches behind
    _projections._PARALLEL_KERNEL_LOCK; resampling stays concurrent.

    Returns list of (row_offset, col_offset, tile_data) tuples.
    """
    max_concurrent = max(1, max_memory_bytes // max(tile_mem, 1))

    def _do_one(job):
        _, _, rchunk, cchunk, cb = job
        return _reproject_chunk_numpy(
            source_data,
            src_bounds, src_shape, y_desc,
            src_wkt, tgt_wkt,
            cb, (rchunk, cchunk),
            resampling, nodata, precision,
            source_x_desc=x_desc,
            band_nodata=band_nodata,
        )

    results = []
    if max_concurrent >= 2 and len(batch) > 1:
        import os
        from concurrent.futures import ThreadPoolExecutor
        n_threads = min(max_concurrent, len(batch), os.cpu_count() or 4)
        with ThreadPoolExecutor(max_workers=n_threads) as pool:
            for sub_start in range(0, len(batch), n_threads):
                sub = batch[sub_start:sub_start + n_threads]
                tiles = list(pool.map(_do_one, sub))
                for job, tile in zip(sub, tiles):
                    ro, co, rchunk, cchunk, _ = job
                    results.append((ro, co, tile))
                del tiles
    else:
        for job in batch:
            ro, co, rchunk, cchunk, _ = job
            tile = _do_one(job)
            results.append((ro, co, tile))
            del tile

    return results


def _reproject_streaming(
    raster, src_bounds, src_shape, y_desc,
    src_wkt, tgt_wkt,
    out_bounds, out_shape,
    resampling, nodata, precision,
    tile_size, max_memory_bytes,
    x_desc=False,
    band_nodata=None,
):
    """Streaming reproject for datasets too large for dask's graph.

    Two modes:
    1. **Local** (no dask.distributed): ThreadPoolExecutor within one
       process, bounded by max_memory.
    2. **Distributed** (dask.distributed active): creates a dask.bag
       with one partition per worker, each partition processes its
       tile batch using threads.  Graph size: O(n_workers), not
       O(n_tiles).

    Memory usage per worker: bounded by max_memory.

    Output dtype follows the same rule as the dask backends: integer
    sources round-trip back to their original dtype (the per-tile worker
    casts tiles back after clamping), floats return float64 (#3093).
    """
    if isinstance(tile_size, int):
        tile_size = (tile_size, tile_size)

    # Match the dask backends: integer sources round-trip back to their
    # original dtype after clamping (the per-tile worker already casts
    # tiles back); floats stay float64. Without this, the streaming path
    # silently promoted integer inputs to float64 while every other
    # backend preserved the source dtype (#3093).
    src_dtype = np.dtype(raster.dtype)
    if np.issubdtype(src_dtype, np.integer):
        out_dtype = src_dtype
    else:
        out_dtype = np.dtype(np.float64)

    # 3-D (y, x, band) sources produce 3-D tiles, so the assembled output
    # needs the band axis too (#3093).
    if raster.data.ndim == 3:
        result_shape = (*out_shape, raster.data.shape[2])
    else:
        result_shape = out_shape

    row_chunks, col_chunks = _compute_chunk_layout(out_shape, tile_size)

    tile_mem = tile_size[0] * tile_size[1] * 8 * 4  # ~4 arrays per tile

    # Build tile job list
    jobs = []
    row_offset = 0
    for rchunk in row_chunks:
        col_offset = 0
        for cchunk in col_chunks:
            cb = _chunk_bounds(
                out_bounds, out_shape,
                row_offset, row_offset + rchunk,
                col_offset, col_offset + cchunk,
            )
            jobs.append((row_offset, col_offset, rchunk, cchunk, cb))
            col_offset += cchunk
        row_offset += rchunk

    # Check if dask.distributed is active
    _use_distributed = False
    try:
        from dask.distributed import get_client
        client = get_client()
        n_distributed_workers = len(client.scheduler_info()['workers'])
        if n_distributed_workers > 0:
            _use_distributed = True
    except (ImportError, ValueError):
        pass

    if _use_distributed and len(jobs) > n_distributed_workers:
        # Distributed: partition tiles across workers via dask.bag
        import dask.bag as db

        # Split jobs into N partitions (one per worker)
        n_parts = min(n_distributed_workers, len(jobs))
        batch_size = math.ceil(len(jobs) / n_parts)
        batches = [jobs[i:i + batch_size] for i in range(0, len(jobs), batch_size)]

        # Create bag and map the batch processor
        bag = db.from_sequence(batches, npartitions=len(batches))
        results_bag = bag.map(
            _process_tile_batch,
            source_data=raster.data,
            src_bounds=src_bounds, src_shape=src_shape, y_desc=y_desc,
            src_wkt=src_wkt, tgt_wkt=tgt_wkt,
            resampling=resampling, nodata=nodata, precision=precision,
            max_memory_bytes=max_memory_bytes, tile_mem=tile_mem,
            x_desc=x_desc,
            band_nodata=band_nodata,
        )

        # Compute all partitions and assemble result
        result = np.full(result_shape, nodata, dtype=out_dtype)
        for batch_results in results_bag.compute():
            for ro, co, tile in batch_results:
                result[ro:ro + tile.shape[0], co:co + tile.shape[1]] = tile
        return result

    # Local: ThreadPoolExecutor within one process
    result = np.full(result_shape, nodata, dtype=out_dtype)
    batch_results = _process_tile_batch(
        jobs, raster.data,
        src_bounds, src_shape, y_desc,
        src_wkt, tgt_wkt,
        resampling, nodata, precision,
        max_memory_bytes, tile_mem,
        x_desc=x_desc,
        band_nodata=band_nodata,
    )
    for ro, co, tile in batch_results:
        result[ro:ro + tile.shape[0], co:co + tile.shape[1]] = tile

    return result


def _reproject_dask_cupy(
    raster, src_bounds, src_shape, y_desc,
    src_wkt, tgt_wkt,
    out_bounds, out_shape,
    resampling, nodata, precision,
    chunk_size,
    x_desc=False,
    band_nodata=None,
):
    """Dask+CuPy backend: process output chunks on GPU.

    Two modes based on available GPU memory:

    **Fast path** (output fits in GPU VRAM): pre-allocates the full output
    on GPU and fills it chunk-by-chunk.  ~22x faster than the map_blocks
    path because CRS/transformer objects are created once and CUDA kernels
    run with minimal launch overhead.

    **Chunked fallback** (output exceeds GPU VRAM): delegates to
    ``_reproject_dask(is_cupy=True)`` which uses ``map_blocks`` and
    processes one chunk at a time with O(chunk_size) GPU memory.
    """
    import cupy as cp

    from ._crs_utils import _crs_from_wkt

    src_crs = _crs_from_wkt(src_wkt)
    tgt_crs = _crs_from_wkt(tgt_wkt)

    # Use larger chunks for GPU to amortize kernel launch overhead
    gpu_chunk = chunk_size or 2048
    if isinstance(gpu_chunk, int):
        gpu_chunk = (gpu_chunk, gpu_chunk)

    row_chunks, col_chunks = _compute_chunk_layout(out_shape, gpu_chunk)
    out_h, out_w = out_shape
    src_left, src_bottom, src_right, src_top = src_bounds
    src_h, src_w = src_shape
    src_res_x = (src_right - src_left) / src_w
    src_res_y = (src_top - src_bottom) / src_h

    # 3-D source: the fast path's inline loop assumes 2-D windows.
    # Delegate to the map_blocks path which handles 3-D via
    # _reproject_chunk_cupy's per-band loop (#2027).
    if raster.data.ndim == 3:
        return _reproject_dask(
            raster, src_bounds, src_shape, y_desc,
            src_wkt, tgt_wkt,
            out_bounds, out_shape,
            resampling, nodata, precision,
            chunk_size or 2048, True,  # is_cupy=True
            x_desc=x_desc,
            band_nodata=band_nodata,
        )

    # Memory check: if the full output doesn't fit in GPU memory,
    # fall back to the map_blocks path which is O(chunk_size) memory.
    estimated = out_shape[0] * out_shape[1] * 8  # float64
    try:
        free_gpu, _total = cp.cuda.Device().mem_info
        fits_in_gpu = estimated < 0.5 * free_gpu
    except (AttributeError, RuntimeError):
        fits_in_gpu = False

    if not fits_in_gpu:
        import warnings
        warnings.warn(
            f"Output ({estimated / 1e9:.1f} GB) exceeds GPU memory; "
            f"falling back to chunked map_blocks path.",
            stacklevel=3,
        )
        return _reproject_dask(
            raster, src_bounds, src_shape, y_desc,
            src_wkt, tgt_wkt,
            out_bounds, out_shape,
            resampling, nodata, precision,
            chunk_size or 2048, True,  # is_cupy=True
            x_desc=x_desc,
            band_nodata=band_nodata,
        )

    # Match the dask+numpy and chunked dask+cupy paths: integer sources
    # round-trip back to their original dtype after clamping; floats stay
    # float64. Without this, the eager dask+cupy fast path silently
    # promoted int16/uint8/etc. inputs to float64 while the other
    # backends preserved the source dtype (#2505).
    src_dtype = np.dtype(raster.dtype)
    if np.issubdtype(src_dtype, np.integer):
        out_dtype = src_dtype
    else:
        out_dtype = np.dtype(np.float64)

    result = cp.full(out_shape, nodata, dtype=out_dtype)

    row_offset = 0
    for i, rchunk in enumerate(row_chunks):
        col_offset = 0
        for j, cchunk in enumerate(col_chunks):
            cb = _chunk_bounds(
                out_bounds, out_shape,
                row_offset, row_offset + rchunk,
                col_offset, col_offset + cchunk,
            )
            chunk_shape = (rchunk, cchunk)

            # CUDA coordinate transform (reuses cached CRS objects).
            # precision == 0 forces the exact pyproj path, so skip CUDA.
            cuda_coords = None
            if precision != 0:
                try:
                    from ._projections_cuda import try_cuda_transform
                    cuda_coords = try_cuda_transform(
                        src_crs, tgt_crs, cb, chunk_shape,
                    )
                except (ImportError, ModuleNotFoundError):
                    cuda_coords = None

            if cuda_coords is not None:
                src_y, src_x = cuda_coords
                if x_desc:
                    src_col_px = (src_right - src_x) / src_res_x - 0.5
                else:
                    src_col_px = (src_x - src_left) / src_res_x - 0.5
                if y_desc:
                    src_row_px = (src_top - src_y) / src_res_y - 0.5
                else:
                    src_row_px = (src_y - src_bottom) / src_res_y - 0.5

                # Batch the four reductions into a single device-to-host
                # transfer instead of four separate synchronous .get() calls.
                mins_maxes = cp.stack([
                    cp.nanmin(src_row_px), cp.nanmax(src_row_px),
                    cp.nanmin(src_col_px), cp.nanmax(src_col_px),
                ])
                r_min_val, r_max_val, c_min_val, c_max_val = (
                    float(v) for v in mins_maxes.get()
                )
                if not (np.isfinite(r_min_val) and np.isfinite(r_max_val)
                        and np.isfinite(c_min_val) and np.isfinite(c_max_val)):
                    col_offset += cchunk
                    continue
                r_min = int(np.floor(r_min_val)) - 2
                r_max = int(np.ceil(r_max_val)) + 3
                c_min = int(np.floor(c_min_val)) - 2
                c_max = int(np.ceil(c_max_val)) + 3
            else:
                # CPU fallback for this chunk
                from ._crs_utils import _require_pyproj
                pyproj = _require_pyproj()
                transformer = pyproj.Transformer.from_crs(
                    tgt_crs, src_crs, always_xy=True
                )
                src_y, src_x = _transform_coords(
                    transformer, cb, chunk_shape, precision,
                    src_crs=src_crs, tgt_crs=tgt_crs,
                )
                if x_desc:
                    src_col_px = (src_right - src_x) / src_res_x - 0.5
                else:
                    src_col_px = (src_x - src_left) / src_res_x - 0.5
                if y_desc:
                    src_row_px = (src_top - src_y) / src_res_y - 0.5
                else:
                    src_row_px = (src_y - src_bottom) / src_res_y - 0.5
                r_min = np.nanmin(src_row_px)
                r_max = np.nanmax(src_row_px)
                c_min = np.nanmin(src_col_px)
                c_max = np.nanmax(src_col_px)
                if not np.isfinite(r_min) or not np.isfinite(r_max):
                    col_offset += cchunk
                    continue
                if not np.isfinite(c_min) or not np.isfinite(c_max):
                    col_offset += cchunk
                    continue
                r_min = int(np.floor(r_min)) - 2
                r_max = int(np.ceil(r_max)) + 3
                c_min = int(np.floor(c_min)) - 2
                c_max = int(np.ceil(c_max)) + 3

            # Check overlap
            if r_min >= src_h or r_max <= 0 or c_min >= src_w or c_max <= 0:
                col_offset += cchunk
                continue

            r_min_clip = max(0, r_min)
            r_max_clip = min(src_h, r_max)
            c_min_clip = max(0, c_min)
            c_max_clip = min(src_w, c_max)

            # Guard: cap source window to prevent GPU OOM (matches numpy path)
            _MAX_WINDOW_PIXELS = 64 * 1024 * 1024  # 64 Mpix
            win_pixels = (r_max_clip - r_min_clip) * (c_max_clip - c_min_clip)
            if win_pixels > _MAX_WINDOW_PIXELS:
                col_offset += cchunk
                continue

            # Fetch only the needed source window from dask
            window = raster.data[r_min_clip:r_max_clip, c_min_clip:c_max_clip]
            if hasattr(window, 'compute'):
                window = window.compute()
            if not isinstance(window, cp.ndarray):
                window = cp.asarray(window)
            window = window.astype(cp.float64)

            if not np.isnan(nodata):
                window[window == nodata] = cp.nan

            local_row = src_row_px - r_min_clip
            local_col = src_col_px - c_min_clip

            # Always use the native CUDA kernels (numpy parity). The
            # cupyx.scipy.ndimage.map_coordinates fallback diverged from
            # numpy at chunk edges and for cubic resampling (#2620).
            # local_row/local_col may be numpy here (pyproj fallback);
            # _resample_cupy_native transfers them to the GPU.
            chunk_data = _resample_cupy_native(
                window, local_row, local_col,
                resampling=resampling, nodata=nodata,
            )

            # Clamp + cast back for integer source dtypes so this fast
            # path returns the same dtype as the other backends (#2505).
            # Matches the per-chunk cast in _reproject_chunk_cupy.
            if np.issubdtype(out_dtype, np.integer):
                info = np.iinfo(out_dtype)
                chunk_data = cp.clip(
                    cp.round(chunk_data), info.min, info.max,
                ).astype(out_dtype)

            result[row_offset:row_offset + rchunk,
                   col_offset:col_offset + cchunk] = chunk_data
            col_offset += cchunk
        row_offset += rchunk

    return result


def _finite_pair_bbox(tx, ty):
    """Bounding box of (tx, ty) pairs where both coordinates are finite.

    The x and y coordinates must be filtered together: a transform can
    send some probe points to NaN/inf, and dropping finite x and finite
    y independently would mix coordinates from different points into one
    box. Returns ``(left, bottom, right, top)`` or ``None`` when no pair
    is finite in both coordinates.
    """
    tx = np.asarray(tx, dtype=np.float64)
    ty = np.asarray(ty, dtype=np.float64)
    mask = np.isfinite(tx) & np.isfinite(ty)
    if not mask.any():
        return None
    tx = tx[mask]
    ty = ty[mask]
    return (float(tx.min()), float(ty.min()), float(tx.max()), float(ty.max()))


def _source_footprint_in_target(src_bounds, src_wkt, tgt_wkt):
    """Compute approximate bounding box of source raster in target CRS."""
    try:
        from ._crs_utils import _crs_from_wkt, _resolve_crs
        try:
            src_crs = _crs_from_wkt(src_wkt)
        except Exception:
            src_crs = _resolve_crs(src_wkt)
        try:
            tgt_crs = _crs_from_wkt(tgt_wkt)
        except Exception:
            tgt_crs = _resolve_crs(tgt_wkt)
    except Exception:
        return None

    sl, sb, sr, st = src_bounds
    mx = (sl + sr) / 2
    my = (sb + st) / 2
    xs = np.array([sl, mx, sr, sl, mx, sr, sl, mx, sr, sl, sr, mx])
    ys = np.array([sb, sb, sb, my, my, my, st, st, st, mx, mx, sb])

    try:
        from ._projections import transform_points
        result = transform_points(src_crs, tgt_crs, xs, ys)
        if result is not None:
            tx, ty = result
            return _finite_pair_bbox(tx, ty)
    except (ImportError, ModuleNotFoundError):
        pass

    try:
        from ._crs_utils import _require_pyproj
        pyproj = _require_pyproj()
        transformer = pyproj.Transformer.from_crs(src_crs, tgt_crs, always_xy=True)
        tx, ty = transformer.transform(xs.tolist(), ys.tolist())
        return _finite_pair_bbox(tx, ty)
    except Exception:
        return None


def _bounds_overlap(a, b):
    """Return True if bounding boxes *a* and *b* overlap."""
    return a[0] < b[2] and a[2] > b[0] and a[1] < b[3] and a[3] > b[1]


def _reproject_block_adapter(
    block, block_info, source_data,
    src_bounds, src_shape, y_desc,
    src_wkt, tgt_wkt,
    out_bounds, out_shape,
    resampling, nodata, precision,
    is_cupy, src_footprint_tgt, n_bands=None,
    x_desc=False,
    band_nodata=None,
):
    """``map_blocks`` adapter for reprojection.

    Derives chunk bounds from *block_info* and delegates to the
    per-chunk worker.

    For 3-D sources the template carries the band axis, so each block
    is ``(rh, rw, n_bands)``. The adapter strips the trailing band axis
    when computing 2-D chunk bounds and the per-chunk worker returns a
    3-D result that fits the template (#2027).
    """
    info = block_info[0]
    # 3-D template -> array-location is 3 entries; spatial dims are the
    # first two. Band dim spans the full axis (single chunk).
    spatial_loc = info['array-location'][:2]
    (row_start, row_end), (col_start, col_end) = spatial_loc
    chunk_shape = (row_end - row_start, col_end - col_start)
    cb = _chunk_bounds(out_bounds, out_shape,
                       row_start, row_end, col_start, col_end)

    is_3d = n_bands is not None

    # Skip chunks that don't overlap the source footprint. The fill dtype
    # must match the template: integer sources advertise their own dtype,
    # so a float64 fill here would promote the assembled output (#3096).
    if src_footprint_tgt is not None and not _bounds_overlap(cb, src_footprint_tgt):
        if np.issubdtype(source_data.dtype, np.integer):
            empty_dtype = source_data.dtype
        else:
            empty_dtype = np.float64
        empty_shape = (*chunk_shape, n_bands) if is_3d else chunk_shape
        if is_cupy:
            # Match the data chunks (and the 3-D branch): dask+cupy
            # blocks should be cupy arrays even when empty.
            import cupy as cp
            return cp.full(empty_shape, nodata, dtype=empty_dtype)
        return np.full(empty_shape, nodata, dtype=empty_dtype)

    chunk_fn = _reproject_chunk_cupy if is_cupy else _reproject_chunk_numpy
    return chunk_fn(
        source_data, src_bounds, src_shape, y_desc,
        src_wkt, tgt_wkt,
        cb, chunk_shape,
        resampling, nodata, precision,
        source_x_desc=x_desc,
        band_nodata=band_nodata,
    )


def _reproject_dask(
    raster, src_bounds, src_shape, y_desc,
    src_wkt, tgt_wkt,
    out_bounds, out_shape,
    resampling, nodata, precision,
    chunk_size, is_cupy,
    x_desc=False,
    band_nodata=None,
):
    """Dask+NumPy backend: ``map_blocks`` over a template array.

    Uses a single ``blockwise`` layer in the HighLevelGraph instead of
    O(N) ``dask.delayed`` nodes, keeping graph metadata O(1).

    The source dask array is bound to the adapter via ``functools.partial``
    rather than passed as a ``map_blocks`` kwarg.  This prevents dask from
    adding the full source as a dependency of every output block (which
    would cause a MemoryError on distributed schedulers when the source
    exceeds the worker memory limit).

    For 3-D sources with shape ``(H, W, n_bands)`` the template is built
    as ``(out_H, out_W, n_bands)`` so the lazy metadata matches the
    actual chunk output shape (#2027). Without this, the lazy DataArray
    advertised 2-D shape while the underlying chunks produced 3-D
    arrays, causing a ``ValueError: replacement data has shape ...``
    crash on ``.compute()``.
    """
    import functools

    import dask.array as da

    row_chunks, col_chunks = _compute_chunk_layout(out_shape, chunk_size)

    # Precompute source footprint in target CRS for empty-chunk skipping
    src_footprint_tgt = _source_footprint_in_target(
        src_bounds, src_wkt, tgt_wkt
    )

    # Detect 3-D source: chunks will return (rh, rw, n_bands) so the
    # template must carry the band axis through to the lazy DataArray.
    is_3d = raster.data.ndim == 3
    n_bands = raster.data.shape[2] if is_3d else None

    # Bind the source dask array and all scalar params via partial so
    # map_blocks doesn't detect them as dask Array kwargs (which would
    # add the full source as a dependency of every output block).
    bound_adapter = functools.partial(
        _reproject_block_adapter,
        source_data=raster.data,
        src_bounds=src_bounds,
        src_shape=src_shape,
        y_desc=y_desc,
        src_wkt=src_wkt,
        tgt_wkt=tgt_wkt,
        out_bounds=out_bounds,
        out_shape=out_shape,
        resampling=resampling,
        nodata=nodata,
        precision=precision,
        is_cupy=is_cupy,
        src_footprint_tgt=src_footprint_tgt,
        n_bands=n_bands,
        x_desc=x_desc,
        band_nodata=band_nodata,
    )

    # Pick the template dtype to match the eager path: integer sources
    # round-trip back to their original dtype after clamping; floats stay
    # float64. Without this, dask claims float64 meta but the chunks
    # actually return the integer dtype, producing inconsistent output.
    src_dtype = np.dtype(raster.dtype)
    if np.issubdtype(src_dtype, np.integer):
        out_dtype = src_dtype
    else:
        out_dtype = np.dtype(np.float64)

    if is_3d:
        # Band axis is one chunk in the template regardless of how the
        # source dask array is chunked along its band axis. The per-block
        # worker reads all bands together (via a 2-D y/x slice that
        # rejoins band chunks on compute) and emits the full band stack
        # for its (rh, rw) tile, so multi-chunk output along bands would
        # never get filled.
        template = da.empty(
            (*out_shape, n_bands),
            dtype=out_dtype,
            chunks=(row_chunks, col_chunks, (n_bands,)),
        )
    else:
        template = da.empty(
            out_shape, dtype=out_dtype, chunks=(row_chunks, col_chunks)
        )

    return da.map_blocks(
        bound_adapter,
        template,
        dtype=out_dtype,
        meta=np.array((), dtype=out_dtype),
        **_dask_task_name_kwargs('xrspatial.reproject'),
    )


# ---------------------------------------------------------------------------
# merge()
# ---------------------------------------------------------------------------

def _merge_inputs_to_host(rasters):
    """Bring cupy-backed merge inputs to the host.

    The merge pipeline is numpy-based end to end (``_place_same_crs``,
    ``_reproject_chunk_numpy``, ``_merge_arrays_numpy``), so cupy and
    dask+cupy inputs are converted before any pixel data is touched.
    Returns ``(any_cupy, host_rasters)``; the input list is returned
    unchanged when no raster is GPU-backed. Dask-of-cupy inputs stay
    lazy: each chunk is converted via ``map_blocks`` at compute time.
    """
    try:
        import cupy as cp
    except ImportError:
        return False, rasters

    from ..utils import has_dask_array, is_dask_cupy

    # is_dask_cupy dereferences dask.array.Array, so it must only run
    # when dask is importable (cupy without dask is a valid install).
    _dask_ok = has_dask_array()

    any_cupy = False
    host = []
    for r in rasters:
        if isinstance(r.data, cp.ndarray):
            any_cupy = True
            host.append(r.copy(data=cp.asnumpy(r.data)))
        elif _dask_ok and is_dask_cupy(r):
            any_cupy = True
            host.append(r.copy(data=r.data.map_blocks(
                cp.asnumpy, dtype=r.dtype,
                meta=np.array((), dtype=r.dtype),
            )))
        else:
            host.append(r)
    return any_cupy, host


def merge(
    rasters,
    *,
    target_crs=None,
    resolution=None,
    bounds=None,
    resampling='bilinear',
    nodata=None,
    strategy='first',
    chunk_size=None,
    transform_precision=16,
    name=None,
    bounds_policy="auto",
):
    """Merge multiple rasters into a single mosaic.

    Each input is reprojected to the target CRS (if needed) and placed
    into a unified output grid. Overlapping regions are resolved using
    the selected *strategy*.

    Accepts numpy, cupy, dask+numpy, and dask+cupy inputs. The merge
    itself runs on the CPU: cupy-backed rasters are copied to the host
    on entry and the mosaic is returned as a cupy (or dask+cupy) array
    when any input was GPU-backed.

    Parameters
    ----------
    rasters : list of xr.DataArray
        Input rasters to merge.
    target_crs : optional
        Target CRS. Defaults to the CRS of the first raster.
    resolution : float or (float, float) or None
        Output resolution in target CRS units.
    bounds : (left, bottom, right, top) or None
        Explicit output extent.
    resampling : str
        Interpolation method: 'nearest', 'bilinear', 'cubic'.
    nodata : float or None
        Nodata value for the output. Auto-detected from the first
        raster if None. When every input shares one integer dtype, an
        explicit value that does not fit that dtype range raises
        ``ValueError``, and an auto-detected NaN falls back to
        ``dtype.min`` for signed or ``dtype.max`` for unsigned -- the
        same rules ``reproject`` applies (#2185/#2572).
    strategy : str
        Merge strategy: 'first', 'last', 'mean', 'max', 'min'.
    chunk_size : int or (int, int) or None
        Output chunk size for dask. If None, defaults to 512 for the
        standard dask path and 2048 for the in-memory streaming and
        dask+cupy paths (chosen to amortize kernel launch overhead).
    transform_precision : int
        Control-grid subdivisions for the coordinate transform (default 16).
        Higher values increase accuracy at the cost of more pyproj calls.
        Set to 0 for exact per-pixel transforms matching GDAL/rasterio.
    name : str or None
        Name for the output DataArray.
    bounds_policy : {"auto", "raw", "clamp", "percentile"}, default "auto"
        How to derive the unified output extent from the input rasters
        when ``bounds`` is not supplied. See :func:`reproject` for the
        full description of each option. Ignored when ``bounds`` is given.

    Returns
    -------
    xr.DataArray
        When every input shares the same integer dtype, the mosaic is
        cast back to that dtype (values rounded and clipped, matching
        ``reproject``'s integer round-trip). Mixed input dtypes and
        float inputs produce a float64 mosaic.

        The output y coordinate is always emitted in descending order
        (top-down, north-up) regardless of the input direction. This
        matches the standard raster convention and the output of common
        GIS libraries.

        Non-spatial coords from the first raster (such as a scalar
        ``time`` coord) are carried through to the output. Coords
        aligned to the spatial dims are dropped because their values
        do not apply to the merged grid.

    Notes
    -----
    There is no streaming in-memory path; for very large output mosaics,
    pass dask-backed inputs (or rely on the automatic promotion to the
    dask path) so that each output chunk is computed independently.

    Examples
    --------
    >>> import xarray as xr
    >>> import numpy as np
    >>> from xrspatial.reproject import merge
    >>> tile_a = xr.DataArray(
    ...     np.full((32, 32), 1.0), dims=['y', 'x'],
    ...     coords={'y': np.linspace(50, 45, 32),
    ...             'x': np.linspace(-5, 0, 32)},
    ...     attrs={'crs': 'EPSG:4326'},
    ... )
    >>> tile_b = xr.DataArray(
    ...     np.full((32, 32), 2.0), dims=['y', 'x'],
    ...     coords={'y': np.linspace(50, 45, 32),
    ...             'x': np.linspace(0, 5, 32)},
    ...     attrs={'crs': 'EPSG:4326'},
    ... )
    >>> mosaic = merge([tile_a, tile_b], resolution=0.5)
    >>> mosaic.shape[1] >= 32
    True
    """
    if not rasters:
        raise ValueError("merge(): rasters list must not be empty")

    for i, r in enumerate(rasters):
        # merge() only supports 2-D rasters: the merge strategies, same-CRS
        # placement, and output DataArray construction all assume (y, x).
        # The 3-D (y, x, band) path was never implemented end-to-end and
        # crashed at DataArray construction with a dims-vs-shape mismatch
        # (#2027). Reject 3-D up front so callers get a clear error.
        _validate_raster(r, func_name='merge', name=f'rasters[{i}]',
                         ndim=(2,))
        # Each input must have strictly monotonic, regularly spaced spatial
        # coords. _source_bounds() infers pixel size from only the first
        # two samples, so irregular inputs would otherwise silently produce
        # wrong georeferencing (#2184).
        _validate_source_coords(r, f'merge[rasters[{i}]]')

    _validate_grid_params(
        resolution=resolution,
        bounds=bounds,
        width=None,
        height=None,
        transform_precision=transform_precision,
        func_name='merge',
    )

    _validate_resampling(resampling)
    _validate_strategy(strategy)

    from ._grid import _validate_bounds_policy
    _validate_bounds_policy(bounds_policy, func_name='merge')

    # GPU inputs: the merge pipeline is numpy-based, so cupy and
    # dask+cupy rasters are merged on the host and the mosaic is moved
    # back to the GPU at the end. Same round-trip _apply_vertical_shift
    # uses for its CPU-only geoid lookup. Without this, cupy inputs
    # crashed at `.values` with xarray's implicit-conversion TypeError
    # (#3095).
    _any_cupy, rasters = _merge_inputs_to_host(rasters)

    # Resolve target CRS
    tgt_crs = _resolve_crs(target_crs)
    if tgt_crs is None:
        tgt_crs = _detect_source_crs(rasters[0])
    if tgt_crs is None:
        raise ValueError(
            "Could not detect target CRS. Pass target_crs explicitly."
        )

    # Output dtype follows the reproject() convention (#3262): when every
    # input shares one integer dtype, the mosaic is cast back to that
    # dtype after the float64 merge; mixed dtypes and float inputs return
    # float64. Resolved before nodata detection so the sentinel gets the
    # same dtype-aware treatment as reproject() (#2185/#2572): NaN is
    # swapped for a representable sentinel and an explicit out-of-range
    # nodata raises.
    in_dtypes = {np.dtype(r.dtype) for r in rasters}
    if len(in_dtypes) == 1 and np.issubdtype(next(iter(in_dtypes)), np.integer):
        out_dtype = next(iter(in_dtypes))
    else:
        out_dtype = np.dtype(np.float64)

    # Detect output nodata (the sentinel the user asked for). The merge
    # runs in float64 with NaN as the canonical sentinel, so NaN is fine
    # when the user didn't supply one and the output dtype is float; for
    # integer output the dtype hint resolves a representable sentinel.
    nd = _detect_nodata(rasters[0], nodata, dtype=out_dtype)

    # Gather source info for each raster
    raster_infos = []
    for r in rasters:
        src_crs = _detect_source_crs(r)
        if src_crs is None:
            raise ValueError(
                f"Could not detect CRS for raster '{r.name}'. "
                "Ensure all rasters have CRS metadata."
            )
        sb = _source_bounds(r)
        r_ydim, r_xdim = _find_spatial_dims(r)
        ss = (r.sizes[r_ydim], r.sizes[r_xdim])
        yd = _is_y_descending(r)
        xd = _is_x_descending(r)
        # Per-raster input nodata sentinel. Detected independently of the
        # user-supplied output nodata so that mixed-sentinel inputs are
        # canonicalized correctly during merge. Pass the raster dtype so
        # integer sources get an integer-compatible sentinel rather than
        # NaN -- the same fix as the reproject() path for #2185.
        r_nd = _detect_nodata(r, None, dtype=r.dtype)
        raster_infos.append({
            'raster': r,
            'src_crs': src_crs,
            'src_bounds': sb,
            'src_shape': ss,
            'y_desc': yd,
            'x_desc': xd,
            'src_wkt': src_crs.to_wkt(),
            'raster_nodata': r_nd,
        })

    # Compute unified output grid
    if bounds is None:
        # Union of all raster bounds in target CRS.
        # Capture bounds_policy warnings during the per-input gather and
        # emit a single deduplicated message afterwards. Otherwise a
        # mosaic of N near-antimeridian rasters yields N identical
        # warnings, which is noise the caller cannot act on individually.
        import warnings as _warnings
        all_bounds = []
        with _warnings.catch_warnings(record=True) as _caught:
            _warnings.simplefilter('always', UserWarning)
            for info in raster_infos:
                grid = _compute_output_grid(
                    info['src_bounds'], info['src_shape'],
                    info['src_crs'], tgt_crs,
                    resolution=resolution,
                    bounds_policy=bounds_policy,
                )
                all_bounds.append(grid['bounds'])
        _policy_msgs = [
            str(w.message) for w in _caught
            if issubclass(w.category, UserWarning)
            and 'bounds_policy' in str(w.message)
        ]
        if _policy_msgs:
            _unique = list(dict.fromkeys(_policy_msgs))
            _summary = (
                f"merge: bounds_policy={bounds_policy!r} altered the "
                f"projected extent for {len(_policy_msgs)} input "
                f"raster(s); {len(_unique)} unique trigger(s). "
                f"First trigger: {_unique[0]}"
            )
            _warnings.warn(_summary, UserWarning, stacklevel=2)
        # Re-emit non-bounds_policy warnings that we captured.
        for _w in _caught:
            if 'bounds_policy' not in str(_w.message):
                _warnings.warn_explicit(
                    _w.message, _w.category, _w.filename, _w.lineno,
                )
        left = min(b[0] for b in all_bounds)
        bottom = min(b[1] for b in all_bounds)
        right = max(b[2] for b in all_bounds)
        top = max(b[3] for b in all_bounds)
        merged_bounds = (left, bottom, right, top)
    else:
        merged_bounds = bounds

    # Detect dask inputs up front. A dask-backed merge runs through the
    # lazy _merge_dask path and never materializes the full output, so its
    # output-size guard must be skipped. The auto-promote-to-dask decision
    # below also needs the output shape, which only exists after the grid
    # is computed -- so compute the grid with the guard disabled here and
    # re-apply it only when the final path is the in-memory merge.
    from ..utils import has_dask_array

    any_dask = False
    if has_dask_array():
        import dask.array as _da
        any_dask = any(isinstance(r.data, _da.Array) for r in rasters)

    # Use first raster's info for resolution estimation if needed
    info0 = raster_infos[0]
    grid = _compute_output_grid(
        info0['src_bounds'], info0['src_shape'],
        info0['src_crs'], tgt_crs,
        resolution=resolution, bounds=merged_bounds,
        bounds_policy=bounds_policy,
        lazy_output=True,
    )
    out_bounds = grid['bounds']
    out_shape = grid['shape']
    tgt_wkt = tgt_crs.to_wkt()

    # Auto-promote to dask path if output would be too large for in-memory merge
    if not any_dask:
        out_nbytes = out_shape[0] * out_shape[1] * 8 * len(rasters)  # float64 per tile
        if out_nbytes > _MERGE_OOM_THRESHOLD:
            any_dask = True

    if not any_dask:
        # The final path is the in-memory merge, which allocates the whole
        # output array. Re-apply the output-size guard that was skipped
        # during grid computation so a genuinely in-memory merge over the
        # pixel limit still raises (the guard skip only applies to the
        # lazy dask path). _MAX_OUTPUT_PIXELS is imported from _grid so the
        # limit stays in sync with _compute_output_grid's own check.
        if out_shape[0] * out_shape[1] > _MAX_OUTPUT_PIXELS:
            raise ValueError(
                f"Computed output grid is too large ({out_shape[1]} x "
                f"{out_shape[0]} = {out_shape[0] * out_shape[1]:,} pixels, "
                f"limit is {_MAX_OUTPUT_PIXELS:,}). Increase the resolution "
                f"parameter or reduce the output extent."
            )

    if any_dask:
        result_data = _merge_dask(
            raster_infos, tgt_wkt, out_bounds, out_shape,
            resampling, nd, strategy, chunk_size, transform_precision,
            out_dtype=out_dtype,
        )
    else:
        result_data = _merge_inmemory(
            raster_infos, tgt_wkt, out_bounds, out_shape,
            resampling, nd, strategy, transform_precision,
            out_dtype=out_dtype,
        )

    y_coords, x_coords = _make_output_coords(out_bounds, out_shape)
    ydim, xdim = _find_spatial_dims(rasters[0])

    out_coords = {ydim: y_coords, xdim: x_coords}
    # Carry forward non-spatial coords from the first raster (e.g. scalar
    # 'time' coord). Skip coords aligned to the rebuilt spatial dims
    # because their values do not apply to the new grid.
    for cname, cval in rasters[0].coords.items():
        if cname in (ydim, xdim):
            continue
        if ydim in cval.dims or xdim in cval.dims:
            continue
        out_coords[cname] = cval

    # Carry the first raster's attrs forward (matches the default
    # strategy='first'). Drop the duplicate `crs_wkt` and re-emit
    # `transform` and `res` for the new output grid (rasterio 6-tuple).
    out_left, _out_bottom, _out_right, out_top = grid['bounds']
    out_res_x = grid['res_x']
    out_res_y = grid['res_y']
    out_attrs = {**rasters[0].attrs}
    out_attrs.pop('crs_wkt', None)
    out_attrs['crs'] = tgt_wkt
    out_attrs['nodata'] = nd
    out_attrs['res'] = (out_res_x, out_res_y)
    out_attrs['transform'] = (
        out_res_x, 0.0, out_left, 0.0, -out_res_y, out_top,
    )
    # If the first raster used `_FillValue`, propagate the resolved
    # nodata to both keys for consistent round-trip. Leave absent
    # otherwise.
    if '_FillValue' in rasters[0].attrs:
        out_attrs['_FillValue'] = nd
    # Same treatment for `nodatavals` (rasterio convention).
    if 'nodatavals' in rasters[0].attrs:
        old_nv = rasters[0].attrs['nodatavals']
        try:
            n_entries = max(1, len(old_nv))
        except TypeError:
            n_entries = 1
        out_attrs['nodatavals'] = tuple(nd for _ in range(n_entries))

    # Move the mosaic back to the GPU when any input was cupy-backed.
    # An eager merge returns a cupy array; a lazy merge keeps the dask
    # graph and converts each chunk on compute.
    if _any_cupy:
        import cupy as cp
        if any_dask:
            result_data = result_data.map_blocks(
                cp.asarray, dtype=result_data.dtype,
                meta=cp.array((), dtype=result_data.dtype),
            )
        else:
            result_data = cp.asarray(result_data)

    result = xr.DataArray(
        result_data,
        dims=[ydim, xdim],
        coords=out_coords,
        name=name or rasters[0].name or 'merged',
        attrs=out_attrs,
    )
    return result


def _place_same_crs(src_data, src_bounds, src_shape, y_desc,
                    out_bounds, out_shape, nodata, x_desc=False):
    """Place a same-CRS tile into the output grid by coordinate alignment.

    No reprojection needed -- just index the output rows/columns that
    overlap with the source tile and copy the data.

    The output grid is always north-up with ascending x (row 0 is the
    top of ``out_bounds``, column 0 is the left edge). When ``y_desc``
    is ``False`` the source is y-ascending (row 0 is the bottom of
    ``src_bounds``); in that case the source window is flipped along y
    before being written so the placed data has the same north-up
    orientation as the rest of the output. Without this,
    ``merge([r])`` of a y-ascending raster silently differs from
    ``reproject(r, target_crs=r.crs)`` (#2186).

    ``x_desc`` mirrors the same logic on the horizontal axis: when
    True, ``src_data`` is laid out with column 0 at the maximum x, so
    the source window is reversed along its column axis before
    placement (#2183).
    """
    out_h, out_w = out_shape
    src_h, src_w = src_shape
    o_left, o_bottom, o_right, o_top = out_bounds
    s_left, s_bottom, s_right, s_top = src_bounds

    o_res_x = (o_right - o_left) / out_w
    o_res_y = (o_top - o_bottom) / out_h
    s_res_x = (s_right - s_left) / src_w
    s_res_y = (s_top - s_bottom) / src_h

    # Output pixel range that this tile covers. The output is always
    # north-up so ``row_start`` is measured from ``o_top`` downward.
    col_start = int(round((s_left - o_left) / o_res_x))
    col_end = int(round((s_right - o_left) / o_res_x))
    row_start = int(round((o_top - s_top) / o_res_y))
    row_end = int(round((o_top - s_bottom) / o_res_y))

    # Clip to output bounds
    col_start_clip = max(0, col_start)
    col_end_clip = min(out_w, col_end)
    row_start_clip = max(0, row_start)
    row_end_clip = min(out_h, row_end)

    if col_start_clip >= col_end_clip or row_start_clip >= row_end_clip:
        return np.full(out_shape, nodata, dtype=np.float64)

    # Resolutions may differ slightly; if close enough, do direct copy
    res_ratio_x = s_res_x / o_res_x
    res_ratio_y = s_res_y / o_res_y
    if abs(res_ratio_x - 1.0) > 0.01 or abs(res_ratio_y - 1.0) > 0.01:
        return None  # resolutions too different, fall back to reproject

    # Source column offset (columns always run left-to-right).
    src_col_start = col_start_clip - col_start
    src_c_end = min(src_col_start + (col_end_clip - col_start_clip), src_w)
    actual_cols = src_c_end - src_col_start

    out_data = np.full(out_shape, nodata, dtype=np.float64)

    if actual_cols <= 0:
        return out_data

    # Resolve the source column slice. When ``x_desc`` is True the
    # source has column 0 at max x, so we read a mirrored slice and
    # reverse it; otherwise we read the natural left-to-right slice.
    # ``src_col_start`` / ``src_c_end`` were computed above in
    # ascending-x source space and stay valid for both branches.
    if x_desc:
        c_lo = src_w - src_c_end
        c_hi = src_w - src_col_start
    else:
        c_lo = src_col_start
        c_hi = src_c_end

    if y_desc:
        # Source is north-up: rows go top-to-bottom, same as output.
        # Output row R corresponds to source row ``R - row_start``.
        src_row_start = row_start_clip - row_start
        src_r_end = min(src_row_start + (row_end_clip - row_start_clip),
                        src_h)
        actual_rows = src_r_end - src_row_start
        if actual_rows <= 0:
            return out_data
        src_window = np.asarray(
            src_data[src_row_start:src_r_end, c_lo:c_hi],
            dtype=np.float64,
        )
        if x_desc:
            src_window = src_window[:, ::-1]
        out_row_start = row_start_clip
    else:
        # Source is south-up: source row 0 sits at the bottom of
        # ``src_bounds``. Output row R corresponds to source row
        # ``row_end - 1 - R`` (with R and row_end both measured from
        # ``o_top`` downward, so larger R means lower latitude and
        # smaller source row index). Concretely, if src_h=4 and
        # row_end=4, output row 0 maps to source row 3, output row 3
        # maps to source row 0 -- a vertical flip. Read a contiguous
        # ascending slice of source rows and reverse it so the
        # placed window comes out north-up.
        src_lo = max(0, row_end - row_end_clip)
        src_hi = min(src_h, row_end - row_start_clip)
        actual_rows = src_hi - src_lo
        if actual_rows <= 0:
            return out_data
        src_window = np.asarray(
            src_data[src_lo:src_hi, c_lo:c_hi],
            dtype=np.float64,
        )[::-1, :]
        if x_desc:
            src_window = src_window[:, ::-1]
        # The reversed slice's first row corresponds to source row
        # ``src_hi - 1``, which maps to output row ``row_end - src_hi``.
        # When the source is not clipped at the top, this equals
        # ``row_start_clip``; clipping at the source top shifts the
        # placement downward by the clipped row count.
        out_row_start = row_end - src_hi

    out_data[out_row_start:out_row_start + actual_rows,
             col_start_clip:col_start_clip + actual_cols] = src_window
    return out_data


def _cast_merged_dtype(merged, out_dtype):
    """Cast a float64 merged mosaic back to an integer output dtype.

    Matches the round/clip/cast convention the reproject() workers use
    for integer sources (#2505/#3093). Must run after the NaN canonical
    sentinel has been converted back to the resolved output nodata --
    for integer ``out_dtype`` that sentinel is guaranteed non-NaN by
    ``_detect_nodata``'s dtype hint, so no NaN survives to the cast.
    No-op for float output dtypes.
    """
    if np.issubdtype(out_dtype, np.integer):
        info = np.iinfo(out_dtype)
        merged = np.clip(np.round(merged), info.min, info.max).astype(out_dtype)
    return merged


def _merge_inmemory(
    raster_infos, tgt_wkt, out_bounds, out_shape,
    resampling, nodata, strategy, transform_precision,
    out_dtype=np.float64,
):
    """In-memory merge using numpy.

    Detects same-CRS tiles and uses fast direct placement instead
    of reprojection.

    Each raster is reprojected using its own nodata sentinel and then
    canonicalized to NaN before the strategy merge so that mixed-sentinel
    inputs do not leak invalid pixels into the mosaic. After merging, the
    NaN canonical sentinel is converted back to the user-requested
    output ``nodata``.
    """
    from ._crs_utils import _crs_from_wkt
    tgt_crs = _crs_from_wkt(tgt_wkt)

    arrays = []
    for info in raster_infos:
        r_nd = info.get('raster_nodata', float('nan'))
        # Check if source CRS matches target (no reprojection needed)
        placed = None
        if info['src_crs'] == tgt_crs:
            placed = _place_same_crs(
                info['raster'].values,
                info['src_bounds'], info['src_shape'], info['y_desc'],
                out_bounds, out_shape, r_nd,
                x_desc=info.get('x_desc', False),
            )
        if placed is not None:
            arr = placed
        else:
            arr = _reproject_chunk_numpy(
                info['raster'].values,
                info['src_bounds'], info['src_shape'], info['y_desc'],
                info['src_wkt'], tgt_wkt,
                out_bounds, out_shape,
                resampling, r_nd, transform_precision,
                source_x_desc=info.get('x_desc', False),
            )
        # Canonicalize this raster's sentinel to NaN before the merge so
        # rasters with different sentinels merge correctly.
        if not np.isnan(r_nd):
            arr = np.asarray(arr, dtype=np.float64)
            arr = np.where(arr == r_nd, np.nan, arr)
        arrays.append(arr)

    merged = _merge_arrays_numpy(arrays, float('nan'), strategy)
    if not np.isnan(nodata):
        merged = np.where(np.isnan(merged), nodata, merged)
    return _cast_merged_dtype(merged, out_dtype)


def _merge_block_adapter(
    block, block_info,
    raster_data_list, src_bounds_list, src_shape_list, y_desc_list,
    src_wkt_list, tgt_wkt,
    out_bounds, out_shape,
    resampling, nodata, strategy, precision,
    src_footprints_tgt, raster_nodata_list, same_crs_list,
    x_desc_list=None,
    out_dtype=np.float64,
):
    """``map_blocks`` adapter for merge.

    Each raster is reprojected with its own input nodata sentinel and
    canonicalized to NaN before the strategy merge. The final result is
    converted back to the user-requested output ``nodata``.
    """
    info = block_info[0]
    (row_start, row_end), (col_start, col_end) = info['array-location']
    chunk_shape = (row_end - row_start, col_end - col_start)
    cb = _chunk_bounds(out_bounds, out_shape,
                       row_start, row_end, col_start, col_end)

    # Only reproject rasters whose footprint overlaps this chunk
    arrays = []
    for i in range(len(raster_data_list)):
        if (src_footprints_tgt[i] is not None
                and not _bounds_overlap(cb, src_footprints_tgt[i])):
            continue
        r_nd = raster_nodata_list[i]
        placed = None
        xd_i = x_desc_list[i] if x_desc_list is not None else False
        if same_crs_list[i]:
            # Same-CRS path: direct pixel placement (no resampling).
            # Pass the dask array straight through -- _place_same_crs
            # slices before np.asarray(), so np.asarray on the slice
            # materializes only the source window for this output chunk.
            # An eager .compute() here would materialize the full source
            # per output chunk, amplifying driver-side data flow by O(N).
            placed = _place_same_crs(
                raster_data_list[i],
                src_bounds_list[i], src_shape_list[i], y_desc_list[i],
                cb, chunk_shape, r_nd,
                x_desc=xd_i,
            )
        if placed is not None:
            arr = placed
        else:
            arr = _reproject_chunk_numpy(
                raster_data_list[i],
                src_bounds_list[i], src_shape_list[i], y_desc_list[i],
                src_wkt_list[i], tgt_wkt,
                cb, chunk_shape,
                resampling, r_nd, precision,
                source_x_desc=xd_i,
            )
        if not np.isnan(r_nd):
            arr = np.asarray(arr, dtype=np.float64)
            arr = np.where(arr == r_nd, np.nan, arr)
        arrays.append(arr)

    if not arrays:
        # Empty-chunk fills must match the template dtype, or a single
        # no-overlap chunk would promote the assembled output (#3096).
        return np.full(chunk_shape, nodata, dtype=out_dtype)
    merged = _merge_arrays_numpy(arrays, float('nan'), strategy)
    if not np.isnan(nodata):
        merged = np.where(np.isnan(merged), nodata, merged)
    return _cast_merged_dtype(merged, out_dtype)


def _merge_dask(
    raster_infos, tgt_wkt, out_bounds, out_shape,
    resampling, nodata, strategy, chunk_size, transform_precision,
    out_dtype=np.float64,
):
    """Dask merge backend using ``map_blocks``."""
    import functools

    import dask.array as da

    row_chunks, col_chunks = _compute_chunk_layout(out_shape, chunk_size)

    # Prepare lists for the worker
    data_list = [info['raster'].data for info in raster_infos]
    bounds_list = [info['src_bounds'] for info in raster_infos]
    shape_list = [info['src_shape'] for info in raster_infos]
    ydesc_list = [info['y_desc'] for info in raster_infos]
    xdesc_list = [info.get('x_desc', False) for info in raster_infos]
    wkt_list = [info['src_wkt'] for info in raster_infos]
    rnodata_list = [
        info.get('raster_nodata', float('nan')) for info in raster_infos
    ]

    # Precompute source footprints in target CRS
    footprints = [
        _source_footprint_in_target(bounds_list[i], wkt_list[i], tgt_wkt)
        for i in range(len(raster_infos))
    ]

    # Precompute CRS-equality flags so per-block adapters can shortcut to
    # direct pixel placement (matches the eager _merge_inmemory path).
    from ._crs_utils import _crs_from_wkt
    tgt_crs = _crs_from_wkt(tgt_wkt)
    same_crs_list = [
        bool(_crs_from_wkt(wkt_list[i]) == tgt_crs)
        for i in range(len(raster_infos))
    ]

    # Bind via partial to prevent map_blocks from adding dask arrays
    # in data_list as whole-array dependencies.
    bound_adapter = functools.partial(
        _merge_block_adapter,
        raster_data_list=data_list,
        src_bounds_list=bounds_list,
        src_shape_list=shape_list,
        y_desc_list=ydesc_list,
        src_wkt_list=wkt_list,
        tgt_wkt=tgt_wkt,
        out_bounds=out_bounds,
        out_shape=out_shape,
        resampling=resampling,
        nodata=nodata,
        strategy=strategy,
        precision=transform_precision,
        src_footprints_tgt=footprints,
        raster_nodata_list=rnodata_list,
        same_crs_list=same_crs_list,
        x_desc_list=xdesc_list,
        out_dtype=out_dtype,
    )

    # The template dtype must match what the chunks actually return, or
    # the graph advertises one dtype and computes another (#3096).
    template = da.empty(
        out_shape, dtype=out_dtype, chunks=(row_chunks, col_chunks)
    )

    return da.map_blocks(
        bound_adapter,
        template,
        dtype=out_dtype,
        meta=np.array((), dtype=out_dtype),
        **_dask_task_name_kwargs('xrspatial.merge'),
    )