Preprocessing API

def resample_image(
    image: Image,
    new_spacing: tuple[float, float, float],
    interpolation: str = "linear",
    boundary_mode: str = "nearest",
    round_intensities: bool = False,
    mask_threshold: Optional[float] = None,
    source_mask: Optional[Image | npt.NDArray[np.bool_]] = None,
    weight_threshold: float = 0.5,
    region: Optional[tuple[slice, slice, slice]] = None,
) -> Image:
    """
    Resample image to new voxel spacing using IBSI-compliant 'Align grid centers' method.

    The common cases (3D float64 or uint8 image, 'nearest' boundary, nearest/linear/cubic
    interpolation) run on parallel numba kernels that give scipy's output bit for bit;
    everything else uses scipy.ndimage.affine_transform.

    Args:
        image: Input Image object.
        new_spacing: Target spacing (x, y, z). Must be positive.
        interpolation: Interpolation method.
            'nearest': Nearest neighbour (order 0).
            'linear': Trilinear (order 1).
            'cubic': Tricubic spline (order 3), with the spline prefilter of scipy.
        boundary_mode: Padding mode for extrapolation.
            'nearest' (default): Replicates edge values (aaaa|abcd|dddd).
            'constant': Pads with constant value (0).
            'reflect': Reflects at boundary.
            'wrap': Wraps around.
        round_intensities: If True, round resulting intensities to nearest integer.
        mask_threshold: If provided, treat output as a binary mask.
                        Interpolated values >= threshold become 1, others 0.
                        An interpolated value of every mask type is rounded at
                        the threshold (value + 1 - threshold >= 1): at 0.5, the
                        rounding of scipy's uint8 output.
                        Commonly 0.5 for partial volume correction.
        source_mask: Optional source validity mask. If provided (or if image.source_mask
                     is set), only valid voxels are used for interpolation. This prevents
                     sentinel values (e.g., -2048 in CT) from contaminating the resampled
                     output. Can be an Image object or a boolean numpy array.
        weight_threshold: Only used with a source mask. Minimum fraction of
                     interpolation weight that must come from valid voxels for an
                     output voxel to be considered valid. Default 0.5 (majority).
        region: Optional slices of the new grid. The output then holds only this region,
                     with the values of the whole grid bit for bit (its origin is the
                     region's first voxel), and the rest of the grid is not computed.

    Returns:
        Resampled Image object. If source_mask was used, the output Image will have
        its source_mask attribute set to the resampled validity mask.

    Note:
        When source_mask is active, the function uses normalized interpolation:
        the contribution of each input voxel is weighted by its validity, and the
        result is normalized by the sum of valid weights. This ensures that sentinel
        voxels do not affect the output. A float64 image with linear or cubic
        interpolation runs this on the numba kernels.

    Raises:
        ValueError: If any element of `new_spacing` is not positive, if
            `interpolation` is not one of 'nearest', 'linear', 'cubic', or if
            `source_mask` does not match the image geometry (when passed as an
            Image) or array shape (when passed as a numpy array).

    Example:
        Resample image to isotropic 1mm spacing using linear interpolation:

        ```python
        from pictologics.preprocessing import resample_image

        # Resample to 1x1x1 mm
        resampled_img = resample_image(
            image,
            new_spacing=(1.0, 1.0, 1.0),
            interpolation="linear"
        )
        ```

        Resample with sentinel-value exclusion:

        ```python
        # Image has -2048 outside ROI
        image_with_sentinel = image.with_source_mask(roi_mask)
        resampled = resample_image(
            image_with_sentinel,
            new_spacing=(1.0, 1.0, 1.0)
        )
        # Sentinel voxels were excluded from interpolation
        ```
    """
    if any(s <= 0 for s in new_spacing):
        raise ValueError(f"New spacing must be positive, got {new_spacing}")

    # Determine effective source mask
    effective_source: Optional[npt.NDArray[np.bool_]] = None

    if source_mask is not None:
        if isinstance(source_mask, Image):
            _validate_geometry(source_mask, image, "source mask", "image")
            arr: npt.NDArray[Any] = source_mask.array
            effective_source = arr if arr.dtype == np.bool_ else arr > 0
        else:
            effective_source = source_mask.astype(bool, copy=False)
    elif image.has_source_mask:
        effective_source = image.source_mask

    if effective_source is not None and effective_source.shape != image.array.shape:
        raise ValueError(
            f"Source mask shape {effective_source.shape} must match image shape {image.array.shape}"
        )

    # An all-valid mask is equivalent to no mask; skip the normalized
    # interpolation path, which costs a second affine_transform.
    source_all_valid = effective_source is not None and bool(effective_source.all())
    if source_all_valid:
        effective_source = None

    # Map interpolation string to spline order
    interpolation_map = {
        "nearest": 0,
        "linear": 1,
        "cubic": 3,
    }

    if interpolation not in interpolation_map:
        raise ValueError(
            f"Unknown interpolation method: {interpolation}. "
            f"Supported: {list(interpolation_map.keys())}"
        )

    order = interpolation_map[interpolation]

    # The new grid, and the transform from an output index to an input coordinate:
    # x_in = matrix * x_out + offset ('align grid centers', see _grid)
    out_shape, scale, shift, origin_shift = _output_grid(
        image.array.shape, image.spacing, new_spacing
    )
    matrix = np.array(scale)
    offset = np.array(shift)

    # Perform resampling
    resampled_array: npt.NDArray[Any]
    new_source_mask: Optional[npt.NDArray[np.bool_]] = None
    start = np.zeros(3, dtype=np.int64)

    # The kernels cover the hot path (3D, 'nearest' boundary, common dtypes) at every
    # size, with serial twins below the parallel gates; anything else falls back to
    # scipy.ndimage. The kernels take the shift form scipy's zoom_shift uses internally
    # (coordinate is (idx + offset/scale) * scale) and follow its steps, so their output
    # is bit for bit the scipy output.
    kernel_ok = image.array.ndim == 3 and boundary_mode == "nearest"
    kernel_shift = offset / matrix
    if region is not None:
        dtype = image.array.dtype
        if not (
            kernel_ok
            and (
                (order == 1 and dtype == np.float64)
                or (effective_source is None and order == 1 and dtype == np.uint8)
                or (effective_source is None and order == 0 and dtype in _NEAREST_DTYPES)
                or (order == 3 and dtype == np.float64)
            )
        ):
            # No kernel for this case: cut the whole grid
            whole = resample_image(
                image,
                new_spacing,
                interpolation,
                boundary_mode,
                round_intensities,
                mask_threshold,
                source_mask,
                weight_threshold,
            )
            return _region_of(whole, region)
        start = np.array([r.start for r in region], dtype=np.int64)
        out_shape = tuple(r.stop - r.start for r in region)
    # A mask threshold applies to the interpolated value: a uint8 or bool mask would
    # otherwise round or truncate the value first and lose the threshold.
    thresholded = False
    # The linear kernel rounds the values itself (no second pass over the output)
    rounded = False

    if effective_source is None:
        if kernel_ok and order == 1 and image.array.dtype in (np.float64, np.uint8):
            src = _row_order_copy(image.array)
            thresholded = mask_threshold is not None
            rounded = bool(round_intensities)
            resampled_array = np.empty(
                out_shape, dtype=np.uint8 if thresholded else image.array.dtype
            )
            linear = sized(
                _resample_trilinear_numba,
                _resample_trilinear_numba_serial,
                resampled_array.size,
                _LINEAR_PARALLEL_MIN,
            )
            linear(
                src,
                matrix,
                kernel_shift,
                start,
                image.array.dtype == np.uint8,
                math.nan if mask_threshold is None else float(mask_threshold),
                rounded,
                resampled_array,
            )
        elif kernel_ok and order == 0 and image.array.dtype in _NEAREST_DTYPES:
            src = _row_order_copy(image.array)
            resampled_array = np.empty(out_shape, dtype=image.array.dtype)
            nearest = sized(
                _resample_nearest_numba,
                _resample_nearest_numba_serial,
                resampled_array.size,
                _NEAREST_PARALLEL_MIN,
            )
            nearest(src, matrix, kernel_shift, start, resampled_array)
        elif kernel_ok and order == 3 and image.array.dtype == np.float64:
            coef = _spline_coefficients(_row_order_copy(image.array))
            resampled_array = np.empty(out_shape, dtype=np.float64)
            cubic = sized(
                _resample_tricubic_numba,
                _resample_tricubic_numba_serial,
                resampled_array.size,
                _CUBIC_PARALLEL_MIN,
            )
            cubic(coef, matrix, kernel_shift, start, float(_SPLINE_NPAD), resampled_array)
            del coef
        else:
            # A mask threshold needs the interpolated value, so a non-float64 mask is
            # interpolated in float64. A uint8 mask at the threshold 0.5 keeps scipy's
            # uint8 output, whose rounding gives that same result.
            source = image.array
            if (
                mask_threshold is not None
                and order > 0
                and not (source.dtype == np.uint8 and mask_threshold == 0.5)
            ):
                source = source.astype(np.float64, copy=False)
            resampled_array = affine_transform(
                source,
                matrix=matrix,
                offset=offset,
                output_shape=out_shape,
                order=order,
                mode=boundary_mode,
            )
        if source_all_valid:
            # Preserve the output contract: a resampled image keeps a source mask
            new_source_mask = np.ones(out_shape, dtype=bool)
    elif kernel_ok and order == 1 and image.array.dtype == np.float64:
        # Masked resampling: one fused pass instead of two affine_transforms
        src = _row_order_copy(image.array)
        valid = _row_order_copy(effective_source)
        resampled_array = np.empty(out_shape, dtype=np.float64)
        new_source_mask = np.empty(out_shape, dtype=np.bool_)
        masked = sized(
            _resample_trilinear_masked_numba,
            _resample_trilinear_masked_numba_serial,
            resampled_array.size,
            _MASKED_PARALLEL_MIN,
        )
        masked(
            src,
            valid,
            matrix,
            kernel_shift,
            start,
            float(weight_threshold),
            resampled_array,
            new_source_mask,
        )
    elif kernel_ok and order == 3 and image.array.dtype == np.float64:
        # Masked cubic resampling: the two sums of the normalized interpolation in one pass
        src = _row_order_copy(image.array)
        valid = _row_order_copy(effective_source)
        coef, wcoef = _spline_coefficients_masked(src, valid)
        resampled_array = np.empty(out_shape, dtype=np.float64)
        new_source_mask = np.empty(out_shape, dtype=np.bool_)
        masked_cubic = sized(
            _resample_tricubic_masked_numba,
            _resample_tricubic_masked_numba_serial,
            resampled_array.size,
            _CUBIC_PARALLEL_MIN,
        )
        masked_cubic(
            coef,
            wcoef,
            matrix,
            kernel_shift,
            start,
            float(_SPLINE_NPAD),
            float(weight_threshold),
            resampled_array,
            new_source_mask,
        )
        del coef, wcoef
    else:
        # Masked resampling using normalized interpolation
        resampled_array, new_source_mask = _resample_with_source_mask(
            image.array,
            effective_source,
            matrix=matrix,
            offset=offset,
            output_shape=out_shape,
            order=order,
            mode=boundary_mode,
            weight_threshold=weight_threshold,
        )

    # Post-processing
    if mask_threshold is not None and not thresholded:
        # Binarize mask
        resampled_array = _mask_threshold(resampled_array, mask_threshold, order > 0)
    elif round_intensities and not rounded:
        # Round intensities, in place: the resampled array is always a new one
        np.round(resampled_array, out=resampled_array)

    # Update origin to maintain center alignment (the shift of _grid, in world axes)
    new_origin = tuple(
        np.array(image.origin) + _direction_matrix(image.direction) @ np.array(origin_shift)
    )
    if region is not None:
        new_origin = _region_origin(new_origin, image.direction, new_spacing, region)

    return Image(
        array=resampled_array,
        spacing=new_spacing,
        origin=new_origin,
        direction=image.direction,
        modality=image.modality,
        source_mask=new_source_mask,
    )

