# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/08_dataset_info.ipynb.

# %% auto #0
__all__ = ['MedDataset', 'suggest_patch_size', 'preprocess_dataset', 'get_class_weights']

# %% ../nbs/08_dataset_info.ipynb #027f016a-a80c-4842-b9dc-0bddb358a00c
from .vision_core import *
from .vision_plot import find_max_slice

from sklearn.utils.class_weight import compute_class_weight
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm.auto import tqdm
from pathlib import Path
import torchio as tio
import pandas as pd
import numpy as np
import torch
import glob
import hashlib
import os
import pickle
import matplotlib.pyplot as plt

# %% ../nbs/08_dataset_info.ipynb #7401beac
import warnings

class MedDataset:
    """A class to extract and present information about the dataset."""

    def __init__(self, dataframe=None, image_col:str=None, mask_col:str="mask_path",
                 path=None, img_list=None, postfix:str='', apply_reorder:bool=True,
                 dtype:(MedImage, MedMask)=MedImage, max_workers:int=1,
                 use_cache:bool=True, cache_path=None):
        """Constructs MedDataset object.

        Args:
            dataframe: DataFrame containing image paths.
            image_col: Column name for image paths (used for visualization).
            mask_col: Column name for mask/label paths when using dataframe mode.
            path: Directory path containing images.
            img_list: List of image file paths to analyze.
            postfix: File postfix filter when using path mode.
            apply_reorder: Whether to reorder images to RAS+ orientation.
            dtype: MedImage for images or MedMask for segmentation masks.
            max_workers: Number of parallel workers for processing.
            use_cache: Enable metadata caching. When True (default) and cache_path
                is not provided, auto-generates a cache path in ~/.cache/fastmonai/.
                Set to False to disable caching entirely.
            cache_path: Explicit path to a pickle file for metadata caching.
                Only used when use_cache=True. If None and use_cache=True,
                a path is auto-generated from the file list and config.
        """
        self.input_df = dataframe
        self.image_col = image_col
        self.mask_col = mask_col
        self.path = path
        self.img_list = img_list
        self.postfix = postfix
        self.apply_reorder = apply_reorder
        self.dtype = dtype
        self.max_workers = max_workers
        self.use_cache = use_cache
        self.cache_path = cache_path
        self.df = self._create_data_frame()

    def _create_data_frame(self):
        """Private method that returns a dataframe with information about the dataset."""

        # Handle img_list (simple list of paths)
        if self.img_list is not None:
            file_list = self.img_list

        # Handle path-based initialization
        elif self.path:
            file_list = glob.glob(f'{self.path}/*{self.postfix}*')
            if not file_list:
                print('Could not find images. Check the image path')
                return pd.DataFrame()

        # Handle dataframe-based initialization
        elif self.input_df is not None and self.mask_col in self.input_df.columns:
            file_list = self.input_df[self.mask_col].tolist()

        else:
            print('Error: Must provide path, img_list, or dataframe with mask_col')
            return pd.DataFrame()

        # Resolve cache path
        if self.use_cache and self.cache_path is None:
            self.cache_path = self._auto_cache_path(file_list)

        # Process images to extract metadata (with optional caching)
        if self.use_cache and self.cache_path is not None:
            data_info_dict = self._process_with_cache(file_list)
        else:
            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                data_info_dict = list(executor.map(self._get_data_info, file_list))

        df = pd.DataFrame(data_info_dict)

        if len(df) > 0 and df.orientation.nunique() > 1 and not self.apply_reorder:
            raise ValueError(
                'Mixed orientations detected in dataset. '
                'Please recreate MedDataset with apply_reorder=True to get correct resample values: '
                'MedDataset(..., apply_reorder=True)'
            )

        return df

    def _auto_cache_path(self, file_list):
        """Generate automatic cache path from file list and config."""
        try:
            cache_dir = Path.home() / '.cache' / 'fastmonai'
            cache_dir.mkdir(parents=True, exist_ok=True)
            config_str = f"reorder={self.apply_reorder}|dtype={self.dtype.__name__}"
            abs_paths = sorted(os.path.abspath(fn) for fn in file_list)
            key_input = config_str + '|' + '|'.join(abs_paths)
            key = hashlib.md5(key_input.encode()).hexdigest()
            return str(cache_dir / f'med_dataset_{key}.pkl')
        except Exception:
            return None

    @property
    def fingerprint(self):
        """Compute a dataset fingerprint from per-file content hashes.

        Returns a deterministic MD5 hex digest representing the dataset contents.
        Returns None if content hashes are not available (e.g., all files failed).
        """
        if 'content_hash' not in self.df.columns:
            return None
        hashes = self.df['content_hash'].dropna().sort_values().tolist()
        if not hashes:
            return None
        combined = ''.join(hashes)
        return hashlib.md5(combined.encode()).hexdigest()

    def summary(self):
        """Summary DataFrame of the dataset with example path for similar data."""

        columns = ['dim_0', 'dim_1', 'dim_2', 'voxel_0', 'voxel_1', 'voxel_2', 'orientation']

        return self.df.groupby(columns, as_index=False).agg(
            example_path=('path', 'min'), total=('path', 'size')
        ).sort_values('total', ascending=False)

    def get_suggestion(self, include_patch_size: bool = False):
        """Returns suggested preprocessing parameters as a dictionary.

        The returned target_spacing is derived from the mode (most common value)
        of voxel spacings in the dataset.

        Note:
            apply_reorder is NOT included in the return value because it is not
            data-derived. Access dataset.apply_reorder directly if needed.

        Args:
            include_patch_size: If True, includes suggested patch_size for
                patch-based training. Requires vision_patch module.

        Returns:
            dict: {'target_spacing': [voxel_0, voxel_1, voxel_2]}
                  If include_patch_size=True, also includes 'patch_size': [dim_0, dim_1, dim_2]
        """
        target_spacing = [float(self.df.voxel_0.mode()[0]), float(self.df.voxel_1.mode()[0]), float(self.df.voxel_2.mode()[0])]
        result = {'target_spacing': target_spacing}

        if include_patch_size:
            result['patch_size'] = suggest_patch_size(self)

        return result

    def _get_data_info(self, fn: str):
        """Private method to collect information about an image file."""
        try:
            _, o, _ = med_img_reader(fn, apply_reorder=self.apply_reorder, only_tensor=False, dtype=self.dtype)

            # Hash loaded tensor data (implicitly captures apply_reorder state)
            content_hash = hashlib.md5(o.data.numpy().tobytes()).hexdigest()

            info_dict = {'path': fn, 'content_hash': content_hash, 'dim_0': o.shape[1], 'dim_1': o.shape[2], 'dim_2': o.shape[3],
                         'voxel_0': round(o.spacing[0], 4), 'voxel_1': round(o.spacing[1], 4), 'voxel_2': round(o.spacing[2], 4),
                         'orientation': f'{"".join(o.orientation)}+'}

            if self.dtype is MedMask:
                # Calculate voxel volume in mm³
                voxel_volume = o.spacing[0] * o.spacing[1] * o.spacing[2]

                # Get voxel counts for each label
                mask_labels_dict = o.count_labels()

                # Calculate volumes for each label > 0 (skip background)
                for key, voxel_count in mask_labels_dict.items():
                    label_int = int(key)
                    if label_int > 0 and voxel_count > 0:  # Skip background (label 0)
                        volume_mm3 = voxel_count * voxel_volume
                        info_dict[f'label_{label_int}_volume_mm3'] = round(volume_mm3, 4)

            return info_dict

        except Exception as e:
            print(f"Warning: Failed to process {fn}: {e}")
            return {'path': fn, 'error': str(e)}

    def _load_cache(self):
        """Load metadata cache from disk. Returns empty dict on any failure."""
        try:
            with open(self.cache_path, 'rb') as f:
                cache = pickle.load(f)
            if not isinstance(cache, dict) or cache.get('version') != 1:
                return {}
            return cache.get('entries', {})
        except Exception:
            return {}

    def _save_cache(self, entries):
        """Save metadata cache to disk."""
        try:
            cache = {'version': 1, 'entries': entries}
            with open(self.cache_path, 'wb') as f:
                pickle.dump(cache, f)
        except Exception as e:
            print(f"Warning: Failed to save metadata cache: {e}")

    def _process_with_cache(self, file_list):
        """Process file list with per-file metadata caching.

        Loads existing cache, identifies files that need reprocessing
        (cache misses), processes only those files, updates and saves
        the cache, then returns all results in original order.
        """
        entries = self._load_cache()

        cached_results = []
        files_to_process = []
        process_indices = []

        for i, fn in enumerate(file_list):
            abs_path = os.path.abspath(fn)
            entry = entries.get(abs_path)
            if entry is not None:
                try:
                    stat = os.stat(fn)
                    if (entry.get('mtime') == stat.st_mtime
                            and entry.get('size') == stat.st_size
                            and entry.get('apply_reorder') == self.apply_reorder
                            and entry.get('dtype') == self.dtype.__name__):
                        cached_results.append((i, entry['info']))
                        continue
                except OSError:
                    pass
            files_to_process.append(fn)
            process_indices.append(i)

        # Process cache misses
        if files_to_process:
            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                fresh_results = list(executor.map(self._get_data_info, files_to_process))
        else:
            fresh_results = []

        # Update cache entries with fresh results (skip errors)
        for fn, info in zip(files_to_process, fresh_results):
            if 'error' not in info:
                abs_path = os.path.abspath(fn)
                try:
                    stat = os.stat(fn)
                    entries[abs_path] = {
                        'mtime': stat.st_mtime,
                        'size': stat.st_size,
                        'apply_reorder': self.apply_reorder,
                        'dtype': self.dtype.__name__,
                        'info': info
                    }
                except OSError:
                    pass

        self._save_cache(entries)

        # Reconstruct results in original file_list order
        all_results = [None] * len(file_list)
        for i, info in cached_results:
            all_results[i] = info
        for i, info in zip(process_indices, fresh_results):
            all_results[i] = info

        n_cached = len(cached_results)
        n_processed = len(files_to_process)
        if n_cached > 0:
            print(f"MedDataset cache: {n_cached} cached, {n_processed} processed")
        elif n_processed > 0:
            print(f"MedDataset: processed {n_processed} files (results cached)")

        return all_results

    def calculate_target_size(self, target_spacing: list = None) -> list:
        """Calculate the target image size for the dataset.

        .. deprecated::
            Use `get_size_statistics(target_spacing)['max']` instead for consistency
            with other size statistics methods.

        Args:
            target_spacing: If provided, calculates size after resampling to this spacing.
                           If None, returns original dimensions.

        Returns:
            list: [dim_0, dim_1, dim_2] largest dimensions in dataset.
        """
        warnings.warn(
            "calculate_target_size() is deprecated. "
            "Use get_size_statistics(target_spacing)['max'] instead.",
            DeprecationWarning,
            stacklevel=2
        )
        if target_spacing is not None:
            org_voxels = self.df[["voxel_0", "voxel_1", 'voxel_2']].values
            org_dims = self.df[["dim_0", "dim_1", 'dim_2']].values

            ratio = org_voxels/target_spacing
            new_dims = (org_dims * ratio).T
            # Use floor() to match TorchIO's Resample dimension calculation
            dims = [float(np.floor(new_dims[0].max())), float(np.floor(new_dims[1].max())), float(np.floor(new_dims[2].max()))]
        else:
            dims = [float(self.df.dim_0.max()), float(self.df.dim_1.max()), float(self.df.dim_2.max())]

        return dims

    def get_size_statistics(self, target_spacing: list = None) -> dict:
        """Calculate comprehensive size statistics for the dataset.

        Args:
            target_spacing: If provided, calculates statistics after
                           simulating resampling to this spacing.

        Returns:
            dict with keys: 'median', 'min', 'max', 'std', 'percentile_10', 'percentile_90'
                  Each value is a list [dim_0, dim_1, dim_2].
        """
        if len(self.df) == 0:
            raise ValueError("Dataset is empty - cannot calculate statistics")

        if target_spacing is not None:
            # Simulate resampled dimensions
            org_voxels = self.df[["voxel_0", "voxel_1", "voxel_2"]].values
            org_dims = self.df[["dim_0", "dim_1", "dim_2"]].values
            ratio = org_voxels / np.array(target_spacing)
            dims = np.floor(org_dims * ratio)
        else:
            dims = self.df[["dim_0", "dim_1", "dim_2"]].values

        return {
            'median': [float(np.median(dims[:, i])) for i in range(3)],
            'min': [float(np.min(dims[:, i])) for i in range(3)],
            'max': [float(np.max(dims[:, i])) for i in range(3)],
            'std': [float(np.std(dims[:, i])) for i in range(3)],
            'percentile_10': [float(np.percentile(dims[:, i], 10)) for i in range(3)],
            'percentile_90': [float(np.percentile(dims[:, i], 90)) for i in range(3)],
        }

    def get_volume_summary(self):
        """Returns DataFrame with volume statistics for each label.

        Returns:
            DataFrame with columns: label, count, mean_mm3, median_mm3, min_mm3, max_mm3
            Returns None if no volume columns found (dtype was not MedMask).
        """
        volume_cols = [col for col in self.df.columns if col.endswith('_volume_mm3')]

        if not volume_cols:
            return None

        summary_data = []
        for col in volume_cols:
            non_zero = self.df[self.df[col] > 0][col]
            if len(non_zero) > 0:
                summary_data.append({
                    'label': col.replace('_volume_mm3', ''),
                    'count': len(non_zero),
                    'mean_mm3': non_zero.mean(),
                    'median_mm3': non_zero.median(),
                    'min_mm3': non_zero.min(),
                    'max_mm3': non_zero.max()
                })

        return pd.DataFrame(summary_data) if summary_data else None

    def _visualize_single_case(self, img_path, mask_path, case_id, anatomical_plane=2, cmap='hot', figsize=(12, 5)):
        """Helper method to visualize a single case."""
        try:
            # Create MedImage and MedMask with current preprocessing settings
            suggestion = self.get_suggestion()
            MedBase.item_preprocessing(target_spacing=suggestion['target_spacing'], apply_reorder=self.apply_reorder)

            img = MedImage.create(img_path)
            mask = MedMask.create(mask_path)

            # Find optimal slice using explicit function
            mask_data = mask.numpy()[0]  # Remove channel dimension
            optimal_slice = find_max_slice(mask_data, anatomical_plane)

            # Create subplot
            fig, axes = plt.subplots(1, 2, figsize=figsize)

            # Show image
            img.show(ctx=axes[0], anatomical_plane=anatomical_plane, slice_index=optimal_slice)
            axes[0].set_title(f"{case_id} - Image (slice {optimal_slice})")

            # Show overlay
            img.show(ctx=axes[1], anatomical_plane=anatomical_plane, slice_index=optimal_slice)
            mask.show(ctx=axes[1], anatomical_plane=anatomical_plane, slice_index=optimal_slice,
                     alpha=0.3, cmap=cmap)
            axes[1].set_title(f"{case_id} - Overlay (slice {optimal_slice})")

            # Adjust spacing to bring plots closer
            plt.subplots_adjust(wspace=0.1)
            plt.tight_layout()
            plt.show()

        except Exception as e:
            print(f"Failed to visualize case {case_id}: {e}")

    def visualize_cases(self, n_cases=4, anatomical_plane=2, cmap='hot', figsize=(12, 5)):
        """Visualize cases from the dataset.

        Args:
            n_cases: Number of cases to show.
            anatomical_plane: 0=sagittal, 1=coronal, 2=axial
            cmap: Colormap for mask overlay
            figsize: Figure size for each case
        """
        if self.input_df is None:
            print("Error: No dataframe provided. Cannot visualize cases.")
            return

        if self.image_col is None:
            print("Error: No image_col specified. Cannot visualize cases.")
            return

        # Check if required columns exist
        if self.image_col not in self.input_df.columns:
            print(f"Error: Column '{self.image_col}' not found in dataframe.")
            return

        if self.mask_col not in self.input_df.columns:
            print(f"Error: Column '{self.mask_col}' not found in dataframe.")
            return

        for idx in range(min(n_cases, len(self.input_df))):
            row = self.input_df.iloc[idx]
            case_id = row.get('case_id', f'Case_{idx}')  # Fallback if no case_id
            img_path = row[self.image_col]
            mask_path = row[self.mask_col]

            self._visualize_single_case(img_path, mask_path, case_id, anatomical_plane, cmap, figsize)
            print("-" * 60)

