"""
Feature extraction module for MedTDA.

This module provides the FeatureExtractor class for extracting vectorized
TDA features from medical images.
"""

import numpy as np
from typing import Union, Optional, List, Dict, Tuple, Any, Callable
from pathlib import Path
from collections import OrderedDict

from .barcodeextractor import BarcodeExtractor, _UNSET
from . import vectorizers


class FeatureExtractor:
    """
    Extract vectorized TDA features from medical images.
    
    This class handles the complete pipeline: image loading, preprocessing,
    persistent homology computation, and vectorization of barcodes into
    fixed-length feature vectors.
    
    Parameters
    ----------
    spacing : tuple or None, default=None
        Target voxel spacing for resampling (3D/4D only).
    window : tuple or None, default=None
        Windowing parameters as (center, width).
    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.
    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.
        If -1, automatically determined based on image dimensionality.
    return_barcodes : bool, default=False
        If True, return computed barcodes along with features.
    vectorization_method : str or list, default='PersStats'
        Single method name or list of method names to apply.
        Options: 'BettiCurve', 'EntropySummary', 'PersStats', 
        'PersTropicalCoordinates', 'PersLandscape', 'PersImage', 
        'PersLifespan', 'PersSilhouette'
        
    Examples
    --------
    >>> from medtda import FeatureExtractor
    >>> extractor = FeatureExtractor(
    ...     normalize=True,
    ...     vectorization_method='PersStats'
    ... )
    >>> features = extractor.execute('image.nii.gz', 'mask.nii.gz')
    >>> print(features.keys())  # e.g., dict_keys(['PersStats_H0_mean', ...])
    """
    
    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,
        return_barcodes: bool = False,
        vectorization_method: Union[str, List[str]] = 'PersStats'
    ):
        """Initialize FeatureExtractor with parameters."""
        # Initialize BarcodeExtractor for PH computation
        self.barcode_extractor = BarcodeExtractor(
            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,
            filtration_type=filtration_type,
            construction=construction,
            max_dimension=max_dimension
        )
        
        # Store vectorization parameters
        self.return_barcodes = return_barcodes
        self._vectorization_methods: Dict[str, bool] = {}
        self._vectorization_params: Dict[str, Dict[str, Any]] = {}
        
        # Initialize vectorization methods
        if isinstance(vectorization_method, str):
            vectorization_method = [vectorization_method]
        
        for method in vectorization_method:
            self._validate_vectorization_method(method)
            self._vectorization_methods[method] = True
            self._vectorization_params[method] = {}
    
    def _validate_vectorization_method(self, method: str) -> None:
        """Validate that vectorization method exists."""
        if method not in vectorizers.VECTORIZATION_METHODS:
            available = list(set([k for k in vectorizers.VECTORIZATION_METHODS.keys() 
                                  if k[0].isupper()]))  # Get PascalCase names
            raise ValueError(
                f"Unknown vectorization method '{method}'. "
                f"Available methods: {available}"
            )
    
    # Preprocessing parameter setters (delegate to BarcodeExtractor)
    def set_spacing(self, spacing: Optional[Tuple[float, ...]]) -> None:
        """Update spacing parameter."""
        self.barcode_extractor.set_spacing(spacing)
    
    def set_window(self, window_center: float, window_width: float) -> None:
        """Set windowing parameters."""
        self.barcode_extractor.set_window(window_center, window_width)
    
    def set_normalize(self, enabled: bool, method: str = 'minmax') -> None:
        """Enable/disable normalization and set method."""
        self.barcode_extractor.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.barcode_extractor.set_clip_percentiles(percentiles)
    
    def set_background_value(self, value: Optional[float]) -> None:
        """Set background value for masked regions."""
        self.barcode_extractor.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.barcode_extractor.set_label(label)
    
    def set_filtration_type(self, filtration_type: str) -> None:
        """Set persistent homology filtration type."""
        self.barcode_extractor.set_filtration_type(filtration_type)
    
    def set_construction(self, construction: str) -> None:
        """Set cubical complex construction type."""
        self.barcode_extractor.set_construction(construction)
    
    def set_max_dimension(self, max_dimension: int) -> None:
        """Set maximum homology dimension to compute."""
        self.barcode_extractor.set_max_dimension(max_dimension)
    
    def set_return_barcodes(self, return_barcodes: bool) -> None:
        """
        Enable/disable barcode return.
        
        Parameters
        ----------
        return_barcodes : bool
            If True, execute() returns (features, barcodes) tuple.
        """
        self.return_barcodes = return_barcodes
    
    # Vectorization method management
    def set_vectorization_method(self, method: str, **params) -> None:
        """
        Set a single vectorization method with optional parameters.
        
        This replaces all currently enabled methods with just this one.
        
        Parameters
        ----------
        method : str
            Vectorization method name.
        **params
            Method-specific parameters.
        """
        self._validate_vectorization_method(method)
        self._vectorization_methods = {method: True}
        self._vectorization_params = {method: params}
    
    def set_vectorization_params(self, method: str, **params) -> None:
        """
        Update parameters for a vectorization method.
        
        Parameters
        ----------
        method : str
            Vectorization method name.
        **params
            Method-specific parameters to update.
        """
        if method not in self._vectorization_methods:
            raise ValueError(
                f"Method '{method}' is not enabled. "
                f"Enable it first using enable_vectorization_methods()."
            )
        self._vectorization_params[method].update(params)
    
    def enable_vectorization_methods(self, methods: List[str]) -> None:
        """
        Enable multiple vectorization methods.
        
        Parameters
        ----------
        methods : list of str
            List of method names to enable.
        """
        for method in methods:
            self._validate_vectorization_method(method)
            self._vectorization_methods[method] = True
            if method not in self._vectorization_params:
                self._vectorization_params[method] = {}
    
    def enable_all_vectorization_methods(self) -> None:
        """Enable all available vectorization methods."""
        # Get unique PascalCase method names
        all_methods = [k for k in vectorizers.VECTORIZATION_METHODS.keys() 
                      if k[0].isupper()]
        self.enable_vectorization_methods(all_methods)
    
    def disable_vectorization_method(self, method: str) -> None:
        """
        Disable a specific vectorization method.
        
        Parameters
        ----------
        method : str
            Method name to disable.
        """
        if method in self._vectorization_methods:
            del self._vectorization_methods[method]
            if method in self._vectorization_params:
                del self._vectorization_params[method]
    
    def get_enabled_vectorization_methods(self) -> List[str]:
        """
        Get list of currently enabled vectorization methods.
        
        Returns
        -------
        list of str
            Enabled method names.
        """
        return list(self._vectorization_methods.keys())
    
    def get_vectorization_params(self, method: str) -> Dict:
        """
        Get current parameters for a vectorization method.
        
        Parameters
        ----------
        method : str
            Method name.
            
        Returns
        -------
        dict
            Current parameters for the method.
        """
        if method not in self._vectorization_params:
            return {}
        return self._vectorization_params[method].copy()
    
    def get_settings(self) -> Dict:
        """
        Get current settings for extraction.
        
        Returns
        -------
        dict
            Dictionary containing all current settings.
        """
        barcode_settings = self.barcode_extractor.get_settings()
        return {
            **barcode_settings,
            'return_barcodes': self.return_barcodes,
            'vectorization_methods': self.get_enabled_vectorization_methods(),
            'vectorization_params': self._vectorization_params.copy()
        }
    
    def execute(
        self,
        image: Union[str, Path, np.ndarray],
        mask: Optional[Union[str, Path, np.ndarray]] = None,
        label: Optional[int] = _UNSET  # type: ignore[assignment]
    ) -> Union[Dict[str, Any], Tuple[Dict[str, Any], Dict[str, np.ndarray]]]:
        """
        Execute feature extraction pipeline.
        
        This method performs:
        1. Preprocess image (via BarcodeExtractor)
        2. Compute persistent homology (get barcodes)
        3. Vectorize barcodes using enabled methods
        4. Format features with naming convention
        
        Parameters
        ----------
        image : str, Path, or numpy.ndarray
            Input image.
        mask : str, Path, numpy.ndarray, or None, default=None
            Optional mask.
        label : int or None, optional
            Label value to extract from mask. If None, uses instance default.
            
        Returns
        -------
        features : OrderedDict
            Feature dictionary with keys following naming convention:
            - Dict-based methods (e.g., PersStats): '{method}_{dimension}_{feature_name}'
            - Array-based methods (e.g., BettiCurve): '{method}_{dimension}_f{index}'
        barcodes : dict (optional)
            If return_barcodes=True, returns (features, barcodes) tuple.
            
        Examples
        --------
        >>> extractor = FeatureExtractor(
        ...     normalize=True,
        ...     vectorization_method=['PersStats', 'BettiCurve']
        ... )
        >>> features = extractor.execute('image.nii.gz')
        >>> print(list(features.keys())[:5])  # First 5 feature names
        """
        # Step 1: Extract barcodes using BarcodeExtractor
        barcodes = self.barcode_extractor.execute(image, mask, label)
        
        # Step 2: Vectorize barcodes for each enabled method
        features = dict()
        
        for method_name in self.get_enabled_vectorization_methods():
            method_params = self._vectorization_params.get(method_name, {})
            method_func: Callable[..., Any] = vectorizers.VECTORIZATION_METHODS[method_name]  # type: ignore[assignment,index]
            
            # Apply vectorization to each dimension
            for dim_key, barcode in barcodes.items():
                # Get feature vector for this dimension
                feature_vector = method_func(barcode, **method_params)
                
                # Add features with proper naming
                if isinstance(feature_vector, dict):
                    # Method returns named features (like PersStats)
                    for feat_name, feat_value in feature_vector.items():
                        full_name = f"{method_name}_{dim_key}_{feat_name}"
                        features[full_name] = feat_value
                else:
                    # Method returns array - create indexed names
                    for i, feat_value in enumerate(feature_vector):
                        full_name = f"{method_name}_{dim_key}_f{i}"
                        features[full_name] = feat_value
        
        # Step 3: Return features (and optionally barcodes)
        if self.return_barcodes:
            return features, barcodes
        else:
            return features