def discretise_image(
    image: Image | npt.NDArray[Any],
    method: str,
    roi_mask: Image | npt.NDArray[Any] | None = None,
    n_bins: Optional[int] = None,
    bin_width: Optional[float] = None,
    min_val: Optional[float] = None,
    max_val: Optional[float] = None,
    cutoffs: Optional[list[float]] = None,
) -> Image | npt.NDArray[np.int32]:
    """
    Discretise image intensities.

    Supports IBSI-compliant Fixed Bin Number (FBN) and Fixed Bin Size (FBS).

    Args:
        image: Input Image object or numpy array.
        method: 'FBN' (Fixed Bin Number), 'FBS' (Fixed Bin Size), or 'FIXED_CUTOFFS'.
        roi_mask: Optional mask to define the ROI for determining min/max values.
        n_bins: Number of bins (required for FBN).
        bin_width: Bin width (required for FBS).
        min_val: Minimum value for discretisation.
                 For FBS, defaults to ROI minimum (or global minimum). To compare
                 images, give FBS the same min_val for every image, as IBSI
                 recommends (the lower bound of the resegmentation range); the
                 pipeline always does.
                 For FBN, defaults to ROI minimum.
        max_val: Maximum value for discretisation (FBN only).
                 Defaults to ROI maximum.
        cutoffs: List of cutoffs (required for FIXED_CUTOFFS). Values below the
                 first cutoff map to bin 1; values >= the last cutoff map to
                 bin len(cutoffs) + 1.

    Returns:
        Discretised Image object or numpy array (depending on input).
        Values are 1-based int32 indices; 0 marks NaN (invalid) voxels.

    Raises:
        ValueError: If `roi_mask` shape does not match `image` shape (or, when
            both are Image objects, if their geometry otherwise differs); if
            `method` is 'FBN' and `n_bins` is missing or not positive; if
            `method` is 'FBS' and `bin_width` is missing or not positive; if
            `method` is 'FIXED_CUTOFFS' and `cutoffs` is missing; or if
            `method` is not one of 'FBN', 'FBS', 'FIXED_CUTOFFS'.

    Example:
        Discretise image into 32 fixed bins (FBN):

        ```python
        from pictologics.preprocessing import discretise_image

        # FBN with 32 bins
        disc_image = discretise_image(
            image,
            method="FBN",
            n_bins=32
        )
        ```
    """
    # Handle input type
    array = image.array if isinstance(image, Image) else image

    mask_arr: Optional[npt.NDArray[Any]] = None
    if roi_mask is not None:
        if isinstance(roi_mask, Image):
            if isinstance(image, Image):
                _validate_geometry(roi_mask, image, "ROI mask", "image")
            mask_arr = roi_mask.array
        else:
            mask_arr = roi_mask

        if mask_arr.shape != array.shape:
            raise ValueError(f"Shape mismatch: Image {array.shape} vs Mask {mask_arr.shape}")

    roi_range: Optional[tuple[Any, ...]] = None  # (min, max) of the finite ROI values

    def _default_bound(index: int, global_reduce: Any) -> Any:
        """Default bin bound (index 0: min, 1: max) over the ROI, falling back to the
        global one. The ROI range is found at the first call, so given bounds skip the
        search. A float64 row-order image with a float64, uint8 or bool row-order mask
        takes one fused pass (a bool mask as its uint8 view); other inputs gather the ROI
        values. Both skip NaN and infinite values."""
        nonlocal roi_range
        if roi_range is None and mask_arr is not None:
            if (
                array.ndim == 3
                and array.dtype == np.float64
                and mask_arr.dtype in (np.float64, np.uint8, np.bool_)
                and array.flags.c_contiguous
                and mask_arr.flags.c_contiguous
            ):
                found = roi_min_max(array, mask_arr)
                roi_range = () if found is None else (np.float64(found[0]), np.float64(found[1]))
            else:
                values = array[mask_arr != 0]
                values = values[np.isfinite(values)]
                roi_range = (values.min(), values.max()) if values.size else ()
        if roi_range:
            return roi_range[index]
        # nanmin/nanmax of an all-NaN image is NaN (warning suppressed); it
        # propagates through the bin math so every voxel ends up invalid (0).
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", RuntimeWarning)
            return global_reduce(array)

    discretised: npt.NDArray[Any]

    if array.size == 0:
        # int32 halves memory vs the default int64; bin counts in radiomics
        # never approach the int32 limit
        discretised = np.zeros(array.shape, dtype=np.int32)
        if isinstance(image, Image):
            return Image(
                array=discretised,
                spacing=image.spacing,
                origin=image.origin,
                direction=image.direction,
                modality=image.modality,
            )
        return discretised

    # The bin maths below run densely on the full array (no boolean
    # gather/scatter): NaN voxels propagate through the float ops and are
    # mapped to bin 0 (invalid) at the end. The kernels read and write a row- or
    # column-order float64 array in its own order, with no copy (another layout from
    # _KERNEL_MIN_SIZE voxels on, as a copy); the numpy chain takes the other arrays.
    order = _shared_order(array)
    use_kernel = array.dtype == np.float64 and (order is not None or array.size >= _KERNEL_MIN_SIZE)
    if method == "FBN":
        if n_bins is None:
            raise ValueError("n_bins required for FBN")
        if n_bins <= 0:
            raise ValueError("n_bins must be positive")

        # Determine min/max
        current_min = min_val if min_val is not None else _default_bound(0, np.nanmin)
        current_max = max_val if max_val is not None else _default_bound(1, np.nanmax)

        if current_max <= current_min:
            # Edge case: flat region or invalid range
            discretised = (~np.isnan(array)).astype(np.int32)
        elif use_kernel:
            # Single-pass kernel; bit-identical to the numpy chain below
            flat = array.ravel(order or "C")
            binned: npt.NDArray[Any] = np.empty(flat.size, dtype=np.int32)
            fbn = sized(
                _discretise_fbn_numba,
                _discretise_fbn_numba_serial,
                flat.size,
                _DISCRETISE_PARALLEL_MIN,
            )
            fbn(flat, float(n_bins), float(current_min), float(current_max), binned)
            discretised = binned.reshape(array.shape, order=order or "C")
        else:
            # IBSI FBN: floor(N_g * (X - X_min) / (X_max - X_min)) + 1
            temp = n_bins * (array - current_min)
            if temp.dtype.kind == "f":
                np.divide(temp, current_max - current_min, out=temp)
            else:
                # integer input: out-of-place division promotes to float
                temp = temp / (current_max - current_min)
            np.floor(temp, out=temp)
            temp += 1

            # Clip to [1, N_g]: the max value falls into N_g + 1 with this
            # formula, and outliers are clipped to the boundary bins
            np.clip(temp, 1, n_bins, out=temp)

            # NaN voxels propagated through the maths; map them to bin 0
            temp[np.isnan(temp)] = 0
            discretised = temp.astype(np.int32)

    elif method == "FBS":
        if bin_width is None:
            raise ValueError("bin_width required for FBS")
        if bin_width <= 0:
            raise ValueError("bin_width must be positive")

        current_min = min_val if min_val is not None else _default_bound(0, np.nanmin)

        if use_kernel:
            # Single-pass kernel; bit-identical to the numpy chain below
            flat = array.ravel(order or "C")
            fbs_binned: npt.NDArray[Any] = np.empty(flat.size, dtype=np.int32)
            fbs = sized(
                _discretise_fbs_numba,
                _discretise_fbs_numba_serial,
                flat.size,
                _DISCRETISE_PARALLEL_MIN,
            )
            fbs(flat, float(bin_width), float(current_min), fbs_binned)
            discretised = fbs_binned.reshape(array.shape, order=order or "C")
        else:
            # IBSI FBS: floor((X - X_min) / w_b) + 1
            temp = array - current_min
            if temp.dtype.kind == "f":
                np.divide(temp, bin_width, out=temp)
            else:
                # integer input: out-of-place division promotes to float
                temp = temp / bin_width
            np.floor(temp, out=temp)
            temp += 1

            # Ensure minimum bin is 1 (NaN propagates through np.maximum)
            np.maximum(temp, 1, out=temp)

            # NaN and +inf voxels propagated through the maths; map them to bin 0
            temp[~np.isfinite(temp)] = 0
            discretised = temp.astype(np.int32)

    elif method == "FIXED_CUTOFFS":
        if cutoffs is None:
            raise ValueError("cutoffs required for FIXED_CUTOFFS")

        # +1 so bins are 1-based: values below the first cutoff map to bin 1,
        # keeping 0 reserved for invalid (NaN) voxels.
        edges = np.asarray(cutoffs, dtype=np.float64)
        steps = np.diff(edges)
        increasing = bool(np.all(steps >= 0))
        if use_kernel and (increasing or bool(np.all(steps <= 0))):
            # The bins of np.digitize in one pass, with no int64 array and no copy
            flat = array.ravel(order or "C")
            cut_binned: npt.NDArray[Any] = np.empty(flat.size, dtype=np.int32)
            cut = sized(
                _discretise_cutoffs_numba,
                _discretise_cutoffs_numba_serial,
                flat.size,
                _DISCRETISE_PARALLEL_MIN,
            )
            cut(flat, edges, increasing, cut_binned)
            discretised = cut_binned.reshape(array.shape, order=order or "C")
        else:  # other arrays, and cutoffs that are not monotonic (digitize raises)
            temp_int = np.digitize(array, bins=np.array(cutoffs)) + 1
            temp_int[np.isnan(array)] = 0
            discretised = temp_int.astype(np.int32)

    else:
        raise ValueError(f"Unknown discretisation method: {method}")

    if isinstance(image, Image):
        return Image(
            array=discretised,
            spacing=image.spacing,
            origin=image.origin,
            direction=image.direction,
            modality=image.modality,
        )
    return discretised

