"""
Base Winnow
Performs core feature selection operations:
- Feature-to-feature correlation analysis
- Feature-to-target correlation analysis
- Variance analysis
"""

import numpy as np
import pandas as pd
from typing import List, Optional, Dict, Tuple, Union
from dataclasses import dataclass
from .logger import logger


@dataclass
class SelectionResult:
    """Class to hold feature selection results"""
    selected_features: List[str]
    removed_features: Dict[str, List[str]]  # reason -> feature list
    statistics: Dict[str, pd.DataFrame]  # statistics


class BaseSelector:
    """
    Base Winnow class.
    
    Uses three main criteria for feature selection:
    1. Inter-feature correlation (removes highly correlated features)
    2. Feature-Target correlation (removes low correlation features)
    3. Variance analysis (removes low variance features)
    
    Parameters
    ----------
    correlation_threshold : float, default=0.8
        Threshold for inter-feature correlation. One of the features 
        in a pair with correlation above this value will be removed.
    
    target_correlation_threshold : float, default=0.05
        Threshold for feature-target correlation. Features with 
        correlation below this value will be removed.
    
    variance_threshold : float, default=0.01
        Threshold for variance. Features with variance below 
        this value will be removed.
    
    Examples
    --------
    >>> import pandas as pd
    >>> from winnow import BaseSelector
    >>> 
    >>> # Sample data
    >>> X = pd.DataFrame({
    ...     'feature1': [1, 2, 3, 4, 5],
    ...     'feature2': [1.1, 2.1, 3.1, 4.1, 5.1],  # high correlation with feature1
    ...     'feature3': [5, 4, 3, 2, 1],
    ...     'feature4': [1, 1, 1, 1, 1]  # zero variance
    ... })
    >>> y = pd.Series([0, 0, 1, 1, 1])
    >>> 
    >>> selector = BaseSelector(
    ...     correlation_threshold=0.9,
    ...     target_correlation_threshold=0.1,
    ...     variance_threshold=0.01
    ... )
    >>> result = selector.fit_transform(X, y)
    >>> print(result.selected_features)
    """
    
    def __init__(
        self,
        correlation_threshold: float = 0.8,
        target_correlation_threshold: float = 0.05,
        variance_threshold: float = 0.01
    ):
        self.correlation_threshold = correlation_threshold
        self.target_correlation_threshold = target_correlation_threshold
        self.variance_threshold = variance_threshold
        
        # Stored values after fit
        self._feature_correlations: Optional[pd.DataFrame] = None
        self._target_correlations: Optional[pd.Series] = None
        self._variances: Optional[pd.Series] = None
        self._selected_features: Optional[List[str]] = None
        self._is_fitted: bool = False
        self._target_column: Optional[str] = None
        self._feature_columns: Optional[List[str]] = None
    
    def fit(
        self, 
        X: pd.DataFrame, 
        y: Optional[Union[pd.Series, str]] = None,
        target_column: Optional[str] = None
    ) -> 'BaseSelector':
        """
        Calculates feature selection criteria.
        
        Two usage patterns are supported:
        1. X and y are provided separately (sklearn style)
        2. X is provided as a single DataFrame, target_column specifies the target
        
        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix or DataFrame containing all data
        y : pd.Series or str, optional
            Target variable (Series) or target column name (str)
            If str, this column from X is used as target
        target_column : str, optional
            Name of the target column. Can be used instead of y.
            
        Returns
        -------
        self : BaseSelector
            Fitted selector instance
        """
        logger.info("Starting feature selection fit process...")
        
        # Parse target
        X_features, y_target = self._parse_target(X, y, target_column)
        
        self._validate_input(X_features, y_target)
        
        # Store metadata
        self._target_column = y_target.name if hasattr(y_target, 'name') else "target"
        self._feature_columns = X_features.columns.tolist()
        
        logger.debug(f"Calculating statistics for {len(self._feature_columns)} features...")
        
        # Calculate statistics
        self._variances = X_features.var()
        self._feature_correlations = X_features.corr()
        self._target_correlations = X_features.corrwith(y_target).abs()
        
        self._is_fitted = True
        logger.info("Fit process completed successfully.")
        return self
    
    def _parse_target(
        self, 
        X: pd.DataFrame, 
        y: Optional[Union[pd.Series, str]], 
        target_column: Optional[str]
    ) -> Tuple[pd.DataFrame, pd.Series]:
        """
        Parses X and y from input parameters.
        
        Returns
        -------
        Tuple[pd.DataFrame, pd.Series]
            (Feature DataFrame, Target Series)
        """
        # If target_column is provided
        if target_column is not None:
            if target_column not in X.columns:
                raise ValueError(f"Column '{target_column}' not found in DataFrame.")
            y_series = X[target_column]
            X_features = X.drop(columns=[target_column])
            return X_features, y_series
        
        # If y is a string (column name)
        if isinstance(y, str):
            if y not in X.columns:
                raise ValueError(f"Column '{y}' not found in DataFrame.")
            y_series = X[y]
            X_features = X.drop(columns=[y])
            return X_features, y_series
        
        # If y is a Series (classic usage)
        if isinstance(y, pd.Series):
            return X, y
        
        # Otherwise error
        raise ValueError(
            "Target not specified. Usage options:\n"
            "  1. fit(X, y) - X and y separately\n"
            "  2. fit(df, target_column='target_col') - DataFrame and column name\n"
            "  3. fit(df, y='target_col') - column name as y parameter"
        )
    
    def transform(self, X: pd.DataFrame, y: Optional[pd.Series] = None) -> SelectionResult:
        """
        Performs feature selection based on calculated criteria.
        
        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix
        y : pd.Series, optional
            Target variable (optional, only for statistics)
            
        Returns
        -------
        SelectionResult
            Selected features and statistics
        """
        self._check_is_fitted()
        logger.info("Starting feature transformation...")
        
        removed_features: Dict[str, List[str]] = {
            'low_variance': [],
            'high_correlation': [],
            'low_target_correlation': []
        }
        
        # Initial feature set
        candidate_features = set(X.columns.tolist())
        initial_count = len(candidate_features)
        
        # 1. Remove low variance features
        low_variance_features = self._get_low_variance_features()
        removed_features['low_variance'] = low_variance_features
        candidate_features -= set(low_variance_features)
        logger.debug(f"Removed {len(low_variance_features)} features due to low variance.")
        
        # 2. Remove low target correlation features
        low_target_corr_features = self._get_low_target_correlation_features(
            list(candidate_features)
        )
        removed_features['low_target_correlation'] = low_target_corr_features
        candidate_features -= set(low_target_corr_features)
        logger.debug(f"Removed {len(low_target_corr_features)} features due to low target correlation.")
        
        # 3. Remove one feature from high correlation pairs
        high_corr_features = self._get_high_correlation_features(
            list(candidate_features)
        )
        removed_features['high_correlation'] = high_corr_features
        candidate_features -= set(high_corr_features)
        logger.debug(f"Removed {len(high_corr_features)} features due to high inter-feature correlation.")
        
        # Results
        selected_features = list(candidate_features)
        
        statistics = {
            'variances': self._variances.to_frame('variance'),
            'feature_correlations': self._feature_correlations,
            'target_correlations': self._target_correlations.to_frame('target_correlation')
        }
        
        self._selected_features = selected_features
        
        logger.info(f"Transformation complete. Selected {len(selected_features)}/{initial_count} features.")
        
        return SelectionResult(
            selected_features=selected_features,
            removed_features=removed_features,
            statistics=statistics
        )
    
    def fit_transform(
        self, 
        X: pd.DataFrame, 
        y: Optional[Union[pd.Series, str]] = None,
        target_column: Optional[str] = None
    ) -> SelectionResult:
        """
        Performs fit and transform in one step.
        """
        self.fit(X, y, target_column)
        X_features, y_target = self._parse_target(X, y, target_column)
        return self.transform(X_features, y=y_target)
    
    def get_selected_dataframe(self, X: pd.DataFrame) -> pd.DataFrame:
        """
        Returns DataFrame containing only selected features.
        
        Parameters
        ----------
        X : pd.DataFrame
            Original feature matrix
            
        Returns
        -------
        pd.DataFrame
            DataFrame with selected features
        """
        self._check_is_fitted()
        if self._selected_features is None:
            raise ValueError("Transform has not been called yet. Call fit_transform() first.")
        return X[self._selected_features]
    
    def _get_low_variance_features(self) -> List[str]:
        """Returns features with variance below threshold."""
        return self._variances[
            self._variances < self.variance_threshold
        ].index.tolist()
    
    def _get_low_target_correlation_features(
        self, 
        candidate_features: List[str]
    ) -> List[str]:
        """Returns features with target correlation below threshold."""
        correlations = self._target_correlations[candidate_features]
        return correlations[
            correlations < self.target_correlation_threshold
        ].index.tolist()
    
    def _get_high_correlation_features(
        self, 
        candidate_features: List[str]
    ) -> List[str]:
        """
        Selects one of the features in high correlation pairs.
        Removes the one with lower target correlation.
        """
        to_remove = []
        remaining = set(candidate_features)
        
        # Get upper triangle of correlation matrix
        corr_matrix = self._feature_correlations.loc[
            candidate_features, 
            candidate_features
        ].abs()
        
        for i, feature1 in enumerate(candidate_features):
            if feature1 not in remaining:
                continue
                
            for feature2 in candidate_features[i+1:]:
                if feature2 not in remaining:
                    continue
                
                if corr_matrix.loc[feature1, feature2] > self.correlation_threshold:
                    # Remove the one with lower target correlation
                    corr1 = self._target_correlations.get(feature1, 0)
                    corr2 = self._target_correlations.get(feature2, 0)
                    
                    feature_to_remove = feature2 if corr1 >= corr2 else feature1
                    to_remove.append(feature_to_remove)
                    remaining.discard(feature_to_remove)
        
        return to_remove
    
    def _validate_input(self, X: pd.DataFrame, y: pd.Series) -> None:
        """Performs input validation."""
        if not isinstance(X, pd.DataFrame):
            raise TypeError("X must be a pandas DataFrame.")
        if not isinstance(y, (pd.Series, np.ndarray, list)):
            raise TypeError("y must be a pandas Series or array-like.")
        if len(X) != len(y):
            raise ValueError("X and y must have the same length.")
        if X.empty:
            raise ValueError("X cannot be empty.")
    
    def _check_is_fitted(self) -> None:
        """Checks if the selector is fitted."""
        if not self._is_fitted:
            raise ValueError("Selector is not fitted yet. Call fit() first.")
    
    def get_feature_statistics(self) -> pd.DataFrame:
        """
        Returns a DataFrame containing statistics for all features.
        
        Returns
        -------
        pd.DataFrame
            Statistics table including variance and target correlation
        """
        self._check_is_fitted()
        
        stats = pd.DataFrame({
            'variance': self._variances,
            'target_correlation': self._target_correlations
        })
        
        return stats.sort_values('target_correlation', ascending=False)
    
    @property
    def selected_features(self) -> Optional[List[str]]:
        """Returns the list of selected features."""
        return self._selected_features
    
    def __repr__(self) -> str:
        return (
            f"BaseSelector("
            f"correlation_threshold={self.correlation_threshold}, "
            f"target_correlation_threshold={self.target_correlation_threshold}, "
            f"variance_threshold={self.variance_threshold})"
        )
