Loaders API

def load_image(
    path: str | Path | Sequence[str | Path] | Any,
    dataset_index: int = 0,
    recursive: bool = False,
    reference_image: Optional[Image] = None,
    transpose_axes: tuple[int, int, int] | None = None,
    fill_value: float = 0.0,
    apply_rescale: bool = True,
    subvoxel_tolerance: float = 0.5,
    subvoxel_warning_threshold: float = 0.01,
    min_overlap_fraction: float = 0.5,
    series_uid: Optional[str] = None,
    suv: Optional[str] = None,
    stored_type: bool = False,
) -> Image:
    """
    Load a medical image from a file path or directory.

    This is the main entry point for loading data. It automatically detects whether
    the input is a NIfTI file, DICOM directory/file (single DICOM or series), or
    a DICOM Segmentation (SEG) object and standardizes it into an `Image` object.

    The resulting image array is always 3D with dimensions (x, y, z).

    Note:
        For DICOM SEG files, this function uses ``pictologics.loaders.load_seg()``
        internally. For more control over segment extraction (e.g., selecting specific
        segments or extracting them separately), use ``load_seg()`` directly.
        ``dataset_index`` and ``fill_value`` do not apply to SEG files and are
        ignored with a ``UserWarning`` if set to non-default values.

    Note:
        Every format gives its geometry in the LPS+ world frame (the NIfTI loader
        converts the RAS+ affine), so an image and its masks can come from different
        formats. See the module docstring ("World Coordinate Frame").

    Args:
        path (str | Path | list | DicomPhaseInfo): The absolute or relative path to the
            image file (e.g., .nii.gz, .dcm or file with no extension) or the directory
            containing DICOM files. It can also be the DICOM files of one series: a list
            of file paths, or a ``DicomPhaseInfo`` from ``get_dicom_phases()``, which
            loads that phase without reading the other files of the folder again.
        dataset_index (int, optional): For multi-volume datasets, specifies which
            volume to extract (0-indexed). This works for:

            - **4D NIfTI, NRRD and MetaImage files**: Selects which time point/volume to
              load (for a 3D Slicer ``.seg.nrrd`` file: the layer of segments).
            - **Multi-phase DICOM series**: Selects which phase to load (e.g., cardiac
              phases, temporal positions, echo numbers). Use
              ``pictologics.utilities.get_dicom_phases()`` to discover available phases.
            - **Enhanced multiframe DICOM files**: Selects the phase when the frames
              hold more than one volume (frames at repeated positions).

            Defaults to 0 (the first volume/phase).
        recursive (bool, optional): If True and `path` is a directory, recursively searches
            subdirectories and loads the DICOM series from the folder containing the most
            DICOM files. Defaults to False.
        reference_image (Optional[Image]): If provided and the loaded image has different
            dimensions than the reference, it will be repositioned into the reference
            coordinate space using spatial metadata (origin, spacing). This is useful for
            loading cropped segmentation masks that need to match a full-sized image. A
            loaded image whose axes are the reference axes in another order or with other
            signs (the same grid in another voxel order) is first turned to the voxel
            order of the reference.
        transpose_axes (tuple[int, int, int] | None): Optional axis transposition to apply
            before repositioning. Use this if the mask's axis order differs from the reference.
            E.g., (0, 2, 1) swaps Y and Z axes. Only used when reference_image is provided;
            when set, repositioning is performed even if the loaded shape already matches
            the reference.
        fill_value (float): Fill value for regions outside the loaded image when
            repositioning (default: 0.0). Only used when reference_image is provided.
        apply_rescale (bool): If True (default), apply RescaleSlope and RescaleIntercept
            transformation for DICOM files to convert stored pixel values to real-world
            values (e.g., Hounsfield Units for CT). The DICOM image is then float64, also
            without rescale tags, like a NIfTI image (nibabel's get_fdata()). Set to False
            if you need raw stored values.
        subvoxel_tolerance (float): Maximum permitted fractional-voxel offset when
            repositioning (default: 0.5). Only used when reference_image is provided.
            See ``_position_in_reference`` for full description.
        subvoxel_warning_threshold (float): Fractional-voxel drift above which a
            ``UserWarning`` is emitted during repositioning (default: 0.01). Only used
            when reference_image is provided.
        min_overlap_fraction (float): Minimum fraction of the mask volume that must
            intersect with the reference image space (default: 0.5). Only used when
            reference_image is provided.
        series_uid (str | None): The SeriesInstanceUID of the series to load when a DICOM
            folder holds more than one image series (for example two reconstructions of
            one scan). Without it, such a folder raises an error that lists the series.
            Only used for DICOM folders.
        suv (str | None): Convert a DICOM PET image (Modality PT) to its standardized
            uptake value: ``"bw"`` (body weight, g/ml), ``"lbm"`` (lean body mass by the
            Janmahasatian formula, g/ml), ``"lbm_james"`` (lean body mass by the James
            formula, g/ml) or ``"bsa"`` (body surface area by the Du Bois formula,
            cm2/ml). The factor follows the QIBA vendor-neutral pseudo-code: it needs
            attenuation and decay corrected images, and the dose, its half-life and the
            injection and scan times from the header; the lean body mass and the body
            surface area also need the height (and the lean body mass the sex). None
            (default): no conversion.
        stored_type (bool): Keep the stored type of a NIfTI file that has an integer type
            (int8, uint8, int16, uint16, int32, uint32), native byte order and no scaling
            (slope 1, intercept 0), for a mask or a label map: the array is then that type
            instead of float64. Any other NIfTI file, and every other format, loads as
            without the keyword. The pipeline loads its masks and label maps with it.
            Default False.

    Returns:
        Image: An `Image` object containing the 3D numpy array and metadata (spacing, origin, etc.).

    Raises:
        ValueError: If the path does not exist, the file format is not supported,
            the file is corrupt/unreadable, a DICOM folder holds more than one
            image series and ``series_uid`` does not choose one, or ``suv`` is set
            for an image that is not DICOM PET or lacks an attribute that the SUV needs.

    Example:
        **Loading a NIfTI file:**
        ```python
        from pictologics.loader import load_image

        # Load a standard brain scan
        img = load_image("data/brain.nii.gz")
        print(f"Image shape: {img.array.shape}")
        # Output: Image shape: (256, 256, 128)
        ```

        **Loading a DICOM series:**
        ```python
        # Load a CT scan from a folder of DICOM files
        img_ct = load_image("data/patients/001/CT_scan/")
        print(f"Voxel spacing: {img_ct.spacing}")
        # Output: Voxel spacing: (0.97, 0.97, 2.5)
        ```

        **Loading a single DICOM file:**
        ```python
        # Load a single DICOM file (even without .dcm extension)
        img_slice = load_image("data/slice_001")
        print(f"Modality: {img_slice.modality}")
        ```

        **Recursive DICOM loading:**
        ```python
        # Finds the deep subfolder with actual DICOM files
        img = load_image("data/patients/001/", recursive=True)
        ```

        **Loading a specific volume from a 4D file:**
        ```python
        # Load the 5th time point from a 4D fMRI file
        fmri_vol = load_image("data/fmri.nii.gz", dataset_index=4)
        ```

        **Loading a cropped mask and repositioning to match main image:**
        ```python
        main_img = load_image("ct_scan/")
        mask = load_image("cropped_mask.dcm", reference_image=main_img)
        # mask now has same shape as main_img
        ```

        **Loading a DICOM SEG file (auto-detected):**
        ```python
        # DICOM SEG files are automatically detected and loaded
        seg = load_image("segmentation.dcm")
        print(f"Modality: {seg.modality}")  # Output: Modality: SEG
        # Segments are combined into a label image by default
        ```

        **Loading a specific phase from a multi-phase DICOM series:**
        ```python
        from pictologics.utilities import get_dicom_phases

        # Discover available phases
        phases = get_dicom_phases("cardiac_ct/")
        print(f"Found {len(phases)} phases")
        for p in phases:
            print(f"  {p.index}: {p.label} ({p.num_slices} slices)")

        # Load the 5th phase (40%)
        img = load_image("cardiac_ct/", dataset_index=4)
        ```

        **Loading a PET series as SUV (body weight):**
        ```python
        pet = load_image("pet_series/", suv="bw")
        ```
    """
    if suv is not None and not apply_rescale:
        raise ValueError("suv needs apply_rescale=True (the activity concentrations).")
    phase_files = getattr(path, "file_paths", None)  # a DicomPhaseInfo
    if phase_files is not None or isinstance(path, (list, tuple)):
        loaded_image = _load_dicom_series(
            phase_files if phase_files is not None else path,
            dataset_index,
            apply_rescale,
            series_uid,
            suv,
        )
        if reference_image is not None:
            _validate_geometry(loaded_image, reference_image, "loaded image", "reference image")
        return loaded_image

    path = str(path)
    path_obj = Path(path)
    if not path_obj.exists():
        raise ValueError(f"The specified path does not exist: {path}")

    try:
        if path_obj.is_dir():
            target_path = path_obj
            if recursive:
                target_path = _find_best_dicom_series_dir(path_obj)
            loaded_image = _load_dicom_series(
                target_path, dataset_index, apply_rescale, series_uid, suv
            )
        elif suv is not None and path.lower().endswith(
            (".nii", ".nii.gz", ".nrrd", ".nhdr", ".mha", ".mhd")
        ):
            raise ValueError(f"suv needs a DICOM PET image, not the file '{path}'.")
        elif path.lower().endswith((".nii", ".nii.gz")):
            loaded_image = _load_nifti(path, dataset_index, stored_type)
        elif path.lower().endswith((".nrrd", ".nhdr")):
            from pictologics.loaders.nrrd_loader import _load_nrrd

            loaded_image = _load_nrrd(path, dataset_index)
        elif path.lower().endswith((".mha", ".mhd")):
            from pictologics.loaders.metaimage_loader import _load_metaimage

            loaded_image = _load_metaimage(path, dataset_index)
        else:
            # Attempt to load as a single DICOM file if extension is not NIfTI
            # Check if it's a DICOM SEG or RTSTRUCT file first
            sop_class = _dicom_sop_class(path)  # one header read for the two tests
            if suv is not None and sop_class in (*_SEG_SOP_CLASSES, _RTSTRUCT_SOP_CLASS):
                raise ValueError(f"suv needs a DICOM PET image, not the mask file '{path}'.")
            if sop_class in _SEG_SOP_CLASSES:
                from pictologics.loaders.seg_loader import load_seg

                if dataset_index != 0 or fill_value != 0.0:
                    warnings.warn(
                        "dataset_index and fill_value are ignored for DICOM SEG files. "
                        "Use pictologics.loaders.load_seg() directly for segment selection.",
                        UserWarning,
                        stacklevel=2,
                    )
                seg_result = load_seg(
                    path,
                    reference_image=reference_image,
                    transpose_axes=transpose_axes,
                    subvoxel_tolerance=subvoxel_tolerance,
                    subvoxel_warning_threshold=subvoxel_warning_threshold,
                    min_overlap_fraction=min_overlap_fraction,
                )
                # load_seg gives one label image with its default combine_segments; a dict
                # of masks gives its first mask. load_seg places the mask on the reference.
                if isinstance(seg_result, dict):
                    return next(iter(seg_result.values()))
                return seg_result
            if sop_class == _RTSTRUCT_SOP_CLASS:
                from pictologics.loaders.rtstruct_loader import load_rtstruct

                if reference_image is None:
                    raise ValueError(
                        f"'{path}' is an RTSTRUCT file: its contours need reference_image "
                        "(the image grid to fill)."
                    )
                return cast(
                    Image,
                    load_rtstruct(
                        path,
                        reference_image,
                        subvoxel_tolerance=subvoxel_tolerance,
                        subvoxel_warning_threshold=subvoxel_warning_threshold,
                    ),
                )

            try:
                loaded_image = _load_dicom_file(path, apply_rescale, dataset_index, suv)
            except _DicomContentError:
                raise
            except Exception as read_error:
                raise ValueError(
                    f"Unsupported file format or unable to read file: {path}. Supported: "
                    "NIfTI (.nii, .nii.gz), NRRD (.nrrd, .nhdr, .seg.nrrd), MetaImage (.mha, "
                    ".mhd), and DICOM files and folders (also SEG, and RTSTRUCT with "
                    "reference_image)."
                ) from read_error
    except Exception as e:
        # Re-raise ValueErrors directly, wrap others
        if isinstance(e, ValueError):
            raise e
        raise ValueError(f"Failed to load image from '{path}': {e}") from e

    # Apply repositioning if reference_image is provided and shapes differ,
    # or if an explicit axis transposition was requested (a transposed mask may
    # coincidentally have the same shape as the reference, e.g. cubic volumes).
    if reference_image is not None:
        if transpose_axes is None:
            loaded_image = _reoriented(loaded_image, reference_image)
        # A mask of the shape of the image with another origin (by whole voxels, or by a
        # part of a voxel) is placed as a mask of another shape is
        same_origin = np.allclose(loaded_image.origin, reference_image.origin, atol=1e-5, rtol=1e-5)
        if (
            loaded_image.array.shape != reference_image.array.shape
            or transpose_axes is not None
            or not same_origin
        ):
            loaded_image = _position_in_reference(
                loaded_image,
                reference_image,
                fill_value,
                transpose_axes,
                subvoxel_tolerance,
                subvoxel_warning_threshold,
                min_overlap_fraction,
            )
        else:
            _warn_if_other_frame(loaded_image.frame_of_reference_uid, reference_image, "The image")
            _validate_geometry(loaded_image, reference_image, "loaded image", "reference image")

    return loaded_image

def save_image(image: Image, path: str | Path) -> None:
    """Save an image, a mask or a response map as a NIfTI file (.nii or .nii.gz).

    The geometry goes back to the RAS+ affine of NIfTI (the X and Y rows change sign),
    so ``load_image`` reads the same array and geometry. NIfTI keeps the geometry in
    float32. A bool array is saved as uint8, and a 64-bit integer array as int32 when its
    values fit (few programs read 64-bit NIfTI integers); other arrays keep their type
    (float64 images, float32 response maps, uint8 masks).

    Args:
        image: The Image to save.
        path: The file path, ending with .nii or .nii.gz. A missing folder is made.

    Raises:
        ValueError: If the path does not end with .nii or .nii.gz.

    Example:
        ```python
        from pictologics import load_image, load_rtstruct, save_image

        ct = load_image("ct_scan/")
        masks = load_rtstruct("rtstruct.dcm", ct, roi_names=["GTV"], combine_rois=False)
        save_image(masks["GTV"], "gtv.nii.gz")
        ```
    """
    path = str(path)
    if not path.lower().endswith((".nii", ".nii.gz")):
        raise ValueError(f"save_image writes NIfTI files (.nii, .nii.gz), not '{path}'.")
    affine = np.eye(4)
    affine[:3, :3] = _direction_matrix(image.direction) * np.asarray(image.spacing, dtype=float)
    affine[:3, 3] = image.origin
    affine[:2] *= -1.0  # LPS+ to RAS+: the X and Y rows change sign
    array: npt.NDArray[Any] = (
        image.array.view(np.uint8) if image.array.dtype == np.bool_ else image.array
    )
    if array.dtype.kind in "iu" and array.dtype.itemsize == 8 and array.size:
        # nibabel refuses 64-bit integers without a dtype (few programs read them): int32
        # when the values fit
        info = np.iinfo(np.int32)
        if info.min <= array.min() and array.max() <= info.max:
            array = array.astype(np.int32)
    if (
        array.ndim == 3
        and array.flags.c_contiguous
        and array.size >= _ROW_ORDER_MIN_SIZE
        and array.dtype in _COLUMN_ORDER_DTYPES
    ):
        # nibabel writes the voxels in column order, one slice at a time. A slice of a
        # row-order array costs a cache line per voxel; a column-order copy is fast.
        column = np.empty(array.shape, dtype=array.dtype, order="F")
        _to_row_order_numba(array, column)
        array = column
    nifti = Nifti1Image(array, affine, dtype=array.dtype)  # type: ignore[no-untyped-call]
    # Code 1: scanner coordinates, as ITK writes
    nifti.set_qform(affine, code=1)  # type: ignore[no-untyped-call]
    nifti.set_sform(affine, code=1)  # type: ignore[no-untyped-call]
    nifti.header.set_xyzt_units("mm")  # type: ignore[no-untyped-call]
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    if (
        path.lower().endswith(".gz")
        and array.ndim == 3
        and array.flags.f_contiguous
        and array.size >= _ROW_ORDER_MIN_SIZE
        and array.dtype in _COLUMN_ORDER_DTYPES
    ):
        _write_gzip_nifti(nifti, array, path)
    else:
        nifti.to_filename(path)

