"""Abstract Base Class for TIAToolbox Deep Learning Engines.

This module defines the `EngineABC` class, which serves as a base for implementing
deep learning inference workflows in TIAToolbox. It supports both patch-based and
whole slide image (WSI) processing, and provides a unified interface for model
initialization, data loading, inference, post-processing, and output saving.

Classes:
    - EngineABC: Abstract base class for deep learning engines.
    - EngineABCRunParams: TypedDict for runtime configuration parameters.

Functions:
    - prepare_engines_save_dir: Utility to create or validate output directories.

Features:
    - Supports patch and WSI modes.
    - Handles caching and memory-efficient inference using Dask.
    - Integrates with TIAToolbox models and IO configurations.
    - Outputs predictions in multiple formats including dict, zarr, and AnnotationStore.

Intended Usage:
    Subclass `EngineABC` to implement specific inference logic by overriding
    abstract methods such as preprocessing, postprocessing, and model-specific behavior.

Example:
    >>> class MyEngine(EngineABC):
    >>>     def __init__(self, model, weights, verbose):
    >>>         super().__init__(model=model, weights=weights, verbose=verbose)
    >>> # Implement base class functions and then call.
    >>> engine = MyEngine(model="resnet18-kather100k")
    >>> output = engine.run(images, patch_mode=True)

"""

from __future__ import annotations

import copy
import gc
import shutil
from abc import ABC
from pathlib import Path
from typing import TYPE_CHECKING, TypedDict

import dask
import dask.array as da
import numpy as np
import torch
import zarr
from dask import compute
from numcodecs import Pickle
from torch import nn
from tqdm.auto import tqdm
from typing_extensions import Unpack

from tiatoolbox import DuplicateFilter, logger, rcParam
from tiatoolbox.models.architecture import get_pretrained_model
from tiatoolbox.models.architecture.utils import compile_model
from tiatoolbox.models.dataset.dataset_abc import PatchDataset, WSIPatchDataset
from tiatoolbox.models.models_abc import load_torch_model
from tiatoolbox.utils.misc import (
    dict_to_store_patch_predictions,
    tqdm_dask_progress_bar,
)
from tiatoolbox.wsicore.wsireader import WSIReader, is_zarr

from .io_config import ModelIOConfigABC

if TYPE_CHECKING:  # pragma: no cover
    import os
    from collections.abc import Callable

    from torch.utils.data import DataLoader

    from tiatoolbox.annotation import AnnotationStore
    from tiatoolbox.models.models_abc import ModelABC
    from tiatoolbox.type_hints import IntPair, Resolution, Units
    from tiatoolbox.wsicore import WSIReaderParams


class EngineABCRunParams(TypedDict, total=False):
    """Parameters for configuring the :func:`EngineABC.run()` method.

    Optional Keys:
        auto_get_mask (bool):
            Whether to automatically generate segmentation masks using
            `wsireader.tissue_mask()` during processing.
        batch_size (int):
            Number of image patches per forward pass.
        class_dict (dict):
            Mapping of classification outputs to class names.
        device (str):
            Device to run the model on (e.g., "cpu", "cuda").
            See https://pytorch.org/docs/stable/tensor_attributes.html#torch.device
            for more details.
        labels (list):
            Optional labels for input images. Only a single label per image
            is supported.
        memory_threshold (int):
            Memory usage threshold (in percentage) to trigger caching behavior.
        num_workers (int):
            Number of workers for DataLoader and post-processing.
        output_file (str):
            Filename for saving output (e.g., "zarr" or "annotationstore").
        return_labels (bool):
            Whether to return labels with predictions.
        scale_factor (tuple[float, float]):
            Scale factor for annotations (model_mpp / slide_mpp).
            Used to convert coordinates from non-baseline to baseline
            resolution.
        stride_shape (IntPair):
            Stride used during WSI processing, at requested read resolution.
            Must be positive. Defaults to `patch_input_shape` if not
            provided.
        wsireader_kwargs (WSIReaderParams):
            Specify processing images with no mpp or power in the metadata.
        verbose (bool):
            Whether to enable verbose logging.

    """

    auto_get_mask: bool
    batch_size: int
    class_dict: dict
    device: str
    labels: list
    memory_threshold: int
    num_workers: int
    output_file: str
    return_labels: bool
    scale_factor: tuple[float, float]
    stride_shape: IntPair
    wsireader_kwargs: dict
    verbose: bool


class EngineABC(ABC):  # noqa: B024
    """Abstract base class for TIAToolbox deep learning engines to run CNN models.

    This class provides a unified interface for running inference on image patches
    or whole slide images (WSIs), handling preprocessing, batching, postprocessing,
    and saving predictions.

    Args:
        model (str | ModelABC):
            Model name from TIAToolbox or a PyTorch model instance.
            The user can request pretrained models from the toolbox model zoo using
            the list of pretrained models available at this `link
            <https://tia-toolbox.readthedocs.io/en/latest/pretrained.html>`_
            By default, the corresponding pretrained weights will also
            be downloaded. However, you can override with your own set
            of weights using the `weights` parameter.
        batch_size (int):
            Number of patches per forward pass. Default is 8.
        num_workers (int):
            Number of workers for data loading. Default is 0.
        weights (str | Path | None):
            Path to model weights. If None, default weights are used.

            >>> engine = EngineABC(
            ...    model="pretrained-model",
            ...    weights="/path/to/pretrained-local-weights.pth"
            ... )

        device (str):
            Device to run the model on (e.g., "cpu", "cuda"). Please see
            https://pytorch.org/docs/stable/tensor_attributes.html#torch.device
            for more details on input parameters for device. Default is "cpu".
        verbose (bool):
            Enable verbose logging. Default is False.

    Attributes:
        images (list[str | Path] | np.ndarray):
            Input images or patches.
            A list of image patches in NHWC format as a numpy array
            or a list of str/paths to WSIs.
        masks (list[str | Path] | np.ndarray):
            Optional masks for WSIs.
            A list of tissue masks or binary masks corresponding to processing area of
            input images. These can be a list of numpy arrays or paths to
            the saved image masks. These are only utilized when patch_mode is False.
            Patches are only generated within a masked area.
            If not provided, then a tissue mask will be automatically
            generated for whole slide images, if auto_get_mask is True.
        patch_mode (bool):
            Whether input is treated as patches. TIAToolbox defines
            an image as a patch if HWC of the input image matches with the HWC expected
            by the model. If HWC of the input image does not match with the HWC expected
            by the model, then the patch_mode must be set to False which will allow the
            engine to extract patches from the input image.
            In this case, when the patch_mode is False the input images are treated
            as WSIs. Default value is True.
        model (ModelABC):
            Loaded PyTorch model. For a full list of pretrained models,
            refer to the `docs
            <https://tia-toolbox.readthedocs.io/en/latest/pretrained.html>`_
            By default, the corresponding pretrained weights will also
            be downloaded. However, you can override with your own set
            of weights via the `weights` argument. Argument
            is case-insensitive.
        ioconfig (ModelIOConfigABC):
            IO configuration (:class:`ModelIOConfigABC`) for model input/output.
        dataloader (DataLoader):
            Torch DataLoader for inference.
        return_labels (bool):
            Whether to return labels with probabilities and predictions.
        input_resolutions (list[dict[Units, Resolution]]):
            Resolution settings for input heads. Supported
            units are `level`, `power` and `mpp`. Keys should be "units" and
            "resolution" e.g., [{"units": "mpp", "resolution": 0.25}]. Please see
            :class:`WSIReader` for details.
        patch_input_shape (IntPair):
            Shape of input patches. Patches are at
            requested read resolution, not with respect to level 0,
            and must be positive.
        stride_shape (IntPair):
            Stride used during WSI processing. Stride is
            at requested read resolution, not with respect to
            level 0, and must be positive. If not provided,
            `stride_shape=patch_input_shape`.
        batch_size (int):
            Number of patches per forward pass.
        labels (list | None):
            Optional labels for input images. Only a single label per image is
            supported.
        mask_bounds (tuple (int, int, int, int)) | None:
            mask bounds for processing. top, left, bottom, right.
            When no mask is set, mask_bounds is None.
        num_workers (int):
            Number of workers for data loading.
        patch_input_shape (IntPair | None):
            Shape of input patches.
        input_resolutions (list[dict[Units, Resolution]] | None):
            Resolution settings for input heads.
        return_labels (bool):
            Whether to return labels with predictions.
        stride_shape (IntPair | None):
            Stride used during WSI processing.
        verbose (bool):
            Whether to enable verbose logging.
        dataloader (DataLoader | None):
            Torch DataLoader for inference.
        drop_keys (list):
            Keys to exclude from model output.
        output_type (Any):
            Format of output ("dict", "zarr", "qupath", "AnnotationStore").
        verbose (bool):
            Whether to enable verbose logging.

    Example:
        >>> # Inherit from EngineABC
        >>> class MyEngine(EngineABC):
        >>>     def __init__(self, model, weights, verbose):
        >>>         super().__init__(model=model, weights=weights, verbose=verbose)
        >>> engine = MyEngine(model="resnet18-kather100k")
        >>> output = engine.run(images, patch_mode=True)

        >>> # Define all the abstract classes

        >>> data = np.array([np.ndarray, np.ndarray])
        >>> engine = TestEngineABC(model="resnet18-kather100k")
        >>> output = engine.run(data, patch_mode=True)

        >>> # array of list of 2 image patches as input
        >>> data = np.array([np.ndarray, np.ndarray])
        >>> engine = TestEngineABC(model="resnet18-kather100k")
        >>> output = engine.run(data, patch_mode=True)

        >>> # list of 2 image files as input
        >>> image = ['path/image1.png', 'path/image2.png']
        >>> engine = TestEngineABC(model="resnet18-kather100k")
        >>> output = engine.run(image, patch_mode=False)

        >>> # list of 2 wsi files as input
        >>> wsi_file = ['path/wsi1.svs', 'path/wsi2.svs']
        >>> engine = TestEngineABC(model="resnet18-kather100k")
        >>> output = engine.run(wsi_file, patch_mode=False)

    """

    def __init__(
        self: EngineABC,
        model: str | ModelABC,
        batch_size: int = 8,
        num_workers: int = 0,
        weights: str | Path | None = None,
        *,
        device: str = "cpu",
        verbose: bool = False,
    ) -> None:
        """Initialize the EngineABC instance.

        Args:
            model (str | ModelABC):
                Model name from TIAToolbox or a PyTorch model instance.
            batch_size (int):
                Number of patches per forward pass. Default is 8.
            num_workers (int):
                Number of workers for data loading. Default is 0.
            weights (str | Path | None):
                Path to model weights. If None, default weights are used.
            device (str):
                Device to run the model on (e.g., "cpu", "cuda"). Default is "cpu".
            verbose (bool):
                Enable verbose logging. Default is False.

        """
        self.images = None
        self.masks = None
        self.patch_mode = None
        self.device = device

        # Initialize model with specified weights and ioconfig.
        self.model, self.ioconfig = self._initialize_model_ioconfig(
            model=model, weights=weights
        )
        self.model.to(device=self.device)
        self.model = (
            compile_model(  # for runtime, such as after wrapping with nn.DataParallel
                self.model,
                mode=rcParam["torch_compile_mode"],
            )
        )
        self._ioconfig = self.ioconfig  # runtime ioconfig
        self.batch_size = batch_size
        self.labels: list | None = None
        self.num_workers = num_workers
        self.patch_input_shape: IntPair | None = None
        self.input_resolutions: list[dict[Units, Resolution]] | None = None
        self.return_labels: bool = False
        self.stride_shape: IntPair | None = None
        self.verbose: bool = verbose
        self.dataloader: DataLoader | None = None
        self.drop_keys: list = []
        self.output_type: str | None = None
        self.mask_bounds: tuple[int, int, int, int] | None = None

    @staticmethod
    def _initialize_model_ioconfig(
        model: str | ModelABC,
        weights: str | Path | None,
    ) -> tuple[nn.Module, ModelIOConfigABC | None]:
        """Helper function to initialize model and IO configuration.

        If a pretrained model from TIAToolbox is specified by name, this function
        loads the model and its associated IO configuration. If a custom model is
        provided, it loads the weights if specified and returns None for IO config.

        Args:
            model (str | ModelABC):
                A model name from TIAToolbox or a PyTorch model instance.
                The user can request pretrained models from the toolbox model zoo using
                the list of pretrained models available at this `link
                <https://tia-toolbox.readthedocs.io/en/latest/pretrained.html>`_
                By default, the corresponding pretrained weights will also
                be downloaded. However, you can override with your own set
                of weights using the `weights` parameter.

            weights (str | Path | None):
                Path to pretrained weights. If None and a TIAToolbox model is used,
                default weights are automatically downloaded.

        Returns:
            tuple[nn.Module, ModelIOConfigABC | None]:
                A tuple containing the loaded PyTorch model and its IO configuration.
                If the model is not from TIAToolbox, IO config will be None.

        Raises:
            TypeError:
                If the model is neither a string (TIAToolbox model)
                nor a torch.nn.Module.

        """
        if not isinstance(model, (str, nn.Module)):
            msg = "Input model must be a string or 'torch.nn.Module'."
            raise TypeError(msg)

        if isinstance(model, str):
            # ioconfig is retrieved from the pretrained model in the toolbox.
            # list of pretrained models in the TIA Toolbox is available here:
            # https://tia-toolbox.readthedocs.io/en/latest/pretrained.html
            # no need to provide ioconfig in EngineABC.run() this case.
            return get_pretrained_model(model, weights)

        if weights is not None:
            model = load_torch_model(model=model, weights=weights)

        return model, None

    def _get_model_attr(self: EngineABC, attr_name: str) -> Callable:
        """Return a model attribute, unwrapping DataParallel if required."""
        try:
            return getattr(self.model, attr_name)
        except AttributeError:
            module = getattr(self.model, "module", None)
            return getattr(module, attr_name)

    def get_dataloader(
        self: EngineABC,
        images: str | Path | list[str | Path] | np.ndarray,
        masks: Path | None = None,
        labels: list | None = None,
        ioconfig: ModelIOConfigABC | None = None,
        *,
        patch_mode: bool = True,
        auto_get_mask: bool = True,
        wsireader_kwargs: WSIReaderParams | None = None,
    ) -> torch.utils.data.DataLoader:
        """Pre-process images and masks and return a DataLoader for inference.

        Args:
            images (list[str | Path] | np.ndarray):
                A list of image patches in NHWC format as a numpy array,
                or a list of file paths to WSIs. When `patch_mode` is False,
                expects file paths to WSIs.
            masks (Path | None):
                Optional list of masks used when `patch_mode` is False.
                Patches are generated only within masked areas. If not provided,
                tissue masks are automatically generated.
            labels (list | None):
                Optional list of labels. Only one label per image is supported.
            ioconfig (ModelIOConfigABC | None):
                IO configuration object specifying patch size, stride, and resolution.
            patch_mode (bool):
                Whether to treat input as patches (`True`) or WSIs (`False`).
            wsireader_kwargs (WSIReaderParams):
                Specify processing images with no mpp or power in the metadata.
            auto_get_mask (bool):
                Whether to automatically generate a tissue mask using
                `wsireader.tissue_mask()` when `patch_mode` is False.
                If `True`, only tissue regions are processed. If `False`,
                all patches are processed. Default is `True`.

        Returns:
            torch.utils.data.DataLoader:
                A PyTorch DataLoader configured for inference.

        """
        if labels:
            # if a labels is provided, then return with the prediction
            self.return_labels = bool(labels)

        if not patch_mode:
            dataset = WSIPatchDataset(
                input_img=images,
                input_mask=masks,
                patch_input_shape=ioconfig.patch_input_shape,
                stride_shape=ioconfig.stride_shape,
                resolution=ioconfig.input_resolutions[0]["resolution"],
                units=ioconfig.input_resolutions[0]["units"],
                auto_get_mask=auto_get_mask,
                wsireader_kwargs=wsireader_kwargs,
            )

            dataset.preproc_func = self._get_model_attr("preproc_func")

            # preprocessing must be defined with the dataset
            return torch.utils.data.DataLoader(
                dataset,
                num_workers=self.num_workers,
                batch_size=self.batch_size,
                drop_last=False,
                shuffle=False,
                persistent_workers=self.num_workers > 0,
            )

        dataset = PatchDataset(
            inputs=images, labels=labels, patch_input_shape=ioconfig.patch_input_shape
        )

        dataset.preproc_func = self._get_model_attr("preproc_func")

        # preprocessing must be defined with the dataset
        return torch.utils.data.DataLoader(
            dataset,
            num_workers=self.num_workers,
            batch_size=self.batch_size,
            drop_last=False,
            shuffle=False,
            persistent_workers=self.num_workers > 0,
        )

    def _get_coordinates(self: EngineABC, batch_data: dict) -> np.ndarray:
        """Extract coordinates for each image patch in a batch.

        This method returns coordinates for each patch, either based on
        the patch dimensions (if in patch mode) or from precomputed values
        (if in WSI mode).

        Args:
            batch_data (dict):
                Dictionary containing batch data, including image and
                optional coordinates.

        Returns:
            np.ndarray:
                Array of coordinates for each patch in the batch.
                Shape is (N, 4), where N is the number of patches.

        """
        if self.patch_mode:
            coordinates = [0, 0, *batch_data["image"].shape[1:3]]
            return np.tile(coordinates, reps=(batch_data["image"].shape[0], 1))
        return np.array(batch_data["coords"])

    def infer_patches(
        self: EngineABC,
        dataloader: DataLoader,
        *,
        return_coordinates: bool = False,
    ) -> dict[str, da.Array]:
        """Run model inference on image patches and return predictions.

        This method performs batched inference using a PyTorch DataLoader,
        and accumulates predictions in Dask arrays. It supports optional inclusion
        of coordinates and labels in the output.

        Args:
            dataloader (DataLoader):
                PyTorch DataLoader containing image patches for inference.
            return_coordinates (bool):
                Whether to include coordinates in the output. Required when
                called by `infer_wsi` and `patch_mode` is False.

        Returns:
            dict[str, dask.array.Array]:
                Dictionary containing prediction results as Dask arrays.
                Keys include:
                    - "probabilities": Model output probabilities.
                    - "labels": Ground truth labels (if `return_labels` is True).
                    - "coordinates": Patch coordinates (if `return_coordinates` is
                      True).

        """
        keys = ["probabilities"]
        probabilities, labels, coordinates = [], [], []

        if self.return_labels:
            keys.append("labels")
            labels = []

        if return_coordinates:
            keys.append("coordinates")
            coordinates = []

        # Main output dictionary
        raw_predictions = {key: [] for key in keys}

        # Inference loop
        tqdm_loop = tqdm(
            dataloader,
            leave=False,
            desc="Inferring patches",
            disable=not self.verbose,
        )

        infer_batch = self._get_model_attr("infer_batch")
        for batch_data in tqdm_loop:
            batch_output = infer_batch(
                self.model,
                batch_data["image"],
                device=self.device,
            )

            probabilities.append(
                da.from_array(
                    batch_output,  # probabilities
                )
            )

            if return_coordinates:
                coordinates.append(
                    da.from_array(
                        self._get_coordinates(batch_data),
                    )
                )

            if self.return_labels:
                labels.append(da.from_array(np.array(batch_data["label"])))

        raw_predictions["probabilities"] = da.concatenate(probabilities, axis=0)

        if return_coordinates:
            raw_predictions["coordinates"] = da.concatenate(coordinates, axis=0)

        if self.return_labels:
            labels = [label.reshape(-1) for label in labels]
            raw_predictions["labels"] = da.concatenate(labels, axis=0)

        return raw_predictions

    def post_process_patches(  # skipcq: PYL-R0201
        self: EngineABC,
        raw_predictions: dict[str, da.Array],
        **kwargs: Unpack[EngineABCRunParams],  # noqa: ARG002
    ) -> dict[str, da.Array]:
        """Post-process raw patch predictions from inference.

        This method applies a post-processing function (e.g., smoothing, filtering)
        to the raw model predictions. It supports delayed execution using Dask
        and returns a Dask array for efficient computation.

        Args:
            raw_predictions (dict[str, da.Array]):
                Dictionary containing raw model predictions as Dask arrays.
            **kwargs (EngineABCRunParams):
                Additional runtime parameters to update engine attributes.

                Optional Keys:
                    auto_get_mask (bool):
                        Automatically generate segmentation masks using
                        `wsireader.tissue_mask()` during processing.
                    batch_size (int):
                        Number of image patches per forward pass.
                    class_dict (dict):
                        Mapping of classification outputs to class names.
                    device (str):
                        Device to run the model on (e.g., "cpu", "cuda").
                        See https://pytorch.org/docs/stable/tensor_attributes.html#torch.device
                        for more details.
                    labels (list):
                        Optional labels for input images. Only a single label per image
                        is supported.
                    memory_threshold (int):
                        Memory usage threshold (percentage) to trigger caching behavior.
                    num_workers (int):
                        Number of workers for DataLoader and post-processing.
                    output_file (str):
                        Filename for saving output (e.g., "zarr" or "annotationstore").
                    return_labels (bool):
                        Whether to return labels with predictions.
                    scale_factor (tuple[float, float]):
                        Scale factor for annotations (model_mpp / slide_mpp).
                        Used to convert coordinates from non-baseline to baseline
                        resolution.
                    stride_shape (IntPair):
                        Stride used during WSI processing, at requested read resolution.
                        Must be positive. Defaults to `patch_input_shape` if not
                        provided.
                    wsireader_kwargs (WSIReaderParams):
                        Specify processing images with no mpp or power in the metadata.
                    verbose (bool):
                        Whether to enable verbose logging.

        Returns:
            dict[str, da.Array]:
                Post-processed predictions as a dictionary of Dask arrays.

        """
        return raw_predictions

    def save_predictions(
        self: EngineABC,
        processed_predictions: dict,
        output_type: str,
        save_path: Path | None = None,
        **kwargs: Unpack[EngineABCRunParams],
    ) -> dict | AnnotationStore | Path:
        """Save model predictions to disk or return them in memory.

        Depending on the output type, this method saves predictions as a zarr group,
        an AnnotationStore (SQLite database), or returns them as a dictionary.

        Args:
            processed_predictions (dict):
                Dictionary containing processed model predictions.
            output_type (str):
                Desired output format.
                Supported values are "dict", "zarr", "qupath" and "annotationstore".
            save_path (Path | None):
                Path to save the output file.
                Required for "zarr" and "annotationstore" formats.
            **kwargs (EngineABCRunParams):
                Additional runtime parameters to update engine attributes.

                Optional Keys:
                    auto_get_mask (bool):
                        Automatically generate segmentation masks using
                        `wsireader.tissue_mask()` during processing.
                    batch_size (int):
                        Number of image patches per forward pass.
                    class_dict (dict):
                        Mapping of classification outputs to class names.
                    device (str):
                        Device to run the model on (e.g., "cpu", "cuda").
                        See https://pytorch.org/docs/stable/tensor_attributes.html#torch.device
                        for more details.
                    labels (list):
                        Optional labels for input images. Only a single label per image
                        is supported.
                    memory_threshold (int):
                        Memory usage threshold (percentage) to trigger caching behavior.
                    num_workers (int):
                        Number of workers for DataLoader and post-processing.
                    output_file (str):
                        Filename for saving output (e.g., "zarr" or "annotationstore").
                    return_labels (bool):
                        Whether to return labels with predictions.
                    scale_factor (tuple[float, float]):
                        Scale factor for annotations (model_mpp / slide_mpp).
                        Used to convert coordinates from non-baseline to baseline
                        resolution.
                    stride_shape (IntPair):
                        Stride used during WSI processing, at requested read resolution.
                        Must be positive. Defaults to `patch_input_shape` if not
                        provided.
                    wsireader_kwargs (WSIReaderParams):
                        Specify processing images with no mpp or power in the metadata.
                    verbose (bool):
                        Whether to enable verbose logging.

        Returns:
            dict | AnnotationStore | Path:
                - If output_type is "dict": returns predictions as a dictionary.
                - If output_type is "zarr": returns path to saved zarr file.
                - If output_type is "qupath": returns a QuPath JSON
                  or path to .json file.
                - If output_type is "annotationstore": returns an AnnotationStore
                  or path to .db file.

        Raises:
            TypeError:
                If an unsupported output_type is provided.

        """
        keys_to_compute = [k for k in processed_predictions if k not in self.drop_keys]

        if output_type.lower() == "zarr":
            return self.save_predictions_as_zarr(
                processed_predictions=processed_predictions,
                save_path=save_path,
                keys_to_compute=keys_to_compute,
            )

        values_to_compute = [processed_predictions[k] for k in keys_to_compute]

        # Compute all at once
        computed_values = compute(*values_to_compute)

        # Assign computed values
        processed_predictions = dict(
            zip(keys_to_compute, computed_values, strict=False)
        )

        if output_type.lower() == "dict":
            return processed_predictions

        if output_type.lower() in ["qupath", "annotationstore"]:
            suffix = ".json" if output_type.lower() == "qupath" else ".db"
            save_path = Path(
                kwargs.get("output_file", save_path.parent / ("output" + suffix))
            )

            # scale_factor set from kwargs
            scale_factor = kwargs.get("scale_factor", (1.0, 1.0))
            # class_dict set from kwargs
            class_dict = kwargs.get("class_dict", self._get_model_attr("class_dict"))

            return dict_to_store_patch_predictions(
                processed_predictions,
                scale_factor,
                class_dict,
                save_path,
                output_type=output_type,
                verbose=self.verbose,
            )

        msg = f"Unsupported output type: {output_type}"
        raise TypeError(msg)

    @staticmethod
    def _get_tasks_for_saving_zarr(
        dask_output: da.Array | list,
        key: str,
        task_name: str | None,
        save_path: Path,
        write_tasks: list,
    ) -> list:
        """Helper function to get dask tasks for saving zarr output."""
        if isinstance(dask_output, da.Array):
            dask_output_dtype = dask_output.dtype
            object_codec = Pickle()
            if dask_output_dtype != "object":
                dask_output = dask_output.rechunk("auto")
                object_codec = None
            component = key if task_name is None else f"{task_name}/{key}"
            task = dask_output.to_zarr(
                url=save_path,
                component=component,
                compute=False,
                object_codec=object_codec,  # zarr kwargs
            )
            write_tasks.append(task)

        if isinstance(dask_output, list) and all(
            isinstance(dask_array, da.Array) for dask_array in dask_output
        ):
            for i, dask_array in enumerate(dask_output):
                object_codec = Pickle() if dask_array.dtype == "object" else None
                component = (
                    f"{key}/{i}" if task_name is None else f"{task_name}/{key}/{i}"
                )
                task = dask_array.to_zarr(
                    url=save_path,
                    component=component,
                    compute=False,
                    object_codec=object_codec,  # zarr kwargs
                )
                write_tasks.append(task)

        return write_tasks

    def save_predictions_as_zarr(
        self: EngineABC,
        processed_predictions: dict,
        save_path: Path,
        keys_to_compute: list,
        task_name: str | None = None,
    ) -> Path:
        """Save model predictions as a zarr file.

        This method saves the processed predictions to a zarr file at the specified
        path.

        Args:
            processed_predictions (dict):
                Dictionary containing processed model predictions.
            save_path (Path):
                Path to save the zarr file.
            keys_to_compute (list):
                List of keys in processed_predictions to save.
            task_name (str):
                Task Name for Multitask outputs.

        Returns:
            save_path (Path):
                Path to the saved zarr file.

        """
        if is_zarr(save_path):
            zarr_group = zarr.open(save_path, mode="r")
            keys_to_compute = [k for k in keys_to_compute if k not in zarr_group]

            # If the task group already exists, only compute missing keys
            if task_name is not None and task_name in zarr_group:
                task_group = zarr_group[task_name]
                keys_to_compute = [k for k in keys_to_compute if k not in task_group]

        write_tasks = []
        for key in keys_to_compute:
            dask_output = processed_predictions[key]
            write_tasks = self._get_tasks_for_saving_zarr(
                dask_output=dask_output,
                key=key,
                task_name=task_name,
                save_path=save_path,
                write_tasks=write_tasks,
            )

        msg = f"Saving output to {save_path}."
        _ = tqdm_dask_progress_bar(
            desc=msg,
            write_tasks=write_tasks,
            num_workers=self.num_workers,
            scheduler="threads",  # tasks are I/O-bound and shared memory use threads
            leave=False,
            verbose=self.verbose,
        )

        zarr_group = zarr.open(save_path, mode="r+")
        for key in self.drop_keys:
            if key in zarr_group:
                del zarr_group[key]

        return save_path

    def infer_wsi(
        self: EngineABC,
        dataloader: DataLoader,
        save_path: Path,
        **kwargs: Unpack[EngineABCRunParams],
    ) -> dict:
        """Run model inference on a whole slide image (WSI).

        This method performs inference on a WSI using the provided DataLoader,
        and accumulates predictions in Dask arrays. Optionally includes
        coordinates and labels in the output.

        Args:
            dataloader (DataLoader):
                PyTorch DataLoader configured for WSI processing.
            save_path (Path):
                Path to save the intermediate output. The intermediate output is saved
                in a zarr file.
            **kwargs (EngineABCRunParams):
                Additional runtime parameters to update engine attributes.

                Optional Keys:
                    auto_get_mask (bool):
                        Automatically generate segmentation masks using
                        `wsireader.tissue_mask()` during processing.
                    batch_size (int):
                        Number of image patches per forward pass.
                    class_dict (dict):
                        Mapping of classification outputs to class names.
                    device (str):
                        Device to run the model on (e.g., "cpu", "cuda").
                        See https://pytorch.org/docs/stable/tensor_attributes.html#torch.device
                        for more details.
                    labels (list):
                        Optional labels for input images. Only a single label per image
                        is supported.
                    memory_threshold (int):
                        Memory usage threshold (percentage) to trigger caching behavior.
                    num_workers (int):
                        Number of workers for DataLoader and post-processing.
                    output_file (str):
                        Filename for saving output (e.g., "zarr" or "annotationstore").
                    return_labels (bool):
                        Whether to return labels with predictions.
                    scale_factor (tuple[float, float]):
                        Scale factor for annotations (model_mpp / slide_mpp).
                        Used to convert coordinates from non-baseline to baseline
                        resolution.
                    stride_shape (IntPair):
                        Stride used during WSI processing, at requested read resolution.
                        Must be positive. Defaults to `patch_input_shape` if not
                        provided.
                    wsireader_kwargs (WSIReaderParams):
                        Specify processing images with no mpp or power in the metadata.
                    verbose (bool):
                        Whether to enable verbose logging.

        Returns:
            dict:
                Dictionary containing prediction results as Dask arrays.

        """
        _ = kwargs.get("patch_mode", False)
        _ = save_path
        return self.infer_patches(
            dataloader=dataloader,
            return_coordinates=True,
        )

    # This is not a static model for child classes.
    def post_process_wsi(  # skipcq: PYL-R0201
        self: EngineABC,
        raw_predictions: dict[str, da.Array],
        save_path: Path,  # noqa: ARG002
        **kwargs: Unpack[EngineABCRunParams],  # noqa: ARG002
    ) -> dict[str, da.Array]:
        """Post-process predictions from whole slide image (WSI) inference.

        This method applies a post-processing function (e.g., smoothing, filtering)
        to the raw model predictions. It supports delayed execution using Dask
        and returns a Dask array for efficient computation.

        Args:
            raw_predictions (dict[str, da.Array]):
                Dictionary containing raw model predictions as Dask arrays.
            save_path (Path):
                Path to save the intermediate output. The intermediate output is saved
                in a zarr file.
            **kwargs (EngineABCRunParams):
                Additional runtime parameters to update engine attributes.

                Optional Keys:
                    auto_get_mask (bool):
                        Automatically generate segmentation masks using
                        `wsireader.tissue_mask()` during processing.
                    batch_size (int):
                        Number of image patches per forward pass.
                    class_dict (dict):
                        Mapping of classification outputs to class names.
                    device (str):
                        Device to run the model on (e.g., "cpu", "cuda").
                        See https://pytorch.org/docs/stable/tensor_attributes.html#torch.device
                        for more details.
                    labels (list):
                        Optional labels for input images. Only a single label per image
                        is supported.
                    memory_threshold (int):
                        Memory usage threshold (percentage) to trigger caching behavior.
                    num_workers (int):
                        Number of workers for DataLoader and post-processing.
                    output_file (str):
                        Filename for saving output (e.g., "zarr" or "annotationstore").
                    return_labels (bool):
                        Whether to return labels with predictions.
                    scale_factor (tuple[float, float]):
                        Scale factor for annotations (model_mpp / slide_mpp).
                        Used to convert coordinates from non-baseline to baseline
                        resolution.
                    stride_shape (IntPair):
                        Stride used during WSI processing, at requested read resolution.
                        Must be positive. Defaults to `patch_input_shape` if not
                        provided.
                    wsireader_kwargs (WSIReaderParams):
                        Specify processing images with no mpp or power in the metadata.
                    verbose (bool):
                        Whether to enable verbose logging.

        Returns:
            dict[str, da.Array]:
                Post-processed predictions as a dictionary of Dask arrays.

        """
        return raw_predictions

    def _load_ioconfig(self: EngineABC, ioconfig: ModelIOConfigABC) -> ModelIOConfigABC:
        """Load or validate the IO configuration for the engine.

        If the model is from TIAToolbox and no IO configuration is provided,
        this method attempts to use the default configuration. Otherwise,
        it validates and sets the provided configuration.

        Args:
            ioconfig (ModelIOConfigABC):
                IO configuration to use for model inference.

        Returns:
            ModelIOConfigABC:
                The IO configuration to be used during inference.

        Raises:
            ValueError:
                If no IO configuration is provided and none is available from the model.

        """
        if self.ioconfig is None and ioconfig is None:
            msg = (
                "Please provide a valid ModelIOConfigABC. "
                "No default ModelIOConfigABC found."
            )
            logger.warning(msg)

        if ioconfig and isinstance(ioconfig, ModelIOConfigABC):
            self.ioconfig = ioconfig

        return self.ioconfig

    def _update_ioconfig(
        self: EngineABC,
        ioconfig: ModelIOConfigABC,
        patch_input_shape: IntPair,
        stride_shape: IntPair,
        input_resolutions: list[dict[Units, Resolution]],
    ) -> ModelIOConfigABC:
        """Update the IO configuration used for patch-based inference.

        This method updates the patch input shape, stride, and input resolutions
        in the IO configuration. If no configuration is provided, it creates a new one.

        Args:
            ioconfig (ModelIOConfigABC):
                Existing IO configuration to update. If None, a new one is created.
            patch_input_shape (IntPair):
                Size of patches input to the model (height, width). Patches are at
                requested read resolution, not with respect to level 0,
                and must be positive.
            stride_shape (IntPair):
                Stride used during patch extraction.
                If None, defaults to patch_input_shape.
                Stride is at requested read resolution, not with respect to
                level 0, and must be positive. If not provided,
                `stride_shape=patch_input_shape`.
            input_resolutions (list[dict[Units, Resolution]]):
                List of dictionaries specifying resolution and units
                for each input head. Supported units are `level`, `power` and `mpp`.
                Keys should be "units" and "resolution"
                e.g., [{"units": "mpp", "resolution": 0.25}]. Please see
                :class:`WSIReader` for details.

        Returns:
            ModelIOConfigABC:
                Updated IO configuration for patch-based inference.

        Raises:
            ValueError:
                If neither an IO configuration nor patch/resolution parameters
                are provided.

        """
        config_flag = (
            patch_input_shape is None,
            input_resolutions is None,
        )
        if isinstance(ioconfig, ModelIOConfigABC):
            return ioconfig

        if self.ioconfig is None and any(config_flag):
            msg = (
                "Must provide either "
                "`ioconfig` or `patch_input_shape` and `input_resolutions`."
            )
            raise ValueError(
                msg,
            )

        if stride_shape is None:
            stride_shape = patch_input_shape

        if self.ioconfig:
            ioconfig = copy.deepcopy(self.ioconfig)
            # ! not sure if there is a nicer way to set this
            if patch_input_shape is not None:
                ioconfig.patch_input_shape = patch_input_shape
            if stride_shape is not None:
                ioconfig.stride_shape = stride_shape
            if input_resolutions is not None:
                for input_resolutions_ in input_resolutions:
                    if input_resolutions_["units"] == "baseline":
                        ioconfig = ioconfig.to_baseline()
                ioconfig.input_resolutions = input_resolutions

            return ioconfig

        return ModelIOConfigABC(
            input_resolutions=input_resolutions,
            patch_input_shape=patch_input_shape,
            stride_shape=stride_shape,
            output_resolutions=[],
        )

    @staticmethod
    def _validate_images_masks(images: list | np.ndarray) -> list | np.ndarray:
        """Validate the format and shape of input images or masks.

        Ensures that the input is either a list of file paths or a 4D NumPy array
        in NHWC format.

        Args:
            images (list | np.ndarray):
                List of image paths or a NumPy array of image patches.

        Returns:
            list | np.ndarray:
                The validated input images or masks.

        Raises:
            TypeError:
                If the input is neither a list nor a NumPy array.
            ValueError:
                If the input is a NumPy array but not 4D (NHWC).

        """
        if not isinstance(images, (list, np.ndarray)):
            msg = "Input must be a list of file paths or a numpy array."
            raise TypeError(
                msg,
            )

        if isinstance(images, np.ndarray) and images.ndim != 4:  # noqa: PLR2004
            msg = (
                "The input numpy array should be four dimensional."
                "The shape of the numpy array should be NHWC."
            )
            raise ValueError(msg)

        return [Path(image) if isinstance(image, str) else image for image in images]

    @staticmethod
    def _validate_input_numbers(
        images: list | np.ndarray,
        masks: list[os.PathLike] | np.ndarray | None = None,
        labels: list | None = None,
    ) -> None:
        """Validate that the number of images, masks, and labels match.

        Ensures that the lengths of masks and labels (if provided) match
        the number of input images.

        Args:
            images (list | np.ndarray):
                List of input images or a NumPy array.
            masks (list[PathLike] | np.ndarray | None):
                Optional list of masks corresponding to the input images.
            labels (list | None):
                Optional list of labels corresponding to the input images.

        Returns:
            None

        Raises:
            ValueError:
                If the number of masks or labels does not match the number of images.

        """
        if masks is None and labels is None:
            return

        len_images = len(images)

        if masks is not None and len_images != len(masks):
            msg = (
                f"len(masks) is not equal to len(images) "
                f": {len(masks)} != {len(images)}"
            )
            raise ValueError(
                msg,
            )

        if labels is not None and len_images != len(labels):
            msg = (
                f"len(labels) is not equal to len(images) "
                f": {len(labels)} != {len(images)}"
            )
            raise ValueError(
                msg,
            )
        return

    def _update_run_params(
        self: EngineABC,
        images: list[os.PathLike | Path | WSIReader] | np.ndarray,
        masks: list[os.PathLike | Path] | np.ndarray | None = None,
        input_resolutions: list[dict[Units, Resolution]] | None = None,
        patch_input_shape: IntPair | None = None,
        save_dir: os.PathLike | Path | None = None,
        ioconfig: ModelIOConfigABC | None = None,
        output_type: str = "dict",
        *,
        overwrite: bool = False,
        patch_mode: bool,
        **kwargs: Unpack[EngineABCRunParams],
    ) -> Path | None:
        """Update runtime parameters for the engine before running inference.

        This method sets internal attributes such as caching, batch size,
        IO configuration, and output format based on user input and keyword arguments.

        Args:
            images (list[PathLike | Path | WSIReader] | np.ndarray):
                List of input images or a NumPy array of patches.
            masks (list[PathLike | Path] | np.ndarray | None):
                Optional list of masks for WSI processing.
            input_resolutions (list[dict[Units, Resolution]] | None):
                Resolution settings for input heads. Supported units are `level`,
                `power`, and `mpp`. Keys should be "units" and "resolution", e.g.,
                [{"units": "mpp", "resolution": 0.25}]. See :class:`WSIReader` for
                details.
            patch_input_shape (IntPair | None):
                Shape of input patches (height, width), requested at read
                resolution. Must be positive.
            save_dir (PathLike | Path | None):
                Directory to save output files. Required for WSI mode.
            ioconfig (ModelIOConfigABC | None):
                IO configuration for patch extraction and resolution settings.
            output_type (str):
                Desired output format: "dict", "zarr", "qupath" or "annotationstore".
            overwrite (bool):
                Whether to overwrite existing output files. Default is False.
            patch_mode (bool):
                Whether to treat input as patches (`True`) or WSIs (`False`).
            **kwargs (EngineABCRunParams):
                Additional runtime parameters to update engine attributes.

                Optional Keys:
                    auto_get_mask (bool):
                        Automatically generate segmentation masks using
                        `wsireader.tissue_mask()` during processing.
                    batch_size (int):
                        Number of image patches per forward pass.
                    class_dict (dict):
                        Mapping of classification outputs to class names.
                    device (str):
                        Device to run the model on (e.g., "cpu", "cuda").
                        See https://pytorch.org/docs/stable/tensor_attributes.html#torch.device
                        for more details.
                    labels (list):
                        Optional labels for input images. Only a single label per image
                        is supported.
                    memory_threshold (int):
                        Memory usage threshold (percentage) to trigger caching behavior.
                    num_workers (int):
                        Number of workers for DataLoader and post-processing.
                    output_file (str):
                        Filename for saving output (e.g., "zarr" or "annotationstore").
                    return_labels (bool):
                        Whether to return labels with predictions.
                    scale_factor (tuple[float, float]):
                        Scale factor for annotations (model_mpp / slide_mpp).
                        Used to convert coordinates from non-baseline to baseline
                        resolution.
                    stride_shape (IntPair):
                        Stride used during WSI processing, at requested read resolution.
                        Must be positive. Defaults to `patch_input_shape` if not
                        provided.
                    wsireader_kwargs (WSIReaderParams):
                        Specify processing images with no mpp or power in the metadata.
                    verbose (bool):
                        Whether to enable verbose logging.

        Returns:
            Path | None:
                Path to the save directory if applicable, otherwise None.

        Raises:
            TypeError:
                If an unsupported output_type is provided.
            ValueError:
                If required configuration or input parameters are missing.
            ValueError:
                If save_dir is not provided and output_type is "zarr", "qupath"
                or "annotationstore".

        """
        for key in kwargs:
            setattr(self, key, kwargs.get(key))

        if input_resolutions:
            self.input_resolutions = input_resolutions

        if patch_input_shape is not None:
            self.patch_input_shape = patch_input_shape

        if self.num_workers > 0:
            dask.config.set(scheduler="threads", num_workers=self.num_workers)
        else:
            dask.config.set(scheduler="threads")

        if not self.return_labels:
            self.drop_keys.append("label")

        self.patch_mode = patch_mode

        self._validate_input_numbers(images=images, masks=masks, labels=self.labels)
        if output_type.lower() not in ["dict", "zarr", "qupath", "annotationstore"]:
            msg = "output_type must be 'dict' or 'zarr', 'qupath' or 'annotationstore'."
            raise TypeError(msg)

        self.output_type = output_type

        if save_dir is not None and output_type.lower() not in [
            "zarr",
            "annotationstore",
            "qupath",
        ]:
            self.output_type = "zarr"
            msg = (
                f"output_type has been updated to 'zarr' "
                f"for saving the file to {save_dir}."
                f"Remove `save_dir` input to return the output as a `dict`."
            )
            logger.info(msg)

        if save_dir is None and output_type.lower() in [
            "zarr",
            "qupath",
            "annotationstore",
        ]:
            msg = f"Please provide save_dir for output_type={output_type}"
            raise ValueError(msg)

        self.images = self._validate_images_masks(images=images)

        if masks is not None:
            self.masks = self._validate_images_masks(images=masks)

        # if necessary move model parameters to "cpu" or "gpu" and update ioconfig
        self._ioconfig = self._load_ioconfig(ioconfig=ioconfig)
        self.model = self.model.to(device=self.device)
        self._ioconfig = self._update_ioconfig(
            ioconfig,
            self.patch_input_shape,
            self.stride_shape,
            self.input_resolutions,
        )

        return prepare_engines_save_dir(
            save_dir=save_dir,
            patch_mode=patch_mode,
            overwrite=overwrite,
        )

    def _run_patch_mode(
        self: EngineABC,
        output_type: str,
        save_dir: Path,
        **kwargs: Unpack[EngineABCRunParams],
    ) -> dict | AnnotationStore | Path:
        """Run the engine in patch mode.

        This method performs inference on image patches, post-processes the predictions,
        and saves the output in the specified format.

        Args:
            output_type (str):
                Desired output format. Supported values are "dict", "zarr",
                "qupath" and "annotationstore".
            save_dir (Path):
                Directory to save the output files.
            **kwargs (EngineABCRunParams):
                Additional runtime parameters to update engine attributes.

                Optional Keys:
                    auto_get_mask (bool):
                        Automatically generate segmentation masks using
                        `wsireader.tissue_mask()` during processing.
                    batch_size (int):
                        Number of image patches per forward pass.
                    class_dict (dict):
                        Mapping of classification outputs to class names.
                    device (str):
                        Device to run the model on (e.g., "cpu", "cuda").
                        See https://pytorch.org/docs/stable/tensor_attributes.html#torch.device
                        for more details.
                    labels (list):
                        Optional labels for input images. Only a single label per image
                        is supported.
                    memory_threshold (int):
                        Memory usage threshold (percentage) to trigger caching behavior.
                    num_workers (int):
                        Number of workers for DataLoader and post-processing.
                    output_file (str):
                        Filename for saving output (e.g., "zarr" or "annotationstore").
                    return_labels (bool):
                        Whether to return labels with predictions.
                    scale_factor (tuple[float, float]):
                        Scale factor for annotations (model_mpp / slide_mpp).
                        Used to convert coordinates from non-baseline to baseline
                        resolution.
                    stride_shape (IntPair):
                        Stride used during WSI processing, at requested read resolution.
                        Must be positive. Defaults to `patch_input_shape` if not
                        provided.
                    wsireader_kwargs (WSIReaderParams):
                        Specify processing images with no mpp or power in the metadata.
                    verbose (bool):
                        Whether to enable verbose logging.

        Returns:
            dict | AnnotationStore | Path:
                - If output_type is "dict": returns predictions as a dictionary.
                - If output_type is "zarr": returns path to saved zarr file.
                - If output_type is "qupath": returns a QuPath JSON
                  or path to .json file.
                - If output_type is "annotationstore": returns an AnnotationStore
                  or path to .db file.

        """
        save_path = None
        if save_dir:
            output_file = Path(kwargs.get("output_file", "output.zarr"))
            save_path = save_dir / (str(output_file.stem) + ".zarr")

        duplicate_filter = DuplicateFilter()
        logger.addFilter(duplicate_filter)
        # Set mask bounds to size of input patch if patch prediction
        # else size of output patch
        self.mask_bounds = (
            self._ioconfig.patch_output_shape
            if hasattr(self._ioconfig, "patch_output_shape")
            else self._ioconfig.patch_input_shape
        )

        self.dataloader = self.get_dataloader(
            images=self.images,
            masks=self.masks,
            labels=self.labels,
            patch_mode=True,
            ioconfig=self._ioconfig,
            wsireader_kwargs=kwargs.get("wsireader_kwargs"),
        )
        raw_predictions = self.infer_patches(
            dataloader=self.dataloader,
            return_coordinates=output_type.lower() in ["annotationstore", "qupath"],
        )

        processed_predictions = self.post_process_patches(
            raw_predictions=raw_predictions,
            **kwargs,
        )

        logger.removeFilter(duplicate_filter)

        out = self.save_predictions(
            processed_predictions=processed_predictions,
            output_type=output_type,
            save_path=save_path,
            **kwargs,
        )

        if isinstance(out, dict):
            return out

        msg = f"Output file saved at {out}."
        logger.info(msg=msg)
        gc.collect()  # Clean up
        return out

    @staticmethod
    def _calculate_scale_factor(dataloader: DataLoader) -> float | tuple[float, float]:
        """Calculate the scale factor for final output based on dataloader resolution.

        This method compares the resolution used during reading with the slide's
        baseline resolution to compute a scale factor for coordinate transformation.

        Args:
            dataloader (DataLoader):
                PyTorch DataLoader used for WSI inference. Must contain resolution
                and unit metadata in its dataset.

        Returns:
            float | tuple[float, float]:
                Scale factor for converting coordinates to baseline resolution.
                - If units are "mpp": returns (model_mpp / slide_mpp).
                - If units are "level": returns downsample ratio.
                - If units are "power": returns objective_power / model_power.
                - If units are "baseline": returns the resolution directly.

        """
        # get units and resolution from dataloader.
        dataloader_units = dataloader.dataset.units
        dataloader_resolution = dataloader.dataset.resolution

        # if dataloader units is baseline slide resolution is 1.0.
        # in this case dataloader resolution / slide resolution will be
        # equal to dataloader resolution.

        if dataloader_units in ["mpp", "level", "power"]:
            wsimeta_dict = dataloader.dataset.reader_info.as_dict()

        if dataloader_units == "mpp":
            slide_resolution = wsimeta_dict[dataloader_units]
            scale_factor = np.divide(dataloader_resolution, slide_resolution)
            return scale_factor[0], scale_factor[1]

        if dataloader_units == "level":
            downsample_ratio = wsimeta_dict["level_downsamples"][dataloader_resolution]
            return downsample_ratio, downsample_ratio

        if dataloader_units == "power":
            slide_objective_power = wsimeta_dict["objective_power"]
            return (
                slide_objective_power / dataloader_resolution,
                slide_objective_power / dataloader_resolution,
            )

        return dataloader_resolution

    def _run_wsi_mode(
        self: EngineABC,
        output_type: str,
        save_dir: Path,
        **kwargs: Unpack[EngineABCRunParams],
    ) -> dict | AnnotationStore | Path:
        """Run the engine in WSI mode (patch_mode = False).

        This method performs inference on each whole slide image (WSI),
        post-processes the predictions, and saves the output in the specified format.

        Args:
            output_type (str):
                Desired output format. Supported values are "dict", "zarr",
                "qupath" and "annotationstore".
            save_dir (Path):
                Directory to save the output files.
            **kwargs (EngineABCRunParams):
                Additional runtime parameters to update engine attributes.

                Optional Keys:
                    auto_get_mask (bool):
                        Automatically generate segmentation masks using
                        `wsireader.tissue_mask()` during processing.
                    batch_size (int):
                        Number of image patches per forward pass.
                    class_dict (dict):
                        Mapping of classification outputs to class names.
                    device (str):
                        Device to run the model on (e.g., "cpu", "cuda").
                        See https://pytorch.org/docs/stable/tensor_attributes.html#torch.device
                        for more details.
                    labels (list):
                        Optional labels for input images. Only a single label per image
                        is supported.
                    memory_threshold (int):
                        Memory usage threshold (percentage) to trigger caching behavior.
                    num_workers (int):
                        Number of workers for DataLoader and post-processing.
                    output_file (str):
                        Filename for saving output (e.g., "zarr" or "annotationstore").
                    return_labels (bool):
                        Whether to return labels with predictions.
                    scale_factor (tuple[float, float]):
                        Scale factor for annotations (model_mpp / slide_mpp).
                        Used to convert coordinates from non-baseline to baseline
                        resolution.
                    stride_shape (IntPair):
                        Stride used during WSI processing, at requested read resolution.
                        Must be positive. Defaults to `patch_input_shape` if not
                        provided.
                    wsireader_kwargs (WSIReaderParams):
                        Specify processing images with no mpp or power in the metadata.
                    verbose (bool):
                        Whether to enable verbose logging.

        Returns:
            dict | AnnotationStore | Path:
                Dictionary mapping each input WSI to its corresponding output path.
                Output may be a zarr file, SQLite database, or in-memory dictionary.

        """
        suffix = ".zarr"
        if output_type.lower() == "annotationstore":
            suffix = ".db"
        if output_type.lower() == "qupath":
            suffix = ".json"

        def get_path(image: Path | WSIReader) -> Path:
            """Return path to output file."""
            return image.input_path if isinstance(image, WSIReader) else image

        out = {
            get_path(image): save_dir / (get_path(image).stem + suffix)
            for image in self.images
        }

        save_path = {
            get_path(image): save_dir / (get_path(image).stem + ".zarr")
            for image in self.images
        }

        tqdm_loop = tqdm(
            self.images,
            leave=False,
            desc="Processing WSIs",
            disable=not self.verbose,
        )

        for image_num, image in enumerate(tqdm_loop):
            duplicate_filter = DuplicateFilter()
            logger.addFilter(duplicate_filter)
            mask = self.masks[image_num] if self.masks is not None else None
            self.dataloader = self.get_dataloader(
                images=image,
                masks=mask,
                patch_mode=False,
                ioconfig=self._ioconfig,
                auto_get_mask=kwargs.get("auto_get_mask", True),
                wsireader_kwargs=kwargs.get("wsireader_kwargs"),
            )

            scale_factor = self._calculate_scale_factor(dataloader=self.dataloader)

            # WSI mode uses WSIPatchDataSet for dataloader
            output_locations = np.asarray(self.dataloader.dataset.outputs)
            self.mask_bounds = (
                None
                if len(output_locations) == 0
                else (
                    output_locations[:, 0].min(),
                    output_locations[:, 1].min(),
                    output_locations[:, 2].max(),
                    output_locations[:, 3].max(),
                )
            )

            raw_predictions = self.infer_wsi(
                dataloader=self.dataloader,
                save_path=save_path[get_path(image)],
                **kwargs,
            )

            processed_predictions = self.post_process_wsi(
                raw_predictions=raw_predictions,
                save_path=save_path[get_path(image)],
                **kwargs,
            )

            kwargs["output_file"] = out[get_path(image)]
            kwargs["scale_factor"] = scale_factor
            out[get_path(image)] = self.save_predictions(
                processed_predictions=processed_predictions,
                output_type=output_type,
                save_path=save_path[get_path(image)],
                **kwargs,
            )
            logger.removeFilter(duplicate_filter)
            msg = f"Output file saved at {out[get_path(image)]}."
            logger.info(msg=msg)
            gc.collect()  # Clean up

        return out

    def run(
        self: EngineABC,
        images: list[os.PathLike | Path | WSIReader] | np.ndarray,
        *,
        masks: list[os.PathLike | Path | np.ndarray] | np.ndarray | None = None,
        input_resolutions: list[dict[Units, Resolution]] | None = None,
        patch_input_shape: IntPair | None = None,
        ioconfig: ModelIOConfigABC | None = None,
        patch_mode: bool = True,
        save_dir: os.PathLike | Path | None = None,
        overwrite: bool = False,
        output_type: str = "dict",
        **kwargs: Unpack[EngineABCRunParams],
    ) -> AnnotationStore | Path | str | dict:
        """Run the engine on input images.

        This method orchestrates the full inference pipeline, including preprocessing,
        model inference, post-processing, and saving results. It supports both patch
        and WSI modes.

        Args:
            images (list[PathLike | Path | WSIReader] | np.ndarray):
                List of input images or a NumPy array of patches.
            masks (list[PathLike | Path] | np.ndarray | None):
                Optional list of masks for WSI processing.
                Only utilised when patch_mode is False.
                Patches are only generated within a masked area.
                If not provided, then a tissue mask will be automatically
                generated for whole slide images.
            input_resolutions (list[dict[Units, Resolution]] | None):
                Resolution settings for input heads. Supported units are `level`,
                `power`, and `mpp`. Keys should be "units" and "resolution", e.g.,
                [{"units": "mpp", "resolution": 0.25}]. See :class:`WSIReader` for
                details.
            patch_input_shape (IntPair | None):
                Shape of input patches (height, width), requested at read
                resolution. Must be positive.
            ioconfig (ModelIOConfigABC | None):
                IO configuration for patch extraction and resolution settings.
            patch_mode (bool):
                Whether to treat input as patches (`True`) or WSIs (`False`).
                Default is True.
            save_dir (PathLike | Path | None):
                Directory to save output files. Required for WSI mode.
            overwrite (bool):
                Whether to overwrite existing output files. Default is False.
            output_type (str):
                Desired output format: "dict", "zarr", "qupath", or "annotationstore".
            **kwargs (EngineABCRunParams):
                Additional runtime parameters to update engine attributes.

                Optional Keys:
                    auto_get_mask (bool):
                        Automatically generate segmentation masks using
                        `wsireader.tissue_mask()` during processing.
                        EngineABC defaults to True.
                    batch_size (int):
                        Number of image patches per forward pass.
                    class_dict (dict):
                        Mapping of classification outputs to class names.
                    device (str):
                        Device to run the model on (e.g., "cpu", "cuda").
                        See https://pytorch.org/docs/stable/tensor_attributes.html#torch.device
                        for more details.
                    labels (list):
                        Optional labels for input images. Only a single label per image
                        is supported.
                    memory_threshold (int):
                        Memory usage threshold (percentage) to trigger caching behavior.
                        Default is 80.
                    num_workers (int):
                        Number of workers for DataLoader and post-processing.
                    output_file (str):
                        Filename for saving output (e.g., "zarr" or "annotationstore").
                    return_labels (bool):
                        Whether to return labels with predictions.
                    scale_factor (tuple[float, float]):
                        Scale factor for annotations (model_mpp / slide_mpp).
                        Used to convert coordinates from non-baseline to baseline
                        resolution.
                    stride_shape (IntPair):
                        Stride used during WSI processing, at requested read resolution.
                        Must be positive. Defaults to `patch_input_shape` if not
                        provided.
                    wsireader_kwargs (WSIReaderParams):
                        Specify processing images with no mpp or power in the metadata.
                    verbose (bool):
                        Whether to enable verbose logging.

        Returns:
            AnnotationStore | Path | str | dict:
                - If patch_mode is True: returns predictions or path to saved output.
                - If patch_mode is False: returns a dictionary mapping each WSI to
                  its output path.

        Examples:
            >>> wsis = ['wsi1.svs', 'wsi2.svs']
            >>> class PatchPredictor(EngineABC):
            >>> # Define all Abstract methods.
            >>>     ...
            >>> predictor = PatchPredictor(model="resnet18-kather100k")
            >>> output = predictor.run(image_patches, patch_mode=True)
            >>> output
            ... "/path/to/Output.db"
            >>> output = predictor.run(
            >>>     image_patches,
            >>>     patch_mode=True,
            >>>     output_type="zarr")
            >>> output
            ... "/path/to/Output.zarr"
            >>> output = predictor.run(wsis, patch_mode=False)
            >>> output.keys()
            ... ['wsi1.svs', 'wsi2.svs']
            >>> output['wsi1.svs']
            ... {'/path/to/wsi1.db'}

        """
        save_dir = self._update_run_params(
            images=images,
            masks=masks,
            input_resolutions=input_resolutions,
            patch_input_shape=patch_input_shape,
            save_dir=save_dir,
            ioconfig=ioconfig,
            overwrite=overwrite,
            patch_mode=patch_mode,
            output_type=output_type,
            **kwargs,
        )

        if patch_mode:
            return self._run_patch_mode(
                output_type=self.output_type,
                save_dir=save_dir,
                **kwargs,
            )

        # All inherited classes will get scale_factors,
        # highest_input_resolution, implement dataloader,
        # pre-processing, post-processing and save_output
        # for WSIs separately.
        return self._run_wsi_mode(
            output_type=self.output_type,
            save_dir=save_dir,
            **kwargs,
        )


def prepare_engines_save_dir(
    save_dir: str | Path | None,
    *,
    patch_mode: bool,
    overwrite: bool = False,
) -> Path | None:
    """Create or validate the save directory for engine outputs.

    Args:
        save_dir (str | Path | None):
            Path to the output directory.
        patch_mode (bool):
            Whether the input is treated as patches.
        overwrite (bool):
            Whether to overwrite existing directory. Default is False.

    Returns:
        Path | None:
            Path to the output directory if created or validated, else None.

    Raises:
        OSError:
            If patch_mode is False and save_dir is not provided.

    """
    if patch_mode:
        if save_dir is not None:
            save_dir = Path(save_dir)
            if save_dir.exists() and overwrite:
                shutil.rmtree(save_dir)
            save_dir.mkdir(parents=True)
            return save_dir
        return None

    if save_dir is None:
        msg = (
            "Input WSIs detected but no save directory provided. "
            "Please provide a 'save_dir'."
        )
        raise OSError(msg)

    logger.info(
        "When providing multiple whole slide images, "
        "the outputs will be saved and the locations of outputs "
        "will be returned to the calling function when `run()` "
        "finishes successfully."
    )

    save_dir = Path(save_dir)
    if save_dir.exists() and overwrite:
        shutil.rmtree(save_dir)
    save_dir.mkdir(parents=True)

    return save_dir