# %% ../nbs/08_dataset_info.ipynb #wp61bidiwwd
def suggest_patch_size(
    dataset: MedDataset,
    target_spacing: list = None,
    min_patch_size: list = None,
    max_patch_size: list = None,
    divisor: int = 16
) -> list:
    """Suggest optimal patch size based on dataset dimensions.

    Uses median shape as the starting point but clamps to the minimum
    volume size per axis, ensuring the suggested patch fits ALL volumes
    without requiring padding during training.

    Algorithm:
    1. Use min(median, min_volume) per axis for safety
    2. Round down to nearest multiple of divisor (16 for UNet compatibility)
    3. Clamp to [min_patch_size, max_patch_size] bounds
    4. Validate: error if min_patch_size exceeds smallest volume

    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. Default [32, 32, 32].
        max_patch_size: Maximum per dimension. Default [256, 256, 256].
        divisor: Ensure divisibility (default 16 for UNet compatibility).

    Returns:
        list: [patch_dim_0, patch_dim_1, patch_dim_2]

    Example:
        >>> from fastMONAI.dataset_info import MedDataset
        >>> dataset = MedDataset(dataframe=df, mask_col='mask_path', dtype=MedMask)
        >>>
        >>> # Use recommended spacing
        >>> patch_size = suggest_patch_size(dataset)
        >>>
        >>> # Use custom spacing
        >>> patch_size = suggest_patch_size(dataset, target_spacing=[1.0, 1.0, 2.0])
    """
    # Defaults
    min_patch_size = min_patch_size or [32, 32, 32]
    max_patch_size = max_patch_size or [256, 256, 256]

    # Use explicit spacing or get from dataset suggestion
    if target_spacing is None:
        suggestion = dataset.get_suggestion()
        target_spacing = suggestion['target_spacing']

    # Get size statistics (resampled to target_spacing)
    stats = dataset.get_size_statistics(target_spacing)
    median_shape = stats['median']
    min_shape = stats['min']

    # Handle single-image edge case
    if len(dataset.df) == 1:
        warnings.warn("Single image dataset - using image dimensions directly")

    def round_to_divisor(val, div):
        """Round down to nearest multiple of divisor."""
        return max(div, int(val // div) * div)

    # Step 1: Clamp to min(median, min_volume) per axis — safety guarantee
    effective_dims = [min(med, mn) for med, mn in zip(median_shape, min_shape)]

    # Step 2: Round down to nearest divisor
    patch_size = [round_to_divisor(dim, divisor) for dim in effective_dims]

    # Step 3: Clamp to [min_patch_size, max_patch_size] bounds
    patch_size = [
        max(min_p, min(max_p, p))
        for p, min_p, max_p in zip(patch_size, min_patch_size, max_patch_size)
    ]

    # Step 4: Final safety check — error if bounds force patch > min volume
    for i, (p, min_dim) in enumerate(zip(patch_size, min_shape)):
        if p > min_dim:
            raise ValueError(
                f"Cannot suggest safe patch_size for dimension {i}: "
                f"smallest volume is {min_dim:.0f} but min_patch_size={min_patch_size[i]}. "
                f"Exclude small volumes or reduce min_patch_size."
            )

    # Warn when suggestion was reduced to fit smallest volume
    for i, (med, mn, p) in enumerate(zip(median_shape, min_shape, patch_size)):
        median_based = min(max_patch_size[i], round_to_divisor(med, divisor))
        if p < median_based:
            warnings.warn(
                f"Dim {i}: patch_size reduced from {median_based} to {p} "
                f"to fit smallest volume (min={mn:.0f}, median={med:.0f})."
            )

    return patch_size

# %% ../nbs/08_dataset_info.ipynb #mbn5svtmzkh
def preprocess_dataset(df, img_col, mask_col=None, output_dir='preprocessed',
                       target_spacing=None, apply_reorder=True, transforms=None,
                       max_workers=4, skip_existing=True):
    """Preprocess dataset to disk, creating new columns for preprocessed paths.

    Processes images (and optionally masks) through a transform pipeline,
    saves to output_dir, then creates new '{col}_preprocessed' columns in
    the DataFrame. Original columns are preserved unchanged.

    Transform pipeline order:
        CopyAffine (if masks) -> ToCanonical (if apply_reorder)
        -> Resample (if target_spacing) -> user transforms

    Args:
        df: DataFrame with file paths.
        img_col: Column name for image paths.
        mask_col: Optional column name for mask paths.
        output_dir: Output directory. Creates images/ and masks/ subdirectories.
        target_spacing: Target voxel spacing for resampling (e.g., [1.0, 1.0, 1.0]).
        apply_reorder: Whether to reorder to RAS+ canonical orientation.
        transforms: Additional TorchIO or fastMONAI transforms to apply after
            reordering and resampling.
        max_workers: Number of parallel workers. Each worker loads a full 3D
            volume into memory, so reduce for large volumes.
        skip_existing: Skip files that already exist on disk (with size > 0).
    """
    # Input validation
    if len(df) == 0:
        raise ValueError("DataFrame is empty")
    if img_col not in df.columns:
        raise ValueError(f"Column '{img_col}' not found in DataFrame")
    if mask_col is not None and mask_col not in df.columns:
        raise ValueError(f"Column '{mask_col}' not found in DataFrame")

    img_names = [Path(p).name for p in df[img_col]]
    if len(set(img_names)) != len(img_names):
        dupes = set(n for n in img_names if img_names.count(n) > 1)
        raise ValueError(f"Duplicate image file names: {dupes}")

    if mask_col is not None:
        mask_names = [Path(p).name for p in df[mask_col]]
        if len(set(mask_names)) != len(mask_names):
            dupes = set(n for n in mask_names if mask_names.count(n) > 1)
            raise ValueError(f"Duplicate mask file names: {dupes}")

    # Build transform pipeline (canonical order)
    all_tfms = []
    if mask_col is not None:
        all_tfms.append(tio.CopyAffine(target='image'))
    if apply_reorder:
        all_tfms.append(tio.ToCanonical())
    if target_spacing is not None:
        all_tfms.append(tio.Resample(target_spacing))
    if transforms:
        all_tfms.extend([getattr(t, 'tio_transform', t) for t in transforms])
    pipeline = tio.Compose(all_tfms) if all_tfms else None

    # Create output directories
    output_dir = Path(output_dir)
    img_dir = output_dir / 'images'
    img_dir.mkdir(parents=True, exist_ok=True)
    if mask_col is not None:
        mask_dir = output_dir / 'masks'
        mask_dir.mkdir(parents=True, exist_ok=True)

    # Build work items, filtering skip_existing
    work_items = []
    skipped = 0
    for idx in range(len(df)):
        img_path = df[img_col].iloc[idx]
        out_img = img_dir / Path(img_path).name

        mask_path = df[mask_col].iloc[idx] if mask_col is not None else None
        out_mask = (mask_dir / Path(mask_path).name) if mask_col is not None else None

        if skip_existing:
            img_ok = out_img.exists() and out_img.stat().st_size > 0
            mask_ok = out_mask is None or (out_mask.exists() and out_mask.stat().st_size > 0)
            if img_ok and mask_ok:
                skipped += 1
                continue

        work_items.append({
            'idx': idx, 'img_path': img_path, 'mask_path': mask_path,
            'out_img': out_img, 'out_mask': out_mask,
        })

    # Process cases
    processed = 0
    failed = 0
    failed_cases = []

    def _process_case(item):
        subject_dict = {'image': tio.ScalarImage(item['img_path'])}
        if item['mask_path'] is not None:
            subject_dict['mask'] = tio.LabelMap(item['mask_path'])

        subject = tio.Subject(**subject_dict)
        if pipeline is not None:
            subject = pipeline(subject)

        # Atomic write: save to temp file (with valid NIfTI extension), then rename
        out_img = item['out_img']
        tmp_img = out_img.parent / f'.tmp_{out_img.name}'
        subject['image'].save(str(tmp_img))
        os.rename(str(tmp_img), str(out_img))

        if item['out_mask'] is not None:
            out_mask = item['out_mask']
            tmp_mask = out_mask.parent / f'.tmp_{out_mask.name}'
            subject['mask'].save(str(tmp_mask))
            os.rename(str(tmp_mask), str(out_mask))

    if work_items:
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = {executor.submit(_process_case, item): item for item in work_items}
            for future in tqdm(as_completed(futures), total=len(futures),
                               desc='Preprocessing'):
                item = futures[future]
                try:
                    future.result()
                    processed += 1
                except Exception as e:
                    failed += 1
                    failed_cases.append(Path(item['img_path']).name)
                    warnings.warn(f"Failed to process {item['img_path']}: {e}")

    # Create new columns for preprocessed paths (preserve originals)
    df[f'{img_col}_preprocessed'] = [str(img_dir / Path(p).name) for p in df[img_col]]

    if mask_col is not None:
        df[f'{mask_col}_preprocessed'] = [str(mask_dir / Path(p).name) for p in df[mask_col]]

    print(f"Preprocessing complete: {processed} processed, {skipped} skipped, {failed} failed")
    if failed_cases:
        print(f"Failed cases: {failed_cases}")

# %% ../nbs/08_dataset_info.ipynb #9b81f6e8-abd7-4bf6-be4c-4118986c308a
def get_class_weights(labels: (np.array, list), class_weight: str = 'balanced') -> torch.Tensor: 
    """Calculates and returns the class weights.

    Args:
        labels: An array or list of class labels for each instance in the dataset.
        class_weight: Defaults to 'balanced'.

    Returns:
        A tensor of class weights.
    """
    
    class_weights =  compute_class_weight(class_weight=class_weight, classes=np.unique(labels), y=labels)
    
    return torch.Tensor(class_weights)
