"""
Barcode extraction module for MedTDA.

This module provides the BarcodeExtractor class for extracting raw persistence
barcodes from medical images without vectorization.
"""

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

from .preprocessor import Preprocessor, _UNSET
from .loaders import load_image, load_mask
from .ph_computer import compute_ph


class BarcodeExtractor:
    """
    Extract raw persistence barcodes from medical images.
    
    This class handles the full pipeline: image loading, preprocessing,
    and persistent homology computation, returning raw barcodes without
    vectorization.
    
    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).
    clip_percentiles : tuple of float, optional
        If provided, clip intensities to percentile range (lower, upper) before
        normalization. For example, (1, 99) clips to 1st and 99th percentiles.
        Only applied 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 based on normalization and dimensionality.
    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.
    filtration_type : {'sublevel', 'superlevel'}, default='sublevel'
        Type of filtration for persistent homology.
    construction : {'T', 'V'}, default='T'
        Type of cubical complex construction.
    max_dimension : int, default=-1
        Maximum homology dimension to compute (computes H0, H1, ..., H_max_dimension).
        If -1, automatically determined based on image dimensionality (uses ndim - 1).
        
    Examples
    --------
    >>> from medtda import BarcodeExtractor
    >>> extractor = BarcodeExtractor(
    ...     normalize=True,
    ...     normalize_method='minmax',
    ...     filtration_type='sublevel',
    ...     max_dimension=2
    ... )
    >>> barcodes = extractor.execute(
    ...     image='path/to/image.nii.gz',
    ...     mask='path/to/mask.nii.gz'
    ... )
    >>> print(barcodes.keys())  # e.g., dict_keys(['H0', 'H1', 'H2'])
    >>> print(barcodes['H1'].shape)  # (n_features, 2) for (birth, death) pairs
    """
    
    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,
        filtration_type: str = 'sublevel',
        construction: str = 'T',
        max_dimension: int = -1
    ):
        """Initialize BarcodeExtractor with parameters."""
        # Initialize preprocessor with preprocessing parameters
        self.preprocessor = Preprocessor(
            spacing=spacing,
            window=window,
            normalize=normalize,
            normalize_method=normalize_method,
            clip_percentiles=clip_percentiles,
            background_value=background_value,
            label=label,
            crop_to_roi=crop_to_roi,
            roi_padding=roi_padding
        )
        
        # Store PH computation parameters
        self.filtration_type = filtration_type
        self.construction = construction
        self.max_dimension = max_dimension
        
        # Validate parameters
        self._validate_parameters()
    
    def _validate_parameters(self) -> None:
        """Validate initialization parameters."""
        if self.filtration_type not in ['sublevel', 'superlevel']:
            raise ValueError(
                f"filtration_type must be 'sublevel' or 'superlevel'. "
                f"Got '{self.filtration_type}'."
            )
        
        if self.construction not in ['T', 'V']:
            raise ValueError(
                f"construction must be 'T' or 'V'. "
                f"Got '{self.construction}'."
            )
        
        if not isinstance(self.max_dimension, int):
            raise ValueError(
                f"max_dimension must be an integer. Got {type(self.max_dimension).__name__}."
            )
    
    def set_spacing(self, spacing: Optional[Tuple[float, ...]]) -> None:
        """
        Update spacing parameter.
        
        Parameters
        ----------
        spacing : tuple or None
            Target voxel spacing for resampling.
        """
        self.preprocessor.set_spacing(spacing)
    
    def set_window(self, window_center: float, window_width: float) -> None:
        """
        Set windowing parameters.
        
        Parameters
        ----------
        window_center : float
            Center of the window.
        window_width : float
            Width of the window.
        """
        self.preprocessor.set_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.preprocessor.set_normalize(enabled, 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.
        """
        self.preprocessor.set_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. If None, automatically determined.
        """
        self.preprocessor.set_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.preprocessor.set_label(label)
    
    def set_filtration_type(self, filtration_type: str) -> None:
        """
        Set persistent homology filtration type.
        
        Parameters
        ----------
        filtration_type : {'sublevel', 'superlevel'}
            Filtration type.
        """
        if filtration_type not in ['sublevel', 'superlevel']:
            raise ValueError(
                f"filtration_type must be 'sublevel' or 'superlevel'. "
                f"Got '{filtration_type}'."
            )
        self.filtration_type = filtration_type
    
    def set_construction(self, construction: str) -> None:
        """
        Set cubical complex construction type.
        
        Parameters
        ----------
        construction : {'T', 'V'}
            Construction type.
        """
        if construction not in ['T', 'V']:
            raise ValueError(
                f"construction must be 'T' or 'V'. "
                f"Got '{construction}'."
            )
        self.construction = construction
    
    def set_max_dimension(self, max_dimension: int) -> None:
        """
        Set maximum homology dimension to compute.
        
        Parameters
        ----------
        max_dimension : int
            Maximum dimension to compute. Use -1 for automatic determination.
        """
        if not isinstance(max_dimension, int):
            raise ValueError(
                f"max_dimension must be an integer. Got {type(max_dimension).__name__}."
            )
        self.max_dimension = max_dimension
    
    def get_settings(self) -> Dict:
        """
        Get current settings for preprocessing and PH computation.
        
        Returns
        -------
        dict
            Dictionary containing all current settings.
        """
        preprocessing_info = self.preprocessor.get_preprocessing_info()
        return {
            'preprocessing': preprocessing_info,
            'filtration_type': self.filtration_type,
            'construction': self.construction,
            'max_dimension': self.max_dimension
        }
    
    def execute(
        self,
        image: Union[str, Path, np.ndarray],
        mask: Optional[Union[str, Path, np.ndarray]] = None,
        label: Optional[int] = _UNSET # type: ignore[assignment]
    ) -> Dict[str, np.ndarray]:
        """
        Execute barcode extraction pipeline.
        
        This method performs the complete extraction pipeline:
        1. Preprocess image (load, resample, window, normalize, crop, mask)
        2. Compute persistent homology
        3. Format barcodes by dimension
        
        Parameters
        ----------
        image : str, Path, or numpy.ndarray
            Input image. Can be a file path or numpy array.
        mask : str, Path, numpy.ndarray, or None, default=None
            Optional mask. Can be a file path or numpy array.
        label : int or None, optional
            Label value to extract from mask. If None, uses instance default.
            
        Returns
        -------
        barcodes : dict
            Dictionary with keys as homology dimensions (e.g., 'H0', 'H1', 'H2')
            and values as barcode arrays of shape (n_features, 2) containing
            (birth, death) pairs.
            
        Examples
        --------
        >>> extractor = BarcodeExtractor(normalize=True)
        >>> barcodes = extractor.execute('image.nii.gz', 'mask.nii.gz')
        >>> print(f"H1 has {len(barcodes['H1'])} features")
        """
        # Step 1: Preprocess image
        # Pass paths directly to preprocessor to preserve metadata (e.g., spacing)
        preprocessed_image, preprocess_metadata = self.preprocessor.preprocess(
            image, mask, label
        )
        
        # Step 2: Determine maxdim for PH computation
        if self.max_dimension == -1:
            maxdim = preprocessed_image.ndim - 1
        else:
            maxdim = self.max_dimension
        
        # Step 3: Compute persistent homology
        ph_array, dgms = compute_ph(
            preprocessed_image,
            maxdim=maxdim,
            construction=self.construction,
            filtration=self.filtration_type,
            max_value=None
        )
        
        # Step 4: Format barcodes by dimension
        barcodes = {}
        
        # Include all dimensions from 0 to maxdim
        for dim in range(maxdim + 1):
            dim_key = f'H{dim}'
            if dim < len(dgms):
                # Extract birth-death pairs for this dimension
                # dgms[dim] has shape (n_features, 2) - already in correct format
                barcodes[dim_key] = dgms[dim]
            else:
                # No features for this dimension
                barcodes[dim_key] = np.empty((0, 2))
        
        return barcodes