def load_and_merge_images(
    image_paths: Sequence[str | Path],
    reference_image: Optional[Image] = None,
    conflict_resolution: str = "max",
    dataset_index: int = 0,
    recursive: bool = False,
    binarize: bool | int | list[int] | tuple[int, int] | None = None,
    reposition_to_reference: bool = False,
    transpose_axes: tuple[int, int, int] | None = None,
    fill_value: float = 0.0,
    relabel_masks: bool = False,
    apply_rescale: bool = True,
    subvoxel_tolerance: float = 0.5,
    subvoxel_warning_threshold: float = 0.01,
    min_overlap_fraction: float = 0.5,
) -> Image:
    """
    Load multiple images (e.g., masks or partial scans) and merge them into a single image.

    This function loads images from the provided paths, validates that they all share
    the same geometry (dimensions, spacing, origin, direction), and merges them
    according to the specified conflict resolution strategy.

    **Use Cases:**
    - Merging multiple segmentation masks into a single ROI.
    - Merging split image volumes (though typically less common than mask merging).
    - Merging cropped/bounding-box segmentation masks (with `reposition_to_reference=True`).

    **Format & Path Support:**
    Since this function uses `load_image` internally for each path, it supports:
    - **NIfTI files** (.nii, .nii.gz).
    - **NRRD and MetaImage files** (.nrrd, .nhdr, .mha, .mhd).
    - **DICOM series** (directories containing DICOM files).
    - **Single DICOM files** (with or without .dcm extension).
    - **Nested directories** (if paths point to folders containing DICOMs).

    Args:
        image_paths (Sequence[str | Path]): List of absolute or relative paths to the images.
            These can be file paths or directory paths.
        reference_image (Optional[Image]): An optional reference image (e.g., the scan
            corresponding to the masks). If provided, the merged image is validated
            against this image's geometry. Required when `reposition_to_reference=True`.
        conflict_resolution (str): Strategy to resolve voxel values when multiple images
            have non-zero values at the same location. Options:
            - 'max': Use the maximum value (default).
            - 'min': Use the minimum value.
            - 'first': Keep the value from the first image encountered (earlier in list).
            - 'last': Overwrite with the value from the last image encountered (later in list).
        dataset_index (int, optional): For multi-volume datasets, specifies which
            volume to extract for all images (0-indexed). This works for:

            - **4D NIfTI, NRRD and MetaImage files**: Selects which time point/volume to
              load (for a 3D Slicer ``.seg.nrrd`` file: the layer of segments).
            - **Multi-phase DICOM series**: Selects which phase to load (e.g., cardiac
              phases, temporal positions, echo numbers). Use
              ``pictologics.utilities.get_dicom_phases()`` to discover available phases.

            Defaults to 0 (the first volume/phase).
        recursive (bool, optional): If True, recursively searches subdirectories
            for each path in `image_paths`. Defaults to False.
        binarize (bool | int | list[int] | tuple[int, int] | None, optional):
            Rules for binarizing the merged image.
            - `None` (default): No binarization.
            - `True`: Sets all voxels > 0 to 1, others to 0.
            - `int` (e.g., 2): Sets voxels == value to 1, others to 0.
            - `list[int]` (e.g., [1, 2]): Sets voxels in list to 1, others to 0.
            - `tuple[int, int]` (e.g., (1, 10)): Sets voxels in inclusive range to 1, others to 0.
        reposition_to_reference (bool): If True and reference_image is provided,
            each loaded image will be repositioned into the reference coordinate
            space before merging. This is required when loading cropped segmentation
            masks that have different dimensions than the reference. Geometry validation
            is performed AFTER repositioning. Defaults to False.
        transpose_axes (tuple[int, int, int] | None): Axis transposition to apply
            when repositioning. E.g., (0, 2, 1) swaps Y and Z axes.
            Only used when `reposition_to_reference=True`.
        fill_value (float): Fill value for regions outside cropped masks when
            repositioning (default: 0.0). Only used when `reposition_to_reference=True`.
        relabel_masks (bool): If True, assigns unique label values (1, 2, 3, ...)
            to each mask file based on its order in `image_paths`. This converts
            binary [0,1] masks into multi-label masks where each file gets a
            distinct label, useful for visualization with different colors.
            Label assignment respects the order of `image_paths`. Defaults to False.
        apply_rescale (bool): If True (default), apply RescaleSlope and RescaleIntercept
            transformation for DICOM files to convert stored pixel values to real-world
            values (e.g., Hounsfield Units for CT). Set to False if you need raw stored values.
        subvoxel_tolerance (float): Maximum permitted fractional-voxel offset when
            repositioning (default: 0.5). Only used when ``reposition_to_reference=True``.
            See ``_position_in_reference`` for full description.
        subvoxel_warning_threshold (float): Fractional-voxel drift above which a
            ``UserWarning`` is emitted during repositioning (default: 0.01). Only used
            when ``reposition_to_reference=True``.
        min_overlap_fraction (float): Minimum fraction of each mask volume that must
            intersect with the reference image space (default: 0.5). Only used when
            ``reposition_to_reference=True``.

    Note:
        The `binarize` parameter is intended for **mask filtering** (e.g., selecting specific
        ROI labels). To filter image intensity values (e.g., HU ranges), use the preprocessing
        steps in the radiomics pipeline configuration instead.

    Returns:
        Image: A new `Image` object containing the merged data. Its array has the common
            type of the loaded arrays (float64 for NIfTI, NRRD and MetaImage files and
            rescaled DICOM; the stored type of SEG files and of DICOM with
            ``apply_rescale=False``). With `binarize`, it is uint8; with `relabel_masks`,
            the smallest unsigned type of the labels. With `reposition_to_reference`, the
            merge starts as `fill_value` in the type of the first image, or float64 when
            that type does not hold `fill_value` exactly.

    Raises:
        ValueError: If `image_paths` is empty, if an invalid `conflict_resolution` is provided,
            if `reposition_to_reference=True` but `reference_image` is not provided,
            or if the images (or reference) have mismatched geometries.

    Example:
        **Merging cropped segmentation masks:**
        ```python
        main_img = load_image("ct_scan/", recursive=True)
        seg_paths = [str(f) for f in Path("masks/").glob("*.dcm")]

        merged = load_and_merge_images(
            seg_paths,
            reference_image=main_img,
            reposition_to_reference=True,
            conflict_resolution="max",
        )
        ```
    """
    if not image_paths:
        raise ValueError("image_paths cannot be empty.")

    valid_strategies = {"max", "min", "first", "last"}
    if conflict_resolution not in valid_strategies:
        raise ValueError(
            f"Invalid conflict_resolution '{conflict_resolution}'. "
            f"Must be one of {valid_strategies}."
        )

    if reposition_to_reference and reference_image is None:
        raise ValueError("reference_image must be provided when reposition_to_reference=True.")

    # The labels of relabel_masks: the smallest unsigned type that holds the count of files
    label_type = np.min_scalar_type(len(image_paths))
    merged_array: Optional[npt.NDArray[Any]] = None

    if reposition_to_reference:
        # Mode: Reposition each image to reference space, then merge
        assert reference_image is not None  # Already validated above

        for i, path in enumerate(image_paths):
            try:
                current_image = load_image(
                    path,
                    dataset_index=dataset_index,
                    recursive=recursive,
                    apply_rescale=apply_rescale,
                    stored_type=relabel_masks,
                )
            except Exception as e:
                raise ValueError(f"Failed to load image '{path}': {e}") from e

            # Where the image lies in reference space; outside its box it is all fill,
            # so only the box merges (the repositioned full-size array is not made)
            data, source, target = _placement(
                current_image,
                reference_image,
                transpose_axes,
                subvoxel_tolerance,
                subvoxel_warning_threshold,
                min_overlap_fraction,
            )
            if merged_array is None:  # the merge, in reference geometry and the first type
                kind = _merge_type(label_type if relabel_masks else data.dtype, fill_value)
                merged_array = _filled(reference_image.array.shape, fill_value, kind)
            if source is None or target is None:
                continue
            current_array = data[source]

            # Merge with conflict resolution, in the common type. With relabel_masks,
            # each value other than the fill value merges as the mask index + 1 (the
            # merge already holds these labels).
            label: Any = None
            if relabel_masks:
                label = label_type.type(i + 1)
            else:
                merged_array = _widened(merged_array, current_array)
            _merge_into(merged_array[target], current_array, fill_value, conflict_resolution, label)
            del current_image, data, current_array  # free the image before the next load
        assert merged_array is not None  # image_paths is not empty

        # Use reference geometry for output
        consensus_spacing = reference_image.spacing
        consensus_origin = reference_image.origin
        consensus_direction = reference_image.direction

    else:
        # Mode: Standard merging with strict geometry validation
        # Load the first image to serve as the consensus geometry
        try:
            consensus_image = load_image(
                image_paths[0],
                dataset_index=dataset_index,
                recursive=recursive,
                apply_rescale=apply_rescale,
                stored_type=relabel_masks,
            )
        except Exception as e:
            raise ValueError(f"Failed to load first image '{image_paths[0]}': {e}") from e

        # The merge goes into the array of the first image, which this function loaded
        merged_array = consensus_image.array

        # Apply relabeling for the first image
        if relabel_masks:
            merged_array = (merged_array != 0).astype(label_type)

        # Iterate through remaining images
        for idx, path in enumerate(image_paths[1:], start=2):
            try:
                current_image = load_image(
                    path,
                    dataset_index=dataset_index,
                    recursive=recursive,
                    apply_rescale=apply_rescale,
                    stored_type=relabel_masks,
                )
            except Exception as e:
                raise ValueError(f"Failed to load image '{path}': {e}") from e

            _validate_geometry(current_image, consensus_image, f"image '{path}'", "consensus image")

            current_array = current_image.array

            # Merge with conflict resolution, in the common type. With relabel_masks,
            # each value other than 0 merges as the mask index (from 2 for the second
            # file).
            label = None
            if relabel_masks:
                label = label_type.type(idx)
            else:
                merged_array = _widened(merged_array, current_array)
            _merge_into(merged_array, current_array, 0, conflict_resolution, label)
            del current_image, current_array  # free the image before the next load

        consensus_spacing = consensus_image.spacing
        consensus_origin = consensus_image.origin
        consensus_direction = consensus_image.direction

        # Validate against reference image if provided (for non-reposition mode)
        if reference_image is not None:
            final_merged_image = Image(
                array=merged_array,
                spacing=consensus_spacing,
                origin=consensus_origin,
                direction=consensus_direction,
                modality="Image",
            )
            _validate_geometry(
                final_merged_image, reference_image, "merged image", "reference image"
            )

    # Apply binarization if requested
    if binarize is not None:
        mask_out: npt.NDArray[Any] = np.zeros_like(merged_array, dtype=np.uint8)
        if isinstance(binarize, bool) and binarize is True:
            mask_out[merged_array > 0] = 1
        elif isinstance(binarize, int) and not isinstance(binarize, bool):
            mask_out[merged_array == binarize] = 1
        elif isinstance(binarize, list):
            mask_out[np.isin(merged_array, binarize)] = 1
        elif isinstance(binarize, tuple) and len(binarize) == 2:
            mask_out[(merged_array >= binarize[0]) & (merged_array <= binarize[1])] = 1
        else:
            if binarize is not False:
                raise ValueError(f"Unsupported binarize value: {binarize}")
            mask_out = merged_array

        if binarize is not False:
            merged_array = mask_out

    return Image(
        array=merged_array,
        spacing=consensus_spacing,
        origin=consensus_origin,
        direction=consensus_direction,
        modality="MergedImage",
    )

def create_full_mask(reference_image: Image, dtype: DTypeLike = np.uint8) -> Image:
    """Create a whole-image ROI mask matching a reference image.

    This utility is primarily used when a user does not provide a segmentation mask.
    The returned mask has the same geometry (shape, spacing, origin, direction) as
    the reference image and contains a value of 1 for every voxel.

    Args:
        reference_image: Image whose geometry should be copied.
        dtype: Numpy dtype to use for the mask array. Defaults to `np.uint8`.

    Returns:
        An `Image` mask with `array == 1` everywhere.

    Raises:
        ValueError: If the reference image does not have a valid 3D array.

    Example:
        ```python
        import numpy as np
        from pictologics.loader import Image, create_full_mask

        image = Image(array=np.zeros((10, 10, 5)), spacing=(1.0, 1.0, 2.0), origin=(0.0, 0.0, 0.0))
        mask = create_full_mask(image)
        print(mask.array.shape, mask.array.dtype)
        # (10, 10, 5) uint8
        print(mask.array.min(), mask.array.max())
        # 1 1
        ```
    """
    if reference_image.array.ndim != 3:
        raise ValueError(
            f"reference_image.array must be 3D, got shape {reference_image.array.shape}"
        )

    mask_array = np.ones(reference_image.array.shape, dtype=dtype)
    return Image(
        array=mask_array,
        spacing=reference_image.spacing,
        origin=reference_image.origin,
        direction=reference_image.direction,
        modality="mask",
        frame_of_reference_uid=reference_image.frame_of_reference_uid,
    )