"""
Utility functions for MedTDA.

This module provides helper functions for image processing, normalization,
and other common operations.
"""

import numpy as np
from typing import Union, Tuple, Optional, List

try:
    import SimpleITK as sitk
    SITK_AVAILABLE = True
except ImportError:
    SITK_AVAILABLE = False


def resample_image(
    image_array: np.ndarray,
    original_spacing: Tuple[float, ...],
    target_spacing: Tuple[float, ...],
    interpolator: str = 'linear'
) -> np.ndarray:
    """
    Resample image to a new voxel spacing using SimpleITK.
    
    This function resamples a 3D or 4D image to match the target spacing,
    which is useful for standardizing images from different scanners or
    acquisition protocols.
    
    Parameters
    ----------
    image_array : numpy.ndarray
        Input image array (3D or 4D).
    original_spacing : tuple of float
        Original voxel spacing (x, y, z) or (x, y, z, t).
        For SimpleITK, this is typically (width, height, depth).
    target_spacing : tuple of float
        Target voxel spacing (x, y, z) or (x, y, z, t).
    interpolator : {'linear', 'nearest', 'bspline'}, default='linear'
        Interpolation method:
        - 'linear': Linear interpolation (good for images)
        - 'nearest': Nearest neighbor (good for masks/labels)
        - 'bspline': B-spline interpolation (smooth but slower)
        
    Returns
    -------
    numpy.ndarray
        Resampled image array.
        
    Raises
    ------
    ImportError
        If SimpleITK is not installed.
    ValueError
        If image dimensions or spacing don't match.
        
    Examples
    --------
    >>> import numpy as np
    >>> image = np.random.rand(100, 100, 50)
    >>> # Resample from 1mm × 1mm × 2mm to isotropic 1mm³
    >>> resampled = resample_image(
    ...     image, 
    ...     original_spacing=(1.0, 1.0, 2.0),
    ...     target_spacing=(1.0, 1.0, 1.0)
    ... )
    >>> print(resampled.shape)
    (100, 100, 100)
    """
    if not SITK_AVAILABLE:
        raise ImportError(
            "SimpleITK is required for image resampling. "
            "Install with: pip install SimpleITK"
        )
    
    # Validate inputs
    if image_array.ndim not in [3, 4]:
        raise ValueError(
            f"Image must be 3D or 4D for resampling. Got {image_array.ndim}D array."
        )
    
    if len(original_spacing) != image_array.ndim:
        raise ValueError(
            f"Original spacing dimension ({len(original_spacing)}) must match "
            f"image dimension ({image_array.ndim})."
        )
    
    if len(target_spacing) != image_array.ndim:
        raise ValueError(
            f"Target spacing dimension ({len(target_spacing)}) must match "
            f"image dimension ({image_array.ndim})."
        )
    
    # Map interpolator string to SimpleITK constant
    interpolator_map = {
        'linear': sitk.sitkLinear,
        'nearest': sitk.sitkNearestNeighbor,
        'bspline': sitk.sitkBSpline
    }
    
    if interpolator not in interpolator_map:
        raise ValueError(
            f"Interpolator must be one of {list(interpolator_map.keys())}. "
            f"Got '{interpolator}'."
        )
    
    sitk_interpolator = interpolator_map[interpolator]
    
    # Convert numpy array to SimpleITK image
    # Note: SimpleITK uses (x, y, z) ordering, numpy uses (z, y, x)
    sitk_image = sitk.GetImageFromArray(image_array)
    sitk_image.SetSpacing(original_spacing)
    
    # Calculate new size based on spacing ratio
    original_size = sitk_image.GetSize()
    new_size = [
        int(np.round(original_size[i] * (original_spacing[i] / target_spacing[i])))
        for i in range(len(original_size))
    ]
    
    # Setup resampler
    resampler = sitk.ResampleImageFilter()
    resampler.SetOutputSpacing(target_spacing)
    resampler.SetSize(new_size)
    resampler.SetOutputDirection(sitk_image.GetDirection())
    resampler.SetOutputOrigin(sitk_image.GetOrigin())
    resampler.SetTransform(sitk.Transform())
    resampler.SetDefaultPixelValue(float(np.min(image_array)))
    resampler.SetInterpolator(sitk_interpolator)
    
    # Execute resampling
    resampled_sitk = resampler.Execute(sitk_image)
    
    # Convert back to numpy array
    resampled_array = sitk.GetArrayFromImage(resampled_sitk)
    
    return resampled_array

def extract_label_from_mask(
    mask: np.ndarray,
    label: Optional[int] = None
) -> np.ndarray:
    """
    Convert multi-label mask to binary mask for specified label.
    
    Parameters
    ----------
    mask : numpy.ndarray
        Input mask array (can be binary or multi-label).
    label : int or None, default=None
        Label value to extract. If None, treat mask as binary
        (any non-zero value is considered foreground).
        
    Returns
    -------
    numpy.ndarray
        Binary mask where specified label = 1, others = 0.
        
    Examples
    --------
    >>> multi_label_mask = np.array([[0, 1, 2], [1, 2, 3]])
    >>> binary_mask = extract_label_from_mask(multi_label_mask, label=2)
    >>> print(binary_mask)
    [[0 0 1]
     [0 1 0]]
    """
    if label is None:
        # Treat as binary mask (any non-zero is foreground)
        return (mask > 0).astype(np.uint8)
    else:
        # Extract specific label
        return (mask == label).astype(np.uint8)


def normalize_image(
    image_array: np.ndarray,
    method: str = 'minmax',
    mask: Optional[np.ndarray] = None,
    label: Optional[int] = 1,
    clip_percentiles: Optional[Tuple[float, float]] = None
) -> np.ndarray:
    """
    Normalize image array.
    
    Parameters
    ----------
    image_array : numpy.ndarray
        Input image array to normalize.
    method : {'minmax', 'zscore', 'robust'}, default='minmax'
        Normalization method:
        - 'minmax': Min-max scaling to [0, 1]
        - 'zscore': Z-score normalization (mean=0, std=1)
        - 'robust': Robust scaling using median and IQR
    mask : numpy.ndarray, optional
        Mask array. If provided, statistics are computed only
        within the masked region.
    label : int or None, default=1
        For multi-label masks, which label value to consider as foreground.
        If None, any non-zero value is considered foreground.
    clip_percentiles : tuple of float, optional
        If provided, clip values to (lower_percentile, upper_percentile) 
        before normalization. For example, (1, 99) clips to 1st and 99th 
        percentiles, removing extreme outliers.
        
    Returns
    -------
    numpy.ndarray
        Normalized image array.
        
    Raises
    ------
    ValueError
        If method is not recognized or arrays have incompatible shapes.
    """
    if method not in ['minmax', 'zscore', 'robust']:
        raise ValueError(
            f"Normalization method must be 'minmax', 'zscore', or 'robust'. "
            f"Got '{method}'."
        )

    if mask is not None and mask.shape != image_array.shape:
        raise ValueError(
            f"Mask shape {mask.shape} does not match image shape {image_array.shape}."
        )

    # Select region for computing statistics
    if mask is not None:
        # Extract binary mask for specified label
        binary_mask = extract_label_from_mask(mask, label)
        values = image_array[binary_mask > 0]
    else:
        values = image_array.ravel()

    if len(values) == 0:
        return image_array.copy()
    
    # Apply percentile clipping if requested
    clipped_image = image_array.copy()
    if clip_percentiles is not None:
        lower_p, upper_p = clip_percentiles
        if lower_p < 0 or upper_p > 100 or lower_p >= upper_p:
            raise ValueError(
                f"clip_percentiles must be (lower, upper) with 0 <= lower <upper <= 100. "
                f"Got {clip_percentiles}."
            )
        lower_val = np.percentile(values, lower_p)
        upper_val = np.percentile(values, upper_p)
        clipped_image = np.clip(clipped_image, lower_val, upper_val)
        # Recompute values after clipping for normalization
        if mask is not None:
            values = clipped_image[binary_mask > 0]
        else:
            values = clipped_image.ravel()

    if method == 'minmax':
        min_val: float = float(np.min(values))
        max_val: float = float(np.max(values))
        if max_val - min_val == 0:
            return np.zeros_like(image_array)
        normalized = (clipped_image - min_val) / (max_val - min_val)

    elif method == 'zscore':
        mean_val = np.mean(values)
        std_val = np.std(values)
        if std_val == 0:
            return np.zeros_like(image_array)
        normalized = (clipped_image - mean_val) / std_val

    elif method == 'robust':
        median_val = np.median(values)
        q1 = np.percentile(values, 25)
        q3 = np.percentile(values, 75)
        iqr = q3 - q1
        if iqr == 0:
            return np.zeros_like(image_array)
        normalized = (clipped_image - median_val) / iqr

    return normalized


def apply_mask(
    image_array: np.ndarray,
    mask: np.ndarray,
    background_value: float = 0,
    label: Optional[int] = 1
) -> np.ndarray:
    """
    Apply mask to image array.
    
    Sets values outside the mask to background_value.
    
    Parameters
    ----------
    image_array : numpy.ndarray
        Input image array.
    mask : numpy.ndarray
        Mask array (same shape as image).
    background_value : float, default=0
        Value to assign to pixels outside the mask.
    label : int or None, default=1
        For multi-label masks, which label value to consider as foreground.
        If None, any non-zero value is considered foreground.
        
    Returns
    -------
    numpy.ndarray
        Masked image array.
        
    Raises
    ------
    ValueError
        If mask shape doesn't match image shape.
    """
    if mask.shape != image_array.shape:
        raise ValueError(
            f"Mask shape {mask.shape} does not match image shape {image_array.shape}."
        )

    # Extract binary mask for specified label
    binary_mask = extract_label_from_mask(mask, label)
    
    masked_image = image_array.copy()
    masked_image[binary_mask == 0] = background_value

    return masked_image


def check_array_dimensions(
    array: np.ndarray,
    expected_dims: Union[int, Tuple[int, ...]],
    name: str = "array"
) -> None:
    """
    Check if array has expected dimensions.
    
    Parameters
    ----------
    array : numpy.ndarray
        Array to check.
    expected_dims : int or tuple of int
        Expected number of dimensions. If int, array must have exactly
        that many dimensions. If tuple, array must have one of those
        dimensions.
    name : str, default='array'
        Name of the array for error messages.
        
    Raises
    ------
    ValueError
        If array dimensions don't match expectations.
    """
    if isinstance(expected_dims, int):
        expected_dims = (expected_dims,)

    if array.ndim not in expected_dims:
        raise ValueError(
            f"{name} must have {expected_dims} dimensions. "
            f"Got {array.ndim}D array."
        )


def safe_divide(numerator: np.ndarray, denominator: np.ndarray) -> np.ndarray:
    """
    Safely divide two arrays, handling division by zero.
    
    Parameters
    ----------
    numerator : numpy.ndarray
        Numerator array.
    denominator : numpy.ndarray
        Denominator array.
        
    Returns
    -------
    numpy.ndarray
        Result of division with zeros where denominator is zero.
    """
    with np.errstate(divide='ignore', invalid='ignore'):
        result = numerator / denominator
        result[~np.isfinite(result)] = 0
    return result


def get_image_statistics(
    image_array: np.ndarray,
    mask: Optional[np.ndarray] = None,
    label: Optional[int] = 1
) -> dict:
    """
    Compute basic statistics for image array.
    
    Parameters
    ----------
    image_array : numpy.ndarray
        Input image array.
    mask : numpy.ndarray, optional
        Mask array. If provided, statistics are computed only
        within the masked region.
    label : int or None, default=1
        For multi-label masks, which label value to consider as foreground.
        If None, any non-zero value is considered foreground.
        
    Returns
    -------
    dict
        Dictionary containing:
        - 'min': Minimum value
        - 'max': Maximum value
        - 'mean': Mean value
        - 'std': Standard deviation
        - 'median': Median value
        - 'q1': First quartile
        - 'q3': Third quartile
    """
    if mask is not None:
        # Extract binary mask for specified label
        binary_mask = extract_label_from_mask(mask, label)
        values = image_array[binary_mask > 0]
    else:
        values = image_array.ravel()

    if len(values) == 0:
        return {
            'min': 0, 'max': 0, 'mean': 0, 'std': 0,
            'median': 0, 'q1': 0, 'q3': 0
        }

    return {
        'min': float(np.min(values)),
        'max': float(np.max(values)),
        'mean': float(np.mean(values)),
        'std': float(np.std(values)),
        'median': float(np.median(values)),
        'q1': float(np.percentile(values, 25)),
        'q3': float(np.percentile(values, 75)),
    }


def crop_to_roi(
    image: np.ndarray,
    mask: np.ndarray,
    padding: int = 1,
    label: Optional[int] = 1
) -> Tuple[np.ndarray, np.ndarray, dict]:
    """
    Crop image and mask to the bounding box of the Region of Interest (ROI).
    
    This function finds the smallest bounding box containing all non-zero 
    mask values and crops both image and mask to this region with optional
    padding. This reduces memory usage and speeds up persistent homology
    computation.
    
    Parameters
    ----------
    image : numpy.ndarray
        Input image array.
    mask : numpy.ndarray
        Mask array (same shape as image).
    padding : int, default=1
        Number of background pixels to include on all sides of the ROI.
        This ensures the ROI is not cut too tightly.
    label : int or None, default=1
        For multi-label masks, which label value to consider as foreground.
        If None, any non-zero value is considered foreground.
        
    Returns
    -------
    cropped_image : numpy.ndarray
        Cropped image array.
    cropped_mask : numpy.ndarray
        Cropped mask array.
    crop_info : dict
        Dictionary containing cropping information:
        - '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 before padding
        - 'reduction_factor': Ratio of cropped size to original size
        
    Raises
    ------
    ValueError
        If mask shape doesn't match image shape or if mask is all zeros.
        
    Examples
    --------
    >>> import numpy as np
    >>> image = np.random.rand(100, 100, 100)
    >>> mask = np.zeros((100, 100, 100))
    >>> mask[40:60, 40:60, 40:60] = 1  # Small ROI
    >>> cropped_img, cropped_mask, info = crop_to_roi(image, mask, padding=1)
    >>> print(f"Original: {image.shape}, Cropped: {cropped_img.shape}")
    Original: (100, 100, 100), Cropped: (22, 22, 22)
    >>> print(f"Reduction: {info['reduction_factor']:.2f}x smaller")
    Reduction: 10.65x smaller
    """
    if mask.shape != image.shape:
        raise ValueError(
            f"Mask shape {mask.shape} does not match image shape {image.shape}."
        )
    
    # Extract binary mask for specified label
    binary_mask = extract_label_from_mask(mask, label)
    
    # Find coordinates of all non-zero mask voxels
    coords = np.argwhere(binary_mask > 0)
    
    if len(coords) == 0:
        raise ValueError(
            "Mask is all zeros. Cannot crop to ROI. "
            "Please provide a mask with at least one non-zero voxel."
        )
    
    # Compute bounding box
    mins = coords.min(axis=0)
    maxs = coords.max(axis=0)
    
    # Add padding
    ndim = image.ndim
    slices: List[slice] = []
    for i in range(ndim):
        start = max(0, mins[i] - padding)
        end = min(image.shape[i], maxs[i] + 1 + padding)
        slices.append(slice(start, end))
    bbox = tuple(slices)
    
    # Crop both image and mask
    cropped_image = image[bbox]
    cropped_mask = mask[bbox]
    
    # Compute statistics
    original_size = np.prod(image.shape)
    cropped_size = np.prod(cropped_image.shape)
    roi_size = len(coords)
    reduction_factor = original_size / cropped_size
    
    crop_info = {
        'bbox': bbox,
        'original_shape': image.shape,
        'cropped_shape': cropped_image.shape,
        'roi_size': int(roi_size),
        'reduction_factor': float(reduction_factor)
    }
    
    return cropped_image, cropped_mask, crop_info