def apply_mask(
    image: Image | npt.NDArray[Any],
    mask: Image | npt.NDArray[Any],
    mask_values: int | list[int] | None = None,
) -> npt.NDArray[np.floating[Any]]:
    """
    Apply mask to image and return flattened array of voxel values.

    A large 3D array (from 2^20 voxels, row order, a float64 or int32 image, a uint8, bool
    or float64 mask) is gathered in parallel when 4 or more threads are set. The values and
    their order are the same.

    Args:
        image: Image object or numpy array.
        mask: Image object (mask) or numpy array.
        mask_values: Optional value(s) in the mask to consider as ROI. When
                     omitted, all nonzero mask values are considered ROI
                     membership. Can be a single integer or a list of integers.

    Returns:
        1D numpy array of values within the mask.

    Raises:
        ValueError: If the image and mask arrays have different shapes.

    Example:
        Extract intensities within a spherical mask:

        ```python
        import numpy as np
        from pictologics.preprocessing import apply_mask

        image = np.arange(27.0).reshape(3, 3, 3)
        mask = np.zeros((3, 3, 3), dtype=np.uint8)
        mask[1, 1, 1] = 1

        values = apply_mask(image, mask)
        # array([13.])
        ```
    """
    # Handle inputs
    img_arr = image.array if isinstance(image, Image) else image
    mask_arr = mask.array if isinstance(mask, Image) else mask

    if isinstance(image, Image) and isinstance(mask, Image):
        _validate_geometry(mask, image, "mask", "image")

    # Ensure shapes match
    if img_arr.shape != mask_arr.shape:
        raise ValueError(
            f"Image shape {img_arr.shape} and mask shape {mask_arr.shape} do not match"
        )

    # Handle mask values
    if mask_values is None:
        # A large array takes the parallel gather where it can
        values = (
            None if img_arr.size < _KERNEL_MIN_SIZE else _threaded_roi_values(img_arr, mask_arr)
        )
        if values is not None:
            return values
        roi_mask = mask_arr != 0
    else:
        if isinstance(mask_values, int):
            mask_values = [mask_values]
        roi_mask = np.isin(mask_arr, mask_values)

    # Apply mask (an all-False mask yields an empty array of the image dtype)
    return cast(npt.NDArray[np.floating[Any]], img_arr[roi_mask])

