"""
Preprocessing module for MedTDA.

This module provides the Preprocessor class for standalone image preprocessing,
including resampling, windowing, normalization, and background masking.
"""

import numpy as np
from typing import Union, Tuple, Optional, Dict
from pathlib import Path

from .loaders import load_image, load_mask, save_image
from .utils import normalize_image, apply_mask, crop_to_roi, resample_image

# Sentinel value to distinguish "not provided" from "None"
_UNSET = object()


class Preprocessor:
    """
    Standalone image preprocessor for inspection and validation.
    
    Handles traditional image preprocessing: resampling, windowing,
    normalization, and background masking. Does NOT handle filtration-specific
    transformations (e.g., negation for superlevel).
    
    Parameters
    ----------
    spacing : tuple or None, default=None
        Target voxel spacing for resampling (3D/4D only).
        If None, no resampling is performed.
    window : tuple or None, default=None
        Windowing parameters as (center, width).
        If None, no windowing is applied.
    normalize : bool, default=False
        Whether to apply normalization.
    normalize_method : {'minmax', 'zscore', 'robust'}, default='minmax'
        Normalization method (only used if normalize=True):
        - 'minmax': Scales to [0, 1] range
        - 'zscore': Z-score normalization (mean=0, std=1)
        - 'robust': Robust scaling using quantiles
    clip_percentiles : tuple of float, optional
        If provided, clip intensities to (lower, upper) percentiles before
        normalization. For example, (1, 99) removes extreme outliers.
        Only used if normalize=True.
    background_value : float or None, default=None
        Value to assign to background pixels when a mask is provided.
        If None, automatically determined:
        - If normalize=True: -1 (below [0,1] range)
        - If normalize=False and 2D: 0
        - If normalize=False and 3D/4D: -3000
    label : int or None, default=1
        For multi-label masks, which label value to consider as foreground.
        If None, treat mask as binary (any non-zero value is foreground).
    crop_to_roi : bool, default=True
        If True and a mask is provided, crop image and mask to the bounding
        box of the ROI with padding. This reduces memory usage and speeds up
        persistent homology computation.
    roi_padding : int, default=1
        Number of background pixels to include on all sides of the ROI when
        crop_to_roi is enabled. This ensures the ROI is not cut too tightly.
        
    Examples
    --------
    Basic preprocessing with normalization:
    
    >>> from medtda import Preprocessor
    >>> preprocessor = Preprocessor(
    ...     normalize=True,
    ...     normalize_method='minmax',
    ...     background_value=None
    ... )
    >>> preprocessed_img, metadata = preprocessor.preprocess(
    ...     image='path/to/image.nii.gz',
    ...     mask='path/to/mask.nii.gz'
    ... )
    >>> print(metadata['transforms'])
    ['normalize_minmax', 'background_set']
    
    Preprocessing with resampling to isotropic spacing:
    
    >>> preprocessor = Preprocessor(
    ...     spacing=(1.0, 1.0, 1.0),  # Resample to 1mm isotropic
    ...     normalize=True,
    ...     normalize_method='minmax',
    ...     crop_to_roi=True
    ... )
    >>> preprocessed_img, metadata = preprocessor.preprocess(
    ...     image='path/to/ct_scan.nii.gz',
    ...     mask='path/to/roi_mask.nii.gz'
    ... )
    >>> print(metadata['transforms'])
    ['resample_1.0x1.0x1.0', 'crop_to_roi_pad1', 'normalize_minmax', 'background_set']
    """
    
    def __init__(
        self,
        spacing: Optional[Tuple[float, ...]] = None,
        window: Optional[Tuple[float, float]] = None,
        normalize: bool = False,
        normalize_method: str = 'minmax',
        clip_percentiles: Optional[Tuple[float, float]] = None,
        background_value: Optional[float] = None,
        label: Optional[int] = 1,
        crop_to_roi: bool = True,
        roi_padding: int = 1
    ):
        """Initialize preprocessor with parameters."""
        self.spacing = spacing
        self.window = window
        self.normalize = normalize
        self.normalize_method = normalize_method
        self.clip_percentiles = clip_percentiles
        self.background_value = background_value
        self.label = label
        self.crop_to_roi = crop_to_roi
        self.roi_padding = roi_padding
        
        # Validate normalize_method
        if normalize_method not in ['minmax', 'zscore', 'robust']:
            raise ValueError(
                f"normalize_method must be 'minmax', 'zscore', or 'robust'. "
                f"Got '{normalize_method}'."
            )
    
    def set_spacing(self, spacing: Optional[Tuple[float, ...]]) -> None:
        """
        Update spacing parameter.
        
        Parameters
        ----------
        spacing : tuple or None
            Target voxel spacing for resampling.
        """
        self.spacing = spacing
    
    def set_window(self, window_center: float, window_width: float) -> None:
        """
        Set windowing parameters.
        
        Parameters
        ----------
        window_center : float
            Window center value.
        window_width : float
            Window width value.
        """
        self.window = (window_center, window_width)
    
    def set_normalize(self, enabled: bool, method: str = 'minmax') -> None:
        """
        Enable/disable normalization and set method.
        
        Parameters
        ----------
        enabled : bool
            Whether to enable normalization.
        method : {'minmax', 'zscore', 'robust'}, default='minmax'
            Normalization method.
        """
        self.normalize = enabled
        if method not in ['minmax', 'zscore', 'robust']:
            raise ValueError(
                f"normalize_method must be 'minmax', 'zscore', or 'robust'. "
                f"Got '{method}'."
            )
        self.normalize_method = method
    
    def set_clip_percentiles(self, percentiles: Optional[Tuple[float, float]]) -> None:
        """
        Set percentile clipping for outlier removal.
        
        Parameters
        ----------
        percentiles : tuple of float or None
            If provided, should be (lower_percentile, upper_percentile) where
            0 <= lower_percentile < upper_percentile <= 100. Values outside
            these percentiles will be clipped. Set to None to disable clipping.
        """
        if percentiles is not None:
            lower_p, upper_p = percentiles
            if lower_p < 0 or upper_p > 100 or lower_p >= upper_p:
                raise ValueError(
                    f"clip_percentiles must satisfy 0 <= lower < upper <= 100. "
                    f"Got ({lower_p}, {upper_p})."
                )
        self.clip_percentiles = percentiles
    
    def set_background_value(self, value: Optional[float]) -> None:
        """
        Set background value for masked regions.
        
        Parameters
        ----------
        value : float or None
            Background value, or None for automatic determination.
        """
        self.background_value = value
    
    def set_label(self, label: Optional[int]) -> None:
        """
        Set label value to extract from multi-label mask.
        
        Parameters
        ----------
        label : int or None
            Label value to extract. If None, treat mask as binary.
        """
        self.label = label
    
    def preprocess(
        self,
        image: Union[str, Path, np.ndarray],
        mask: Optional[Union[str, Path, np.ndarray]] = None,
        label: Optional[int] = _UNSET  # type: ignore[assignment]
    ) -> Tuple[np.ndarray, Dict]:
        """
        Main preprocessing method.
        
        Parameters
        ----------
        image : str, Path, or numpy.ndarray
            Input image (file path or array).
        mask : str, Path, or numpy.ndarray, optional
            Input mask (file path or array).
        label : int or None, optional
            Label value to extract from mask. If None, uses instance default.
            
        Returns
        -------
        preprocessed_image : numpy.ndarray
            Preprocessed image ready for PH computation.
        metadata : dict
            Dictionary containing:
            
            - 'original_range': (min, max) of original image
            - 'final_range': (min, max) of preprocessed image
            - 'transforms': list of applied transformations
            - 'background_value': actual background value used
            - 'original_shape': shape of original image
            - 'final_shape': shape after preprocessing
            - 'label': label value used for mask extraction
            - 'crop_info': dict with cropping details (None if not cropped)
              
              - 'bbox': tuple of slices defining the bounding box
              - 'original_shape': original image shape
              - 'cropped_shape': shape after cropping
              - 'roi_size': number of voxels in ROI
              - 'reduction_factor': ratio of cropped size to original size
        """
        # Use instance label if not provided, otherwise use the provided value (even if None)
        if label is _UNSET:
            label = self.label
        
        # Load image if path is provided
        if isinstance(image, (str, Path)):
            image_array, img_metadata = load_image(image)
        else:
            image_array = np.array(image, dtype=np.float64)
            img_metadata = {}
        
        # Load mask if provided
        mask_array = None
        if mask is not None:
            if isinstance(mask, (str, Path)):
                mask_array, _ = load_mask(mask)
            else:
                mask_array = np.array(mask, dtype=np.uint8)
        
        # Store original statistics
        original_min = float(np.min(image_array))
        original_max = float(np.max(image_array))
        original_shape = image_array.shape
        
        transforms = []
        crop_info = None
        
        # Validate that the specified label exists in the mask (if provided)
        if mask_array is not None and label is not None:
            unique_labels = np.unique(mask_array)
            if label not in unique_labels:
                raise ValueError(
                    f"Specified label {label} does not exist in the mask. "
                    f"Available labels: {unique_labels.tolist()}"
                )

        # 1. Resampling (if spacing is specified and image is 3D/4D)
        if self.spacing is not None and image_array.ndim >= 3:
            # Check if we have original spacing information
            if 'spacing' not in img_metadata or img_metadata['spacing'] is None:
                import warnings
                warnings.warn(
                    "Cannot resample: original spacing information not available. "
                    "Resampling requires images loaded from file formats that store "
                    "spacing metadata (NIfTI, NRRD, MHA, etc.). Skipping resampling."
                )
            else:
                original_spacing = img_metadata['spacing']
                
                # Ensure target spacing has same dimensions as image
                target_spacing = self.spacing
                if len(target_spacing) != len(original_spacing):
                    import warnings
                    warnings.warn(
                        f"Target spacing dimension ({len(target_spacing)}) does not match "
                        f"image dimension ({len(original_spacing)}). Skipping resampling."
                    )
                else:
                    # Resample image
                    image_array = resample_image(
                        image_array,
                        original_spacing,
                        target_spacing,
                        interpolator='linear'
                    )
                    
                    # Resample mask if provided
                    if mask_array is not None:
                        mask_array = resample_image(
                            mask_array,
                            original_spacing,
                            target_spacing,
                            interpolator='nearest'  # Use nearest neighbor for masks
                        )
                        
                        # Verify the specified label still exists after resampling
                        if label is not None:
                            unique_labels_after = np.unique(mask_array)
                            if label not in unique_labels_after:
                                raise ValueError(
                                    f"Label {label} was lost during resampling. "
                                    f"This can happen when the ROI is very small relative to the "
                                    f"resampling ratio (original spacing: {original_spacing}, "
                                    f"target spacing: {target_spacing}). "
                                    f"Labels remaining after resampling: {unique_labels_after.tolist()}. "
                                    f"Consider: (1) using a different target spacing, "
                                    f"(2) checking if the ROI is too small, or "
                                    f"(3) using a different label value."
                                )
                    
                    transforms.append(f'resample_{"x".join(map(str, target_spacing))}')
        
        # 2. Crop to ROI (if mask is provided and crop_to_roi is enabled)
        if mask_array is not None and self.crop_to_roi:
            try:
                from .utils import crop_to_roi as crop_func
                image_array, mask_array, crop_info = crop_func(
                    image_array, mask_array, padding=self.roi_padding, label=label
                )
                transforms.append(f'crop_to_roi_pad{self.roi_padding}')
            except ValueError as e:
                # If cropping fails (e.g., empty mask for the specified label)
                error_msg = (
                    f"ROI cropping failed: {e}. "
                    f"This usually means the mask is empty for label={label}. "
                )
                # Check if it's a label issue
                if label is not None:
                    unique_labels = np.unique(mask_array)
                    if label not in unique_labels:
                        error_msg += (
                            f"Label {label} does not exist in mask. "
                            f"Available labels: {unique_labels.tolist()}"
                        )
                raise ValueError(error_msg) from e

        # 3. Windowing (if specified)
        if self.window is not None:
            window_center, window_width = self.window
            window_min = window_center - window_width / 2
            window_max = window_center + window_width / 2
            image_array = np.clip(image_array, window_min, window_max)
            transforms.append(f'window_{window_center}_{window_width}')
        
        # 4. Normalization (if enabled)
        if self.normalize:
            image_array = normalize_image(
                image_array,
                method=self.normalize_method,
                mask=mask_array,
                label=label,
                clip_percentiles=self.clip_percentiles
            )
            transforms.append(f'normalize_{self.normalize_method}')
        
        # 5. Background masking (if mask is provided)
        actual_background_value = None
        if mask_array is not None:
            # Determine background value
            if self.background_value is not None:
                actual_background_value = self.background_value
            else:
                # Auto-determine based on normalization and dimensions
                if self.normalize:
                    actual_background_value = -1.0
                elif image_array.ndim == 2:
                    actual_background_value = 0.0
                else:  # 3D or 4D
                    actual_background_value = -3000.0
            
            # Apply mask
            image_array = apply_mask(
                image_array,
                mask_array,
                background_value=actual_background_value,
                label=label
            )
            transforms.append('background_set')
        
        # Store final statistics
        final_min = float(np.min(image_array))
        final_max = float(np.max(image_array))
        final_shape = image_array.shape
        
        metadata = {
            'original_range': (original_min, original_max),
            'final_range': (final_min, final_max),
            'transforms': transforms,
            'background_value': actual_background_value,
            'original_shape': original_shape,
            'final_shape': final_shape,
            'label': label,
            'crop_info': crop_info,  # None if cropping not performed
        }
        
        return image_array, metadata
    
    def preprocess_and_save(
        self,
        image: Union[str, Path, np.ndarray],
        mask: Optional[Union[str, Path, np.ndarray]] = None,
        label: Optional[int] = _UNSET,  # type: ignore[assignment]
        output_path: Optional[Union[str, Path]] = None
    ) -> Tuple[np.ndarray, Dict]:
        """
        Preprocess and save result to disk.
        
        Useful for inspection and debugging.
        
        Parameters
        ----------
        image : str, Path, or numpy.ndarray
            Input image.
        mask : str, Path, or numpy.ndarray, optional
            Input mask.
        label : int or None, optional
            Label value to extract from mask. If None, uses instance default.
        output_path : str or Path, optional
            Path to save preprocessed image.
            If None, uses 'preprocessed_{original_name}'.
            
        Returns
        -------
        preprocessed_image : numpy.ndarray
            Preprocessed image.
        metadata : dict
            Preprocessing metadata.
        """
        preprocessed_image, metadata = self.preprocess(image, mask, label)
        
        # Determine output path
        if output_path is None:
            if isinstance(image, (str, Path)):
                image_path = Path(image)
                output_path = image_path.parent / f"preprocessed_{image_path.name}"
            else:
                output_path = "preprocessed_image.nii.gz"
        
        # Save preprocessed image
        save_image(preprocessed_image, output_path)
        
        return preprocessed_image, metadata
    
    def get_preprocessing_info(self) -> Dict:
        """
        Return dictionary of current preprocessing settings.
        
        Returns
        -------
        dict
            Dictionary with current parameter values.
        """
        return {
            'spacing': self.spacing,
            'window': self.window,
            'normalize': self.normalize,
            'normalize_method': self.normalize_method,
            'background_value': self.background_value,
            'label': self.label,
            'crop_to_roi': self.crop_to_roi,
            'roi_padding': self.roi_padding,
        }
    

