"""Patch-based training and inference for 3D medical image segmentation using TorchIO's Queue mechanism."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/10_vision_patch.ipynb.

# %% auto #0
__all__ = ['normalize_patch_transforms', 'PatchConfig', 'med_to_subject', 'create_subjects_dataset', 'create_patch_sampler',
           'MedPatchDataLoader', 'MedPatchDataLoaders', 'PatchInferenceEngine', 'patch_inference']

# %% ../nbs/10_vision_patch.ipynb #cell-2
import torch
import torchio as tio
import pandas as pd
import numpy as np
import warnings
import matplotlib.pyplot as plt
from pathlib import Path
from dataclasses import dataclass, field
from typing import Callable
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from fastai.data.all import *
from fastai.learner import Learner
from .vision_core import MedImage, MedMask, MedBase, med_img_reader
from .vision_plot import find_max_slice
from .vision_inference import _to_original_orientation, _do_resize
from .dataset_info import MedDataset, suggest_patch_size

# %% ../nbs/10_vision_patch.ipynb #aw2pkvm2ibe
def _get_default_device() -> torch.device:
    """Get the default device (CUDA if available, else CPU)."""
    return torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def _warn_config_override(param_name: str, config_value, explicit_value):
    """Warn when explicit argument overrides config value.

    Args:
        param_name: Name of the parameter (e.g., 'apply_reorder', 'target_spacing')
        config_value: Value from PatchConfig
        explicit_value: Explicitly provided value
    """
    if explicit_value is not None and config_value is not None:
        if explicit_value != config_value:
            warnings.warn(
                f"{param_name} mismatch: explicit={explicit_value}, config={config_value}. "
                f"Using explicit argument."
            )


def _extract_tio_transform(tfm):
    """Extract TorchIO transform from fastMONAI wrapper or return as-is.

    This function enables using fastMONAI wrappers (e.g., RandomAffine, RandomGamma)
    in patch-based workflows where raw TorchIO transforms are needed for tio.Compose().

    Uses the explicit `.tio_transform` property when available on fastMONAI wrappers.
    Falls back to returning the transform unchanged for raw TorchIO transforms.

    Args:
        tfm: fastMONAI wrapper (e.g., RandomAffine) or raw TorchIO transform

    Returns:
        The underlying TorchIO transform

    Example:
        >>> from fastMONAI.vision_augmentation import RandomAffine
        >>> wrapped = RandomAffine(degrees=10)
        >>> raw = _extract_tio_transform(wrapped)  # Returns tio.RandomAffine
    """
    return getattr(tfm, 'tio_transform', tfm)


def normalize_patch_transforms(tfms: list) -> list:
    """Normalize transforms for patch-based workflow.

    Extracts underlying TorchIO transforms from fastMONAI wrappers.
    Also accepts raw TorchIO transforms for backward compatibility.

    This enables using the same transform syntax in both standard and
    patch-based workflows:

        >>> from fastMONAI.vision_augmentation import RandomAffine, RandomGamma
        >>>
        >>> # Same syntax works in both contexts
        >>> item_tfms = [RandomAffine(degrees=10), RandomGamma(p=0.5)]   # Standard
        >>> patch_tfms = [RandomAffine(degrees=10), RandomGamma(p=0.5)]  # Patch-based

    Args:
        tfms: List of fastMONAI wrappers or raw TorchIO transforms

    Returns:
        List of raw TorchIO transforms suitable for tio.Compose()
    """
    if tfms is None:
        return None
    return [_extract_tio_transform(t) for t in tfms]

# %% ../nbs/10_vision_patch.ipynb #cell-5
@dataclass
class PatchConfig:
    """Configuration for patch-based training and inference.
    
    Args:
        patch_size: Size of patches [x, y, z].
        patch_overlap: Overlap for inference GridSampler (int, float 0-1, or list).
            - Float 0-1: fraction of patch_size (e.g., 0.5 = 50% overlap)
            - Int >= 1: pixel overlap (e.g., 48 = 48 pixel overlap)
            - List: per-dimension overlap in pixels
        samples_per_volume: Number of patches to extract per volume during training.
        sampler_type: Type of sampler ('uniform', 'label', 'weighted').
        label_probabilities: For LabelSampler, dict mapping label values to probabilities.
        queue_length: Maximum number of patches to store in queue.
        queue_num_workers: Number of workers for parallel patch extraction.
        aggregation_mode: For inference, how to combine overlapping patches ('crop', 'average', 'hann').
        apply_reorder: Whether to reorder to RAS+ canonical orientation. Must match between
            training and inference. Defaults to True (the common case).
        target_spacing: Target voxel spacing [x, y, z] for resampling. Must match between
            training and inference.
        preprocessed: If True, data has been preprocessed externally (e.g., via
            preprocess_dataset()). Training will skip reorder, resample, AND
            pre_patch_tfms (e.g., normalization) since they were already applied.
            Inference is unaffected and always applies pre_inference_tfms to raw
            images. Defaults to False.
        padding_mode: Padding mode for CropOrPad when image < patch_size. Default is 0 (zero padding)
            to align with nnU-Net's approach. Can be int, float, or string (e.g., 'minimum', 'mean').
        keep_largest_component: If True, keep only the largest connected component
            in binary segmentation predictions. Only applies during inference when
            return_probabilities=False. Defaults to False.
    
    Example:
        >>> config = PatchConfig(
        ...     patch_size=[96, 96, 96],
        ...     samples_per_volume=16,
        ...     sampler_type='label',
        ...     label_probabilities={0: 0.1, 1: 0.9},
        ...     target_spacing=[0.5, 0.5, 0.5]
        ... )
    """
    patch_size: list = field(default_factory=lambda: [96, 96, 96])
    patch_overlap: int | float | list = 0
    samples_per_volume: int = 8
    sampler_type: str = 'uniform'
    label_probabilities: dict = None
    queue_length: int = 300
    queue_num_workers: int = 4
    aggregation_mode: str = 'hann'
    # Preprocessing parameters - must match between training and inference
    apply_reorder: bool = True  # Defaults to True (the common case)
    target_spacing: list = None
    preprocessed: bool = False  # True = data already preprocessed, skip all preprocessing during training
    padding_mode: int | float | str = 0  # Zero padding (nnU-Net standard)
    # Post-processing (binary segmentation only)
    keep_largest_component: bool = False
    
    def __post_init__(self):
        """Validate configuration."""
        valid_samplers = ['uniform', 'label', 'weighted']
        if self.sampler_type not in valid_samplers:
            raise ValueError(f"sampler_type must be one of {valid_samplers}")
        
        valid_aggregation = ['crop', 'average', 'hann']
        if self.aggregation_mode not in valid_aggregation:
            raise ValueError(f"aggregation_mode must be one of {valid_aggregation}")
        
        # Validate patch_overlap
        # Negative overlap doesn't make sense
        if isinstance(self.patch_overlap, (int, float)):
            if self.patch_overlap < 0:
                raise ValueError("patch_overlap cannot be negative")
            # Check if overlap as pixels would exceed patch_size (causes step_size=0)
            if self.patch_overlap >= 1:  # Pixel value, not fraction
                for ps in self.patch_size:
                    if self.patch_overlap >= ps:
                        raise ValueError(
                            f"patch_overlap ({self.patch_overlap}) must be less than patch_size ({ps}). "
                            f"Overlap >= patch_size creates step_size <= 0 (infinite patches)."
                        )
        elif isinstance(self.patch_overlap, (list, tuple)):
            for i, (overlap, ps) in enumerate(zip(self.patch_overlap, self.patch_size)):
                if overlap < 0:
                    raise ValueError(f"patch_overlap[{i}] cannot be negative")
                if overlap >= ps:
                    raise ValueError(
                        f"patch_overlap[{i}] ({overlap}) must be less than patch_size[{i}] ({ps}). "
                        f"Overlap >= patch_size creates step_size <= 0 (infinite patches)."
                    )

        # Warn if patch_size dimensions are not divisible by 16
        non_div = [s for s in self.patch_size if s % 16 != 0]
        if non_div:
            warnings.warn(
                f"patch_size {self.patch_size} has dimensions not divisible by 16. "
                f"Most encoder-decoder architectures (e.g., U-Net) require patch sizes "
                f"divisible by 16 (2^4 for 4 downsampling levels)."
            )

    @classmethod
    def from_dataset(
        cls,
        dataset: 'MedDataset',
        target_spacing: list = None,
        min_patch_size: list = None,
        max_patch_size: list = None,
        divisor: int = 16,
        **kwargs
    ) -> 'PatchConfig':
        """Create PatchConfig with automatic patch_size from dataset analysis.

        Combines dataset preprocessing suggestions with patch size calculation
        for a complete, DRY configuration.

        Args:
            dataset: MedDataset instance with analyzed images.
            target_spacing: Target voxel spacing [x, y, z]. If None, uses
                dataset.get_suggestion()['target_spacing'].
            min_patch_size: Minimum per dimension [32, 32, 32].
            max_patch_size: Maximum per dimension [256, 256, 256].
            divisor: Divisibility constraint (default 16 for UNet compatibility).
            **kwargs: Additional PatchConfig parameters (samples_per_volume,
                sampler_type, label_probabilities, etc.).

        Returns:
            PatchConfig with suggested patch_size, apply_reorder, target_spacing.

        Example:
            >>> from fastMONAI.dataset_info import MedDataset
            >>> dataset = MedDataset(dataframe=df, mask_col='mask_path', dtype=MedMask)
            >>> 
            >>> # Use recommended spacing
            >>> config = PatchConfig.from_dataset(dataset, samples_per_volume=16)
            >>> 
            >>> # Use custom spacing
            >>> config = PatchConfig.from_dataset(
            ...     dataset,
            ...     target_spacing=[1.0, 1.0, 2.0],
            ...     samples_per_volume=16
            ... )
        """
        # Get preprocessing suggestion from dataset
        suggestion = dataset.get_suggestion()

        # Use explicit spacing or dataset suggestion
        _target_spacing = target_spacing if target_spacing is not None else suggestion['target_spacing']

        # Calculate patch size for the target spacing
        patch_size = suggest_patch_size(
            dataset,
            target_spacing=_target_spacing,
            min_patch_size=min_patch_size,
            max_patch_size=max_patch_size,
            divisor=divisor
        )

        # Merge with explicit kwargs (kwargs override defaults)
        # Use dataset.apply_reorder directly (not from get_suggestion() since it's not data-derived)
        config_kwargs = {
            'patch_size': patch_size,
            'apply_reorder': dataset.apply_reorder,
            'target_spacing': _target_spacing,
        }
        config_kwargs.update(kwargs)

        return cls(**config_kwargs)

# %% ../nbs/10_vision_patch.ipynb #cell-8
def med_to_subject(
    img: Path | str,
    mask: Path | str = None,
) -> tio.Subject:
    """Create TorchIO Subject with LAZY loading (paths only, no tensor loading).
    
    This function stores file paths in the Subject, allowing TorchIO's Queue
    workers to load volumes on-demand during training. This is memory-efficient
    as volumes are not loaded into RAM until needed.
    
    Args:
        img: Path to image file.
        mask: Path to mask file (optional).
    
    Returns:
        TorchIO Subject with 'image' and optionally 'mask' keys (lazy loaded).
    
    Example:
        >>> subject = med_to_subject('image.nii.gz', 'mask.nii.gz')
        >>> # Volume NOT loaded yet - only path stored
        >>> data = subject['image'].data  # NOW volume is loaded
    """
    subject_dict = {
        'image': tio.ScalarImage(path=str(img))  # Lazy - stores path only
    }
    
    if mask is not None:
        subject_dict['mask'] = tio.LabelMap(path=str(mask))  # Lazy
    
    return tio.Subject(**subject_dict)

# %% ../nbs/10_vision_patch.ipynb #cell-9
def create_subjects_dataset(
    df: pd.DataFrame,
    img_col: str,
    mask_col: str = None,
    pre_tfms: list = None,
    ensure_affine_consistency: bool = True
) -> tio.SubjectsDataset:
    """Build TorchIO SubjectsDataset with LAZY loading from DataFrame.

    This function creates a SubjectsDataset that stores only file paths,
    not loaded tensors. Volumes are loaded on-demand by Queue workers,
    keeping memory usage constant regardless of dataset size.

    Args:
        df: DataFrame with image (and optionally mask) paths.
        img_col: Column name containing image paths.
        mask_col: Column name containing mask paths (optional).
        pre_tfms: List of TorchIO transforms to apply before patch extraction.
                  Use tio.ToCanonical() for reordering and tio.Resample() for resampling.
        ensure_affine_consistency: If True and mask_col is provided, automatically
            prepends tio.CopyAffine(target='image') to ensure spatial metadata
            consistency between image and mask. This prevents "More than one value
            for direction found" errors. Defaults to True.

    Returns:
        TorchIO SubjectsDataset with lazy-loaded subjects.

    Example:
        >>> # Preprocessing via transforms (applied by workers on-demand)
        >>> pre_tfms = [
        ...     tio.ToCanonical(),           # Reorder to RAS+
        ...     tio.Resample([0.5, 0.5, 0.5]),  # Resample
        ...     tio.ZNormalization(),        # Intensity normalization
        ... ]
        >>> dataset = create_subjects_dataset(
        ...     df, img_col='image', mask_col='label',
        ...     pre_tfms=pre_tfms
        ... )
        >>> # Memory: ~0 MB (only paths stored, not volumes)
    """
    subjects = []
    for idx, row in df.iterrows():
        img_path = row[img_col]
        mask_path = row[mask_col] if mask_col else None

        # Create subject with lazy loading (paths only)
        subject = med_to_subject(img=img_path, mask=mask_path)
        subjects.append(subject)

    # Build transform pipeline
    all_transforms = []

    # Add CopyAffine as FIRST transform when mask is present
    # This ensures spatial metadata consistency before other transforms
    if mask_col is not None and ensure_affine_consistency:
        all_transforms.append(tio.CopyAffine(target='image'))

    # Add user-provided transforms
    if pre_tfms:
        all_transforms.extend(pre_tfms)

    transform = tio.Compose(all_transforms) if all_transforms else None

    return tio.SubjectsDataset(subjects, transform=transform)

# %% ../nbs/10_vision_patch.ipynb #cell-11
def create_patch_sampler(config: PatchConfig) -> tio.data.PatchSampler:
    """Create appropriate TorchIO sampler based on config.
    
    Args:
        config: PatchConfig with sampler settings.
    
    Returns:
        TorchIO PatchSampler instance.
    
    Example:
        >>> config = PatchConfig(patch_size=[96, 96, 96], sampler_type='label')
        >>> sampler = create_patch_sampler(config)
    """
    patch_size = config.patch_size
    
    if config.sampler_type == 'uniform':
        return tio.UniformSampler(patch_size)
    
    elif config.sampler_type == 'label':
        return tio.LabelSampler(
            patch_size,
            label_name='mask',
            label_probabilities=config.label_probabilities
        )
    
    elif config.sampler_type == 'weighted':
        raise NotImplementedError(
            "WeightedSampler requires a pre-computed probability map which is not currently supported. "
            "Use 'label' sampler with label_probabilities for weighted sampling based on segmentation labels, "
            "or 'uniform' for random patch extraction."
        )
    
    raise ValueError(f"Unknown sampler type: {config.sampler_type}")

# %% ../nbs/10_vision_patch.ipynb #cell-14
class MedPatchDataLoader:
    """DataLoader wrapper for patch-based training with TorchIO Queue.

    This class wraps a TorchIO Queue to provide a fastai-compatible DataLoader
    interface for patch-based training.

    Args:
        subjects_dataset: TorchIO SubjectsDataset.
        config: PatchConfig with queue and sampler settings.
        batch_size: Number of patches per batch. Must be positive.
        patch_tfms: Transforms to apply to extracted patches (training only).
            Accepts both fastMONAI wrappers (e.g., RandomAffine, RandomGamma) and
            raw TorchIO transforms. fastMONAI wrappers are automatically normalized
            to raw TorchIO for internal use. Mutually exclusive with gpu_augmentation.
        gpu_augmentation: GpuPatchAugmentation instance for GPU-batched augmentation.
            Operates on [B,C,D,H,W] tensors already on GPU, avoiding per-sample CPU
            overhead. Mutually exclusive with patch_tfms. Training only.
        shuffle: Whether to shuffle subjects and patches.
        drop_last: Whether to drop last incomplete batch.
    """

    def __init__(
        self,
        subjects_dataset: tio.SubjectsDataset,
        config: PatchConfig,
        batch_size: int = 4,
        patch_tfms: list = None,
        gpu_augmentation=None,
        shuffle: bool = True,
        drop_last: bool = False
    ):
        if batch_size <= 0:
            raise ValueError(f"batch_size must be positive, got {batch_size}")

        self.subjects_dataset = subjects_dataset
        self.config = config
        self.bs = batch_size
        self.shuffle = shuffle
        self.drop_last = drop_last
        self._device = _get_default_device()
        self.gpu_augmentation = gpu_augmentation

        # Create sampler
        self.sampler = create_patch_sampler(config)

        # Create patch transforms
        # Normalize transforms - accepts both fastMONAI wrappers and raw TorchIO
        normalized_tfms = normalize_patch_transforms(patch_tfms)
        self.patch_tfms = tio.Compose(normalized_tfms) if normalized_tfms else None

        # Create queue
        self.queue = tio.Queue(
            subjects_dataset,
            max_length=config.queue_length,
            samples_per_volume=config.samples_per_volume,
            sampler=self.sampler,
            num_workers=config.queue_num_workers,
            shuffle_subjects=shuffle,
            shuffle_patches=shuffle
        )

        # Create torch DataLoader
        self._dl = DataLoader(
            self.queue,
            batch_size=batch_size,
            num_workers=0,  # Queue handles workers
            drop_last=drop_last
        )

        # Track cleanup state
        self._closed = False

    def __iter__(self):
        """Iterate over batches, yielding (image, mask) tuples."""
        for batch in self._dl:
            # Extract image and mask tensors
            img = batch['image'][tio.DATA]  # [B, C, H, W, D]
            has_mask = 'mask' in batch

            # Apply CPU patch transforms if provided (per-sample TorchIO loop)
            if self.patch_tfms is not None:
                transformed_imgs = []
                transformed_masks = [] if has_mask else None

                for i in range(img.shape[0]):
                    subject_dict = {'image': tio.ScalarImage(tensor=batch['image'][tio.DATA][i])}
                    if has_mask:
                        subject_dict['mask'] = tio.LabelMap(tensor=batch['mask'][tio.DATA][i])

                    subject = tio.Subject(subject_dict)
                    transformed = self.patch_tfms(subject)
                    transformed_imgs.append(transformed['image'].data)
                    if has_mask:
                        transformed_masks.append(transformed['mask'].data)

                img = torch.stack(transformed_imgs)
                mask = torch.stack(transformed_masks) if has_mask else None
            else:
                mask = batch['mask'][tio.DATA] if has_mask else None

            # Move to device
            img = img.to(self._device)
            if mask is not None:
                mask = mask.to(self._device)

            # Apply GPU augmentation if provided (batched, on-device)
            if self.gpu_augmentation is not None:
                img, mask = self.gpu_augmentation(img, mask)

            # Wrap as MedImage/MedMask
            img = MedImage(img)
            if mask is not None:
                mask = MedMask(mask)

            yield img, mask

    def __len__(self):
        """Return number of batches per epoch."""
        n_patches = len(self.subjects_dataset) * self.config.samples_per_volume
        if self.drop_last:
            return n_patches // self.bs
        return (n_patches + self.bs - 1) // self.bs

    @property
    def dataset(self):
        """Return the underlying queue as dataset."""
        return self.queue

    @property
    def device(self):
        """Return current device."""
        return self._device

    def to(self, device):
        """Move DataLoader to device."""
        self._device = device
        return self

    def one_batch(self):
        """Return one batch from the DataLoader.

        Required for fastai compatibility - used for device detection
        and batch shape validation during Learner initialization.

        Returns:
            Tuple of (image, mask) tensors on the correct device.
        """
        return next(iter(self))

    def close(self):
        """Shut down TorchIO Queue workers. Safe to call multiple times."""
        if self._closed:
            return
        self._closed = True
        try:
            # Trigger cleanup of Queue's internal DataLoader iterator
            if hasattr(self, 'queue') and hasattr(self.queue, '_subjects_iterable'):
                self.queue._subjects_iterable = None
        except Exception:
            pass  # Suppress cleanup errors

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        return False

    def __del__(self):
        try:
            self.close()
        except Exception:
            pass

# %% ../nbs/10_vision_patch.ipynb #cell-15
class MedPatchDataLoaders:
    """fastai-compatible DataLoaders for patch-based training with LAZY loading.

    This class provides train and validation DataLoaders that work with
    fastai's Learner for patch-based training on 3D medical images.

    Memory-efficient: Volumes are loaded on-demand by Queue workers,
    keeping memory usage constant (~150 MB) regardless of dataset size.

    Note: Validation uses the same sampling as training (pseudo Dice).
    For true validation metrics, use PatchInferenceEngine with GridSampler
    for full-volume sliding window inference.

    Example:
        >>> import torchio as tio
        >>>
        >>> # New pattern: preprocessing params in config (DRY)
        >>> config = PatchConfig(
        ...     patch_size=[96, 96, 96],
        ...     apply_reorder=True,
        ...     target_spacing=[0.5, 0.5, 0.5]
        ... )
        >>> dls = MedPatchDataLoaders.from_df(
        ...     df, img_col='image', mask_col='label',
        ...     valid_pct=0.2,
        ...     patch_config=config,
        ...     pre_patch_tfms=[tio.ZNormalization()],
        ...     bs=4
        ... )
        >>> learn = Learner(dls, model, loss_func=DiceLoss())
    """

    def __init__(
        self,
        train_dl: MedPatchDataLoader,
        valid_dl: MedPatchDataLoader,
        device: torch.device = None
    ):
        self._train_dl = train_dl
        self._valid_dl = valid_dl
        self._device = device or _get_default_device()

        # Move to device
        self._train_dl.to(self._device)
        self._valid_dl.to(self._device)

        # Track cleanup state
        self._closed = False

    @classmethod
    def from_df(
        cls,
        df: pd.DataFrame,
        img_col: str,
        mask_col: str = None,
        valid_pct: float = 0.2,
        valid_col: str = None,
        patch_config: PatchConfig = None,
        pre_patch_tfms: list = None,
        patch_tfms: list = None,
        gpu_augmentation=None,
        apply_reorder: bool = None,
        target_spacing: list = None,
        bs: int = 4,
        seed: int = None,
        device: torch.device = None,
        ensure_affine_consistency: bool = True
    ) -> 'MedPatchDataLoaders':
        """Create train/valid DataLoaders from DataFrame with LAZY loading.

        Memory-efficient: Only file paths are stored at creation time.
        Volumes are loaded on-demand by Queue workers during training.

        Note: Both train and valid use the same sampling strategy from patch_config.
        This gives pseudo Dice during training. For true validation metrics,
        use PatchInferenceEngine with full-volume sliding window inference.

        Args:
            df: DataFrame with image paths.
            img_col: Column name for image paths.
            mask_col: Column name for mask paths.
            valid_pct: Fraction of data for validation.
            valid_col: Column name for train/valid split (if pre-defined).
            patch_config: PatchConfig instance. Preprocessing params (apply_reorder,
                target_spacing) can be set here for DRY usage with PatchInferenceEngine.
            pre_patch_tfms: TorchIO transforms applied before patch extraction
                           (after reorder/resample). Example: [tio.ZNormalization()].
                           Accepts both fastMONAI wrappers and raw TorchIO transforms.
                           Skipped when preprocessed=True (include in preprocess_dataset()
                           transforms instead). Still needed for inference via pre_inference_tfms.
            patch_tfms: TorchIO transforms applied to extracted patches (training only).
                Mutually exclusive with gpu_augmentation.
            gpu_augmentation: GpuPatchAugmentation instance for GPU-batched augmentation
                (training only). Mutually exclusive with patch_tfms.
            apply_reorder: If True, reorder to RAS+ orientation. If None, uses
                patch_config.apply_reorder. Explicit value overrides config.
            target_spacing: Target voxel spacing [x, y, z]. If None, uses
                patch_config.target_spacing. Explicit value overrides config.
            bs: Batch size.
            seed: Random seed for splitting.
            device: Device to use.
            ensure_affine_consistency: If True and mask_col is provided, automatically
                adds tio.CopyAffine(target='image') as the first transform to prevent
                spatial metadata mismatch errors. Defaults to True.

        Returns:
            MedPatchDataLoaders instance.

        Example:
            >>> # CPU augmentation path (existing)
            >>> dls = MedPatchDataLoaders.from_df(
            ...     df, img_col='image', mask_col='label',
            ...     patch_config=config,
            ...     patch_tfms=[tio.RandomAffine(degrees=10), tio.RandomFlip()],
            ...     bs=4
            ... )
            >>>
            >>> # GPU augmentation path (new, faster for long training runs)
            >>> from fastMONAI.vision_augmentation import gpu_patch_augmentations
            >>> gpu_aug = gpu_patch_augmentations(config.patch_size, config.target_spacing)
            >>> dls = MedPatchDataLoaders.from_df(
            ...     df, img_col='image', mask_col='label',
            ...     patch_config=config,
            ...     gpu_augmentation=gpu_aug,
            ...     bs=4
            ... )
        """
        # Validate mutual exclusivity
        if gpu_augmentation is not None and patch_tfms is not None:
            raise ValueError(
                "Cannot use both gpu_augmentation and patch_tfms. "
                "gpu_augmentation operates on GPU tensors batch-wise, while "
                "patch_tfms uses per-sample CPU TorchIO transforms. Choose one."
            )

        if patch_config is None:
            patch_config = PatchConfig()

        # Use config values, allow explicit overrides for backward compatibility
        _apply_reorder = apply_reorder if apply_reorder is not None else patch_config.apply_reorder
        _target_spacing = target_spacing if target_spacing is not None else patch_config.target_spacing

        # Warn if both config and explicit args provided with different values
        _warn_config_override('apply_reorder', patch_config.apply_reorder, apply_reorder)
        _warn_config_override('target_spacing', patch_config.target_spacing, target_spacing)

        # Split data
        if valid_col is not None:
            train_df = df[df[valid_col] == False].reset_index(drop=True)
            valid_df = df[df[valid_col] == True].reset_index(drop=True)
        else:
            if seed is not None:
                np.random.seed(seed)
            n = len(df)
            valid_idx = np.random.choice(n, size=int(n * valid_pct), replace=False)
            train_idx = np.setdiff1d(np.arange(n), valid_idx)
            train_df = df.iloc[train_idx].reset_index(drop=True)
            valid_df = df.iloc[valid_idx].reset_index(drop=True)

        # Build preprocessing transforms
        all_pre_tfms = []

        # Skip all preprocessing if data was already preprocessed externally
        if not patch_config.preprocessed:
            # Add reorder transform (reorder to RAS+ orientation)
            if _apply_reorder:
                all_pre_tfms.append(tio.ToCanonical())

            # Add resample transform
            if _target_spacing is not None:
                all_pre_tfms.append(tio.Resample(_target_spacing))

            # Add user-provided transforms (normalize to raw TorchIO transforms)
            if pre_patch_tfms:
                all_pre_tfms.extend(normalize_patch_transforms(pre_patch_tfms))

        # Create subjects datasets with lazy loading (paths only, ~0 MB)
        train_subjects = create_subjects_dataset(
            train_df, img_col, mask_col,
            pre_tfms=all_pre_tfms if all_pre_tfms else None,
            ensure_affine_consistency=ensure_affine_consistency
        )
        valid_subjects = create_subjects_dataset(
            valid_df, img_col, mask_col,
            pre_tfms=all_pre_tfms if all_pre_tfms else None,
            ensure_affine_consistency=ensure_affine_consistency
        )

        # Create DataLoaders (both use same patch_config for consistent sampling)
        train_dl = MedPatchDataLoader(
            train_subjects, patch_config, bs,
            patch_tfms=patch_tfms,
            gpu_augmentation=gpu_augmentation,
            shuffle=True, drop_last=True
        )
        valid_dl = MedPatchDataLoader(
            valid_subjects, patch_config, bs,
            patch_tfms=None,  # No augmentation for validation
            gpu_augmentation=None,  # No augmentation for validation
            shuffle=False, drop_last=False
        )

        # Create instance and store metadata
        instance = cls(train_dl, valid_dl, device)
        instance._img_col = img_col
        instance._mask_col = mask_col
        instance._pre_patch_tfms = pre_patch_tfms
        instance._apply_reorder = _apply_reorder
        instance._target_spacing = _target_spacing
        instance._ensure_affine_consistency = ensure_affine_consistency
        instance._patch_config = patch_config
        instance._train_source_df = train_df
        instance._valid_source_df = valid_df
        return instance

    @property
    def train(self):
        """Training DataLoader."""
        return self._train_dl

    @property
    def valid(self):
        """Validation DataLoader."""
        return self._valid_dl

    @property
    def train_ds(self):
        """Training subjects dataset."""
        return self._train_dl.subjects_dataset

    @property
    def valid_ds(self):
        """Validation subjects dataset."""
        return self._valid_dl.subjects_dataset

    @property
    def device(self):
        """Current device."""
        return self._device

    @property
    def bs(self):
        """Batch size."""
        return self._train_dl.bs

    @property
    def apply_reorder(self):
        """Whether reordering to RAS+ is enabled."""
        return getattr(self, '_apply_reorder', False)

    @property
    def target_spacing(self):
        """Target voxel spacing for resampling."""
        return getattr(self, '_target_spacing', None)

    @property
    def patch_config(self):
        """The PatchConfig used for this DataLoaders."""
        return getattr(self, '_patch_config', None)

    @property
    def split_df(self):
        """DataFrame recording train/valid split for reproducibility logging."""
        train = self._train_source_df.assign(is_valid=False)
        valid = self._valid_source_df.assign(is_valid=True)
        return pd.concat([train, valid], ignore_index=True)

    def to(self, device):
        """Move DataLoaders to device."""
        self._device = device
        self._train_dl.to(device)
        self._valid_dl.to(device)
        return self

    def __iter__(self):
        """Iterate over training DataLoader."""
        return iter(self._train_dl)

    def one_batch(self):
        """Return one batch from the training DataLoader.

        Required for fastai Learner compatibility - used for device
        detection and batch shape validation.
        """
        return self._train_dl.one_batch()

    def __len__(self):
        """Return number of batches in training DataLoader."""
        return len(self._train_dl)

    def __getitem__(self, idx):
        """Get DataLoader by index. Required for fastai Learner compatibility.

        Args:
            idx: 0 for training DataLoader, 1 for validation DataLoader.

        Returns:
            MedPatchDataLoader instance.
        """
        if idx == 0:
            return self._train_dl
        elif idx == 1:
            return self._valid_dl
        else:
            raise IndexError(f"Index {idx} out of range. Use 0 (train) or 1 (valid).")

    def cuda(self):
        """Move DataLoaders to CUDA device."""
        return self.to(torch.device('cuda'))

    def cpu(self):
        """Move DataLoaders to CPU."""
        return self.to(torch.device('cpu'))

    def show_batch(self, dl_idx=0, max_n=6, figsize=None, channel=0,
                   slice_index=None, anatomical_plane=0, overlay=False,
                   voxel_size=None, **kwargs):
        """Show a batch of patch samples for visualization."""

        dl = self[dl_idx]
        x, y = dl.one_batch()
        x = x.cpu()
        if y is not None: y = y.cpu()

        nrows = min(x.shape[0], max_n)
        has_mask = y is not None

        if overlay and has_mask:
            ncols = x.shape[1]
        else:
            ncols = x.shape[1] + (1 if has_mask else 0)

        if figsize is None:
            figsize = (ncols * 3, nrows * 3)
        fig, axs = plt.subplots(nrows, ncols, figsize=figsize, squeeze=False)
        flat_axs = axs.flatten()

        imgs, masks_for_overlay, slice_idxs = [], [], []
        for i in range(nrows):
            img = x[i]
            im_channels = [MedImage(c_img[None]) for c_img in img]

            if has_mask:
                mask = y[i]
                idx = find_max_slice(mask[0].numpy(), anatomical_plane) if slice_index is None else slice_index
                if overlay:
                    masks_for_overlay.extend([MedMask(mask)] * len(im_channels))
                else:
                    im_channels.append(MedMask(mask))
            else:
                idx = slice_index

            imgs.extend(im_channels)
            slice_idxs.extend([idx] * len(im_channels))

        _voxel_size = voxel_size if voxel_size is not None else self.target_spacing
        ctxs = [im.show(ax=ax, slice_index=idx, anatomical_plane=anatomical_plane,
                        voxel_size=_voxel_size)
                for im, ax, idx in zip(imgs, flat_axs, slice_idxs)]

        if overlay and has_mask:
            for mask, ax, idx in zip(masks_for_overlay, flat_axs, slice_idxs):
                mask.show(ax=ax, slice_index=idx, anatomical_plane=anatomical_plane,
                          voxel_size=_voxel_size)

        plt.tight_layout()
        plt.show()

    def new_empty(self):
        """Create a new empty version of self for learner export.

        Required for fastai Learner.export() compatibility - creates a
        lightweight placeholder that can be pickled without the full dataset.

        Returns:
            A minimal MedPatchDataLoaders-like object with no data.
        """
        class EmptyMedPatchDataLoaders:
            """Minimal placeholder for exported learner."""
            def __init__(self, device):
                self._device = device
            @property
            def device(self): return self._device
            def to(self, device):
                self._device = device
                return self
            def cpu(self):
                """Move to CPU. Required for load_learner compatibility."""
                return self.to(torch.device('cpu'))

        return EmptyMedPatchDataLoaders(self._device)

    def close(self):
        """Shut down all DataLoader workers. Safe to call multiple times."""
        if self._closed:
            return
        self._closed = True
        if hasattr(self, '_train_dl') and self._train_dl is not None:
            self._train_dl.close()
        if hasattr(self, '_valid_dl') and self._valid_dl is not None:
            self._valid_dl.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        return False

    def __del__(self):
        try:
            self.close()
        except Exception:
            pass

# %% ../nbs/10_vision_patch.ipynb #cell-17
import numbers

def _normalize_patch_overlap(patch_overlap, patch_size):
    """Convert patch_overlap to integer pixel values for TorchIO compatibility.

    TorchIO's GridSampler expects patch_overlap as a tuple of even integers.
    This function handles:
    - Fractional overlap (0-1): converted to pixel values based on patch_size
    - Numpy scalar types: converted to native Python types
    - Sequences: converted to tuple of integers

    Note: Input validation (negative values, overlap >= patch_size) is handled
    by PatchConfig.__post_init__(). This function focuses on format conversion.

    Args:
        patch_overlap: int, float (0-1 for fraction), or sequence
        patch_size: list/tuple of patch dimensions [x, y, z]

    Returns:
        Tuple of even integers suitable for TorchIO GridSampler
    """
    # Handle scalar fractional overlap (0 < x < 1)
    # Note: excludes 1.0 as 100% overlap creates step_size=0 (infinite patches)
    if isinstance(patch_overlap, (int, float, numbers.Number)) and 0 < float(patch_overlap) < 1:
        # Convert fraction to pixel values, ensure even
        result = []
        for ps in patch_size:
            pixels = int(int(ps) * float(patch_overlap))
            # Ensure even (required by TorchIO)
            if pixels % 2 != 0:
                pixels = pixels - 1 if pixels > 0 else 0
            result.append(pixels)
        return tuple(result)

    # Handle scalar integer (including numpy scalars) - values > 1 are pixel counts
    if isinstance(patch_overlap, (int, float, numbers.Number)):
        val = int(patch_overlap)
        # Ensure even
        if val % 2 != 0:
            val = val - 1 if val > 0 else 0
        return tuple(val for _ in patch_size)

    # Handle sequences (list, tuple, ndarray)
    result = []
    for val in patch_overlap:
        pixels = int(val)
        if pixels % 2 != 0:
            pixels = pixels - 1 if pixels > 0 else 0
        result.append(pixels)
    return tuple(result)


# nnU-Net-style mirror TTA: all 2^3 = 8 flip combinations for 3D.
# Batch tensor shape: [B, C, D, H, W], spatial dims are 2, 3, 4.
_TTA_FLIP_AXES = (
    (),         # original
    (4,),       # flip LR (W)
    (3,),       # flip AP (H)
    (2,),       # flip IS (D)
    (3, 4),     # flip LR+AP
    (2, 4),     # flip LR+IS
    (2, 3),     # flip AP+IS
    (2, 3, 4),  # flip all
)


def _predict_patch_tta(model, patch_input):
    """nnU-Net-style mirror TTA: average probabilities over 8 flip combinations.

    Runs 8 forward passes with a running sum for memory efficiency (2x memory,
    not 9x). Each pass: flip input -> forward -> activate -> flip back -> accumulate.

    Args:
        model: PyTorch model in eval mode (already on device).
        patch_input: Batch tensor [B, C, D, H, W] already on device.

    Returns:
        Averaged probability tensor [B, C, D, H, W] on CPU.
    """
    summed_probs = None
    for axes in _TTA_FLIP_AXES:
        flipped = torch.flip(patch_input, list(axes)) if axes else patch_input
        logits = model(flipped)
        n_classes = logits.shape[1]
        probs = torch.sigmoid(logits) if n_classes == 1 else torch.softmax(logits, dim=1)
        if axes:
            probs = torch.flip(probs, list(axes))
        summed_probs = probs if summed_probs is None else summed_probs + probs
    return (summed_probs / len(_TTA_FLIP_AXES)).cpu()


class PatchInferenceEngine:
    """Patch-based inference with automatic volume reconstruction.
    
    Uses TorchIO's GridSampler to extract overlapping patches and
    GridAggregator to reconstruct the full volume from predictions.
    
    Args:
        learner: fastai Learner or PyTorch model (nn.Module). When passing a raw
            PyTorch model, load weights first with model.load_state_dict().
        config: PatchConfig with inference settings. Preprocessing params (apply_reorder,
            target_spacing, padding_mode) can be set here for DRY usage.
        apply_reorder: Whether to reorder to RAS+ orientation. If None, uses config value.
        target_spacing: Target voxel spacing. If None, uses config value.
        batch_size: Number of patches to predict at once. Must be positive.
        pre_inference_tfms: List of TorchIO transforms to apply before patch extraction.
            IMPORTANT: Should match the pre_patch_tfms used during training (e.g., [tio.ZNormalization()]).
            This ensures preprocessing consistency between training and inference.
            Accepts both fastMONAI wrappers and raw TorchIO transforms.
    
    Example:
        >>> # Option 1: From fastai Learner
        >>> engine = PatchInferenceEngine(learn, config, pre_inference_tfms=[ZNormalization()])
        >>> pred = engine.predict('image.nii.gz')
        
        >>> # Option 2: From raw PyTorch model (recommended for deployment)
        >>> model = UNet(spatial_dims=3, in_channels=1, out_channels=2, ...)
        >>> model.load_state_dict(torch.load('final_weights.pth'))
        >>> model.cuda().eval()
        >>> engine = PatchInferenceEngine(model, config, pre_inference_tfms=[ZNormalization()])
        >>> pred = engine.predict('image.nii.gz')
    """
    
    def __init__(
        self,
        learner,
        config: PatchConfig,
        apply_reorder: bool = None,
        target_spacing: list = None,
        batch_size: int = 4,
        pre_inference_tfms: list = None
    ):
        if batch_size <= 0:
            raise ValueError(f"batch_size must be positive, got {batch_size}")
        
        # Extract model from Learner if needed (use isinstance for robust detection)
        # Note: We check for Learner explicitly because some models (e.g., MONAI UNet)
        # have a .model attribute that is NOT the full model but an internal Sequential.
        if isinstance(learner, Learner):
            self.model = learner.model
        else:
            self.model = learner  # Assume it's already a PyTorch model
        
        self.config = config
        self.batch_size = batch_size
        
        # Normalize transforms to raw TorchIO (accepts both fastMONAI wrappers and raw TorchIO)
        normalized_tfms = normalize_patch_transforms(pre_inference_tfms)
        self.pre_inference_tfms = tio.Compose(normalized_tfms) if normalized_tfms else None
        
        # Use config values, allow explicit overrides for backward compatibility
        self.apply_reorder = apply_reorder if apply_reorder is not None else config.apply_reorder
        self.target_spacing = target_spacing if target_spacing is not None else config.target_spacing
        
        # Warn if explicit args provided but differ from config (potential mistake)
        _warn_config_override('apply_reorder', config.apply_reorder, apply_reorder)
        _warn_config_override('target_spacing', config.target_spacing, target_spacing)
        
        # Get device from model parameters, with fallback for parameter-less models
        try:
            self._device = next(self.model.parameters()).device
        except StopIteration:
            self._device = _get_default_device()
    
    def predict(
        self,
        img_path: Path | str,
        return_probabilities: bool = False,
        return_affine: bool = False,
        tta: bool = False
    ) -> torch.Tensor | tuple[torch.Tensor, np.ndarray]:
        """Predict on a single volume using patch-based inference.

        Args:
            img_path: Path to input image.
            return_probabilities: If True, return probability map instead of argmax.
            return_affine: If True, return (prediction, affine) tuple instead of just prediction.
            tta: If True, apply nnU-Net-style mirror test-time augmentation
                (8 flip combinations, averaged probabilities). Requires ~8x inference
                time but improves prediction quality. Works best when training used
                RandomFlip(axes='LRAPIS', p=0.5). Defaults to False.

        Returns:
            Predicted segmentation mask tensor, or tuple (prediction, affine) if return_affine=True.
        """
        # Load image - keep org_img and org_size for post-processing
        # Note: med_img_reader handles reorder/resample internally, no global state needed
        org_img, input_img, org_size = med_img_reader(
            img_path, apply_reorder=self.apply_reorder, target_spacing=self.target_spacing, only_tensor=False
        )

        # Create TorchIO Subject from preprocessed image
        subject = tio.Subject(
            image=tio.ScalarImage(tensor=input_img.data.float(), affine=input_img.affine)
        )

        # Apply pre-inference transforms (e.g., ZNormalization) to match training
        if self.pre_inference_tfms is not None:
            subject = self.pre_inference_tfms(subject)

        # Pad dimensions smaller than patch_size, keep larger dimensions intact
        # GridSampler handles large images via overlapping patches
        img_shape = subject['image'].shape[1:]  # Exclude channel dim
        target_size = [max(s, p) for s, p in zip(img_shape, self.config.patch_size)]
        
        # Warn if volume needed padding (may cause artifacts if training didn't cover similar sizes)
        if any(s < p for s, p in zip(img_shape, self.config.patch_size)):
            padded_dims = [f"dim{i}: {s}<{p}" for i, (s, p) in enumerate(zip(img_shape, self.config.patch_size)) if s < p]
            warnings.warn(
                f"Image size {list(img_shape)} smaller than patch_size {self.config.patch_size} "
                f"in {padded_dims}. Padding with mode={self.config.padding_mode}. "
                "Ensure training data covered similar sizes to avoid artifacts."
            )
        
        # Use padding_mode from config (default: 0 for zero padding, nnU-Net standard)
        subject = tio.CropOrPad(target_size, padding_mode=self.config.padding_mode)(subject)

        # Convert patch_overlap to integer pixel values for TorchIO compatibility
        patch_overlap = _normalize_patch_overlap(self.config.patch_overlap, self.config.patch_size)

        # Create GridSampler
        grid_sampler = tio.GridSampler(
            subject,
            patch_size=self.config.patch_size,
            patch_overlap=patch_overlap
        )

        # Create GridAggregator
        aggregator = tio.GridAggregator(
            grid_sampler,
            overlap_mode=self.config.aggregation_mode
        )

        # Create patch loader
        patch_loader = DataLoader(
            grid_sampler,
            batch_size=self.batch_size,
            num_workers=0
        )

        # Predict patches
        self.model.eval()
        with torch.no_grad():
            for patches_batch in patch_loader:
                patch_input = patches_batch['image'][tio.DATA].to(self._device)
                locations = patches_batch[tio.LOCATION]

                if tta:
                    probs = _predict_patch_tta(self.model, patch_input)
                else:
                    # Forward pass - get logits
                    logits = self.model(patch_input)

                    # Convert logits to probabilities BEFORE aggregation
                    # This is critical: softmax is non-linear, so we must aggregate
                    # probabilities, not logits, to get correct boundary handling
                    n_classes = logits.shape[1]
                    if n_classes == 1:
                        probs = torch.sigmoid(logits)
                    else:
                        probs = torch.softmax(logits, dim=1)  # dim=1 for batch [B, C, H, W, D]

                    probs = probs.cpu()

                # Add probabilities to aggregator
                aggregator.add_batch(probs, locations)

        # Get reconstructed output (now contains probabilities, not logits)
        output = aggregator.get_output_tensor()

        # Convert to prediction mask (only if not returning probabilities)
        if return_probabilities:
            result = output  # Keep as float probabilities
        else:
            n_classes = output.shape[0]
            if n_classes == 1:
                result = (output > 0.5).float()
            else:
                result = output.argmax(dim=0, keepdim=True).float()

        # Apply keep_largest post-processing for binary segmentation
        if not return_probabilities and self.config.keep_largest_component:
            from fastMONAI.vision_inference import keep_largest
            result = keep_largest(result.squeeze(0)).unsqueeze(0)

        # Post-processing: resize back to original size and reorient
        # This matches the workflow in vision_inference.py
        
        # Wrap result in TorchIO Image for resizing
        # Use ScalarImage for probabilities, LabelMap for masks
        if return_probabilities:
            pred_img = tio.ScalarImage(tensor=result.float(), affine=input_img.affine)
        else:
            pred_img = tio.LabelMap(tensor=result.float(), affine=input_img.affine)
        
        # Resize back to original size (before resampling)
        pred_img = _do_resize(pred_img, org_size, image_interpolation='nearest')
        
        # Reorient to original orientation (if reorder was applied)
        # Use explicit .cpu() for consistent device handling
        if self.apply_reorder:
            reoriented_array = _to_original_orientation(
                pred_img.as_sitk(),
                ('').join(org_img.orientation)
            )
            result = torch.from_numpy(reoriented_array).cpu()
            # Only convert to long for masks, not probabilities
            if not return_probabilities:
                result = result.long()
        else:
            result = pred_img.data.cpu()
            # Only convert to long for masks, not probabilities
            if not return_probabilities:
                result = result.long()

        # Use original affine matrix for correct spatial alignment
        # org_img.affine is always available from med_img_reader
        if not (hasattr(org_img, 'affine') and org_img.affine is not None):
            raise RuntimeError(
                "org_img.affine not available. This should never happen - please report this bug."
            )
        affine = org_img.affine.copy()

        if return_affine:
            return result, affine
        return result
    
    def to(self, device):
        """Move engine to device."""
        self._device = device
        self.model.to(device)
        return self

# %% ../nbs/10_vision_patch.ipynb #cell-18
def patch_inference(
    learner,
    config: PatchConfig,
    file_paths: list,
    apply_reorder: bool = None,
    target_spacing: list = None,
    batch_size: int = 4,
    return_probabilities: bool = False,
    progress: bool = True,
    save_dir: str = None,
    pre_inference_tfms: list = None,
    tta: bool = False
) -> list:
    """Batch patch-based inference on multiple volumes.
    
    Args:
        learner: PyTorch model or fastai Learner.
        config: PatchConfig with inference settings. Preprocessing params (apply_reorder,
            target_spacing) can be set here for DRY usage.
        file_paths: List of image paths.
        apply_reorder: Whether to reorder to RAS+ orientation. If None, uses config value.
        target_spacing: Target voxel spacing. If None, uses config value.
        batch_size: Patches per batch.
        return_probabilities: Return probability maps.
        progress: Show progress bar.
        save_dir: Directory to save predictions as NIfTI files. If None, predictions are not saved.
        pre_inference_tfms: List of TorchIO transforms to apply before patch extraction.
            IMPORTANT: Should match the pre_patch_tfms used during training (e.g., [tio.ZNormalization()]).
        tta: If True, apply nnU-Net-style mirror TTA (8 flip combinations).
    
    Returns:
        List of predicted tensors.
    
    Example:
        >>> config = PatchConfig(
        ...     patch_size=[96, 96, 96],
        ...     apply_reorder=True,
        ...     target_spacing=[0.4102, 0.4102, 1.5]
        ... )
        >>> predictions = patch_inference(
        ...     learner=learn,
        ...     config=config,  # apply_reorder and target_spacing from config
        ...     file_paths=val_paths,
        ...     pre_inference_tfms=[tio.ZNormalization()],
        ...     save_dir='predictions/patch_based'
        ... )
    """
    # Use config values if not explicitly provided
    _apply_reorder = apply_reorder if apply_reorder is not None else config.apply_reorder
    _target_spacing = target_spacing if target_spacing is not None else config.target_spacing
    
    engine = PatchInferenceEngine(
        learner, config, _apply_reorder, _target_spacing, batch_size, pre_inference_tfms
    )
    
    # Create save directory if specified
    if save_dir is not None:
        save_path = Path(save_dir)
        save_path.mkdir(parents=True, exist_ok=True)
    
    predictions = []
    desc = 'Patch inference (TTA)' if tta else 'Patch inference'
    iterator = tqdm(file_paths, desc=desc) if progress else file_paths
    
    for path in iterator:
        # Get prediction and affine when saving is needed
        if save_dir is not None:
            pred, affine = engine.predict(path, return_probabilities, return_affine=True, tta=tta)
        else:
            pred = engine.predict(path, return_probabilities, tta=tta)
        predictions.append(pred)
        
        # Save prediction if save_dir specified
        if save_dir is not None:
            input_path = Path(path)
            # Create output filename based on input using suffix-based approach
            # This handles .nii.gz correctly without corrupting filenames with .nii elsewhere
            stem = input_path.stem
            if input_path.suffix == '.gz' and stem.endswith('.nii'):
                # Handle .nii.gz files: stem is "filename.nii", strip the .nii
                stem = stem[:-4]
                out_name = f"{stem}_pred.nii.gz"
            elif input_path.suffix == '.nii':
                # Handle .nii files
                out_name = f"{stem}_pred.nii"
            else:
                # Fallback for other formats
                out_name = f"{stem}_pred.nii.gz"
            out_path = save_path / out_name
            
            # affine is guaranteed to be valid from engine.predict() with return_affine=True
            # Save as NIfTI using TorchIO with correct type
            # Use ScalarImage for probabilities (float), LabelMap for masks (int)
            if return_probabilities:
                pred_img = tio.ScalarImage(tensor=pred, affine=affine)
            else:
                pred_img = tio.LabelMap(tensor=pred, affine=affine)
            pred_img.save(out_path)
    
    return predictions