def resegment_mask(
    image: Image,
    mask: Image,
    range_min: Optional[float] = None,
    range_max: Optional[float] = None,
) -> Image:
    """
    Update mask to exclude voxels where image intensity is outside the specified range.
    Used for IBSI re-segmentation (e.g. [-1000, 400] HU). A NaN intensity is in no
    range, so its voxel always leaves the mask.

    Args:
        image: Image object.
        mask: Image object (mask).
        range_min: Minimum intensity value (inclusive). If None, no lower bound.
        range_max: Maximum intensity value (inclusive). If None, no upper bound.

    Returns:
        Updated Image object (mask) with re-segmentation applied.

    Example:
        Resegment mask to keep only values between -1000 and 400 (e.g. HU range):

        ```python
        from pictologics.preprocessing import resegment_mask

        # Keep voxels in range [-1000, 400]
        new_mask = resegment_mask(
            image,
            mask,
            range_min=-1000,
            range_max=400
        )
        ```
    """
    try:
        _validate_geometry(mask, image, "mask", "image")
    except ValueError as exc:
        if "Dimension mismatch" in str(exc):
            raise ValueError(
                "Image and mask must have the same shape for re-segmentation."
            ) from exc
        raise

    order = _shared_order(image.array, mask.array)
    # The kernel reads and writes the mask as the integer type of its size (the same
    # bits), so a float mask keeps its NaN, infinite and negative values, and 0 is +0.0
    bits = _SAME_SIZE.get(mask.array.dtype.itemsize) if mask.array.dtype.kind in "biuf" else None
    if (
        image.array.dtype == np.float64
        and bits is not None
        and (order is not None or image.array.size >= _KERNEL_MIN_SIZE)
    ):
        # Single-pass kernel; bit-identical to the numpy chain below. It keeps the
        # memory order that the image and the mask share.
        lo = float("-inf") if range_min is None else float(range_min)
        hi = float("inf") if range_max is None else float(range_max)
        flat_out = np.empty(mask.array.size, dtype=mask.array.dtype)
        resegment = sized(
            _resegment_numba, _resegment_numba_serial, flat_out.size, _RESEGMENT_PARALLEL_MIN
        )
        resegment(
            image.array.ravel(order or "C"),
            mask.array.ravel(order or "C").view(bits),
            lo,
            hi,
            flat_out.view(bits),
        )
        new_mask_array = flat_out.reshape(mask.array.shape, order=order or "C")
    else:
        new_mask_array = mask.array.copy()

        # Keep the voxels in [range_min, range_max]; NaN compares False on both bounds
        lo = -np.inf if range_min is None else range_min
        hi = np.inf if range_max is None else range_max
        new_mask_array[~((image.array >= lo) & (image.array <= hi))] = 0

    return Image(
        array=new_mask_array,
        spacing=mask.spacing,
        origin=mask.origin,
        direction=mask.direction,
        modality=mask.modality,
    )

def keep_largest_component(mask: Image) -> Image:
    """
    Keep only the largest connected component in the mask.

    Uses full (26-)connectivity in 3D, i.e. voxels touching at a face, edge, or
    vertex belong to the same component. This matches the connectivity IBSI
    prescribes for 3D zones (GLSZM/GLDZM) and the convention used elsewhere in
    this library, rather than scipy's face-only (6-connectivity) default. A mask
    that fills its bounding box is one component, so the function does not label it.
    """
    mask_array = mask.array
    new_array: npt.NDArray[Any] = np.zeros(mask_array.shape, dtype=np.uint8)

    # Label only inside the nonzero bounding box: for a compact ROI in a
    # CT-sized volume, one fast bbox scan plus a small labeling costs much less
    # than a full-volume connected-component pass. Cropping removes whole
    # all-zero slabs, so components and their raster-scan label order (which
    # breaks size ties) are unchanged.
    bbox: Optional[tuple[slice, ...]]
    if mask_array.ndim == 3:
        bbox = compute_nonzero_bbox(mask_array)
    else:
        bbox = tuple(slice(None) for _ in mask_array.shape)

    if bbox is not None:
        inside = mask_array[bbox] != 0
        if inside.all():  # A full box is one component: labeling is not necessary.
            new_array[bbox] = 1
        else:
            # Full connectivity: a block of 3 voxels on each axis, all True, connects
            # all diagonal neighbours.
            structure = np.ones((3,) * inside.ndim, dtype=bool)
            labeled_sub, num_features = label(inside, structure=structure)
            if num_features <= 1:
                new_array[bbox] = inside
            else:
                del inside  # Release it before the counts
                if labeled_sub.ndim == 3:  # no intp copy of the labels, in threads
                    parts = max(1, min(64, _LABEL_COUNT_CELLS // (num_features + 1)))
                    counts = _label_counts_numba(labeled_sub, num_features, parts)
                else:
                    counts = np.bincount(labeled_sub.ravel())
                counts[0] = 0  # ignore background
                max_label = int(counts.argmax())
                np.equal(labeled_sub, max_label, out=new_array[bbox])  # no bool temporary

    return Image(
        array=new_array,
        spacing=mask.spacing,
        origin=mask.origin,
        direction=mask.direction,
        modality=mask.modality,
    )

def grow_mask(mask: Image, to_mm: float, from_mm: Optional[float] = None) -> Image:
    """
    Grow or shrink a mask by a distance in mm, or keep a ring at its edge.

    Outside the mask, the distance of a voxel is the distance from its center to the
    nearest center of a mask voxel. Inside the mask, it is the distance to the nearest
    center of a voxel outside the mask (the voxels past the image edge count as outside).
    The distances are in mm, with the spacing of each axis, so the mask changes by the
    same amount in every direction, straight out from its surface.

    A positive `to_mm` adds the voxels up to `to_mm` from the mask. A negative `to_mm`
    removes the mask voxels up to -`to_mm` from the outside. With `from_mm`, the result
    leaves out the mask grown (or shrunk) by `from_mm`, which leaves a ring. On a grid, the
    mask changes by whole voxels: a grow by 1 mm on 0.4 mm voxels adds 2 voxels along an
    axis.

    Args:
        mask: A 3D mask. Every voxel that is not 0 is in the mask.
        to_mm: The distance (mm) to grow the mask by. A negative distance shrinks it.
        from_mm: None, or a distance (mm) below `to_mm`: the result leaves out the mask
            grown (or shrunk, when negative) by this distance.

    Returns:
        The new mask (uint8: 1 in the mask, 0 outside), with the geometry of `mask`.

    Raises:
        ValueError: If a distance is not a finite number, or `from_mm` is not below
            `to_mm`.

    Example:
        ```python
        from pictologics.preprocessing import grow_mask

        ring = grow_mask(mask, to_mm=3, from_mm=0)  # the 3 mm around the mask
        core = grow_mask(mask, to_mm=-1)  # the mask without its outer 1 mm
        rim = grow_mask(mask, to_mm=0, from_mm=-1)  # the outer 1 mm of the mask
        ```
    """
    problem = _grow_problem(to_mm, from_mm)
    if problem:
        raise ValueError(problem)
    array = mask.array
    new_array: npt.NDArray[Any] = np.zeros(array.shape, dtype=np.uint8)
    bbox = compute_nonzero_bbox(array)
    if bbox is not None:
        spacing = tuple(float(s) for s in mask.spacing)
        reach = max(float(to_mm), 0.0) + _GROW_TOLERANCE_MM
        box = tuple(
            slice(max(b.start - int(reach // s), 0), min(b.stop + int(reach // s), n))
            for b, s, n in zip(bbox, spacing, array.shape, strict=True)
        )
        # One voxel more on each side, outside the mask (also past the image edge)
        inside = np.pad(array[box] != 0, 1)
        distances: dict[bool, npt.NDArray[np.float64]] = {}

        def grown(x: float) -> npt.NDArray[np.bool_]:
            """The voxels of the mask grown by x mm (shrunk by -x mm when x < 0)."""
            if x == 0:
                return inside
            outside = x > 0
            if outside not in distances:  # to the nearest voxel in (out of) the mask
                squares = _squared_distances(inside if outside else ~inside, spacing)
                distances[outside] = np.sqrt(squares, out=squares)
            if outside:
                return distances[outside] <= x + _GROW_TOLERANCE_MM
            return distances[outside] > _GROW_TOLERANCE_MM - x

        result = grown(float(to_mm))
        if from_mm is not None:
            result = result & ~grown(float(from_mm))
        new_array[box] = result[1:-1, 1:-1, 1:-1]
    return Image(
        array=new_array,
        spacing=mask.spacing,
        origin=mask.origin,
        direction=mask.direction,
        modality=mask.modality,
    )