"""
File contains class used for feature selection.

Author: Dominik Macko
"""

from androidcrypto.analysis.preparation.utils import ignore_cols
from androidcrypto.analysis.preparation.config import FeatureSelectionConfig
from androidcrypto.analysis.preparation.saving import save_to_hdf
from typing import Union, Optional, List, Dict, Tuple

import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.pipeline import Pipeline
from sklearn.ensemble import RandomForestClassifier
from sklearn.feature_selection import VarianceThreshold
from boruta import BorutaPy

class FeatureSelectionCollinearTransformer(BaseEstimator, TransformerMixin):
    """
    Removes correlated features.
    
    from: https://github.com/WillKoehrsen/feature-selector/blob/master/feature_selector/feature_selector.py
    taken and slightly tweaked for this use case and to not require given dependencies
    """

    def __init__(self,
                 correlation_threshold: float=0.95,
                 method: str="pearson" # pearson/kendall/spearman
        ):

        self._correlation_threshold = correlation_threshold
        self._method = method
        self._to_drop = None

        self.corr_matrix = None
        self.record_collinear = None
        
    def fit(self,
            X: Union[pd.DataFrame, np.array],
            y: Optional[Union[pd.DataFrame, pd.Series]]=None
           ) -> "FeatureSelectionCollinearTransformer":
        """Fits given training data, must be used only after being fit.
        
        After fit these attributes can be accessed:
            self.corr_matrix
            self.record_collinear
        """
    
        corr_matrix = X.corr(method=self._method)
        self.corr_matrix = corr_matrix

        # Extract the upper triangle of the correlation matrix
        upper = corr_matrix.where(np.triu(np.ones(corr_matrix.shape), k = 1).astype(np.bool))

        # Select the features with correlations above the threshold
        # Need to use the absolute value
        to_drop = [column for column in upper.columns if any(upper[column].abs() > self._correlation_threshold)]
        self._to_drop = to_drop

        # Dataframe to hold correlated pairs
        record_collinear = pd.DataFrame(columns = ['drop_feature', 'corr_feature', 'corr_value'])

        for column in to_drop:

            # Find the correlated features
            corr_features = list(upper.index[upper[column].abs() > self._correlation_threshold])

            # Find the correlated values
            corr_values = list(upper[column][upper[column].abs() > self._correlation_threshold])
            drop_features = [column for _ in range(len(corr_features))]    

            # Record the information (need a temp df for now)
            temp_df = pd.DataFrame.from_dict({"drop_feature": drop_features,
                                             "corr_feature": corr_features,
                                             "corr_value": corr_values})

            # Add to dataframe
            record_collinear = record_collinear.append(temp_df, ignore_index = True)

        self.record_collinear = record_collinear

        print("%d features with a correlation magnitude greater than %0.2f.\n" % (len(to_drop),
                                                                                  self._correlation_threshold))

        return self
        
    def transform(self,
                  X: pd.DataFrame,
                  y: Optional[Union[pd.DataFrame, pd.Series]]=None
                 ) -> pd.DataFrame:
        """Transforms given training data, must be used only after being fit."""
        
        if self._to_drop is None:
            print("FeatureSelectionCollinearTransformer was not fit before transform!")
            return X
        return X.drop(columns=self._to_drop)

class FeatureSelectionTransformer(BaseEstimator, TransformerMixin):
    """Feature selection transformer that uses variance threshold and Boruta.

    https://towardsdatascience.com/boruta-explained-the-way-i-wish-someone-explained-it-to-me-4489d70e154a
    """
    
    def __init__(self,
                 variance_threshold: float=0,
                 seed: int =42,
                 n_jobs: int=-1,
                 n_estimators: int=200,
                 verbose: int=1,
                 max_depth: int=5,
                 max_iter: int=100,
                 use_weak_boruta_support: bool=False,
                 correlation_threshold: float=0.95
        ):
        """Initializes feature selection transformer.
        
        variance_threshold - specifies cutoff threshold for variance
        seed - random state seed
        n_jobs - cores to use, -1 can be used for maximum, default -1
        n_estimators - number of trees to use, default 200
        max_depth - max depth of trees, default 5
        max_iter - maximum iterations of Boruta, default 50
        use_weak_boruta_support - specifies whether selection should be conservative and weaker features kept
        
        note: to not keep features with the same value in 99.5% of samples use variance_threshold=0.995 * (1 - 0.995)
        """
        
        self._variance_threshold = VarianceThreshold(variance_threshold)
        self._collinear_transformer = FeatureSelectionCollinearTransformer(correlation_threshold)
        self._boruta = BorutaPy(
            RandomForestClassifier(
                n_estimators=n_estimators, 
                n_jobs=n_jobs,
                max_depth=max_depth,
                class_weight="balanced"
            ),
            random_state=seed, 
            verbose=verbose, 
            max_iter=max_iter
        )
        self._use_weak_boruta_support: bool = use_weak_boruta_support   
        
    def fit(self,
            X: pd.DataFrame,
            y: Union[pd.DataFrame, pd.Series]
           ) -> "FeatureSelectionTransformer":
        """Fits transformer to given training data."""

        self._variance_threshold.fit(X, y.values.ravel())
        X_after_variance_threshold = X.loc[:, self._get_variance_threshold_mask()]

        self._collinear_transformer.fit(X_after_variance_threshold)
        X_after_collinear = self._collinear_transformer.transform(X_after_variance_threshold)
        
        self._boruta.fit(X_after_collinear.values, y.values.ravel())

        self._fit_columns = X.columns
        return self

    def transform(self,
                  X: pd.DataFrame,
                  y: Optional[Union[pd.DataFrame, pd.Series]]=None
                 ) -> pd.DataFrame:
        """Transforms given training data, must be used only after being fit."""
        
        # remove features with with low variance
        result = X.loc[:, self._get_variance_threshold_mask()] 
        # remove correlated features
        result = self._collinear_transformer.transform(result) 
        # remove features with low importance
        return result.loc[:, self._get_boruta_mask()] 

    def fit_transform(self,
                  X: pd.DataFrame,
                  y: Optional[Union[pd.DataFrame, pd.Series]]=None
                 ) -> pd.DataFrame:
        """Fits and transforms given data."""

        print(X.shape)
        self._variance_threshold.fit(X, y.values.ravel())
        X_after_variance_threshold = X.loc[:, self._get_variance_threshold_mask()]

        print(X_after_variance_threshold.shape)
        self._collinear_transformer.fit(X_after_variance_threshold)
        X_after_collinear = self._collinear_transformer.transform(X_after_variance_threshold)
        print(X_after_collinear.shape)
        self._boruta.fit(X_after_collinear.values, y.values.ravel())

        self._fit_columns = X.columns

        return X_after_collinear.loc[:, self._get_boruta_mask()] 
    
    def selected_columns(self) -> List[str]:
        """Gets list of selected columns."""
        
        return [*self._fit_columns[self._get_variance_threshold_mask()][self._get_boruta_mask()]]
        
    def _get_variance_threshold_mask(self) -> np.array:
        """Gets variance threshold mask."""
        
        return self._variance_threshold.get_support()
    
    def _get_boruta_mask(self) -> np.array:
        """Gets boruta support mask based on whether weak boruta support should be included."""
        
        if not self._use_weak_boruta_support:
            return self._boruta.support_ 
        return self._boruta.support_weak_ | self._boruta.support_

def feature_selection_based_on_config(
    dfs: Optional[Dict[str, Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]]],
    config: FeatureSelectionConfig
    ) -> Dict[str, Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]]:
    """Selects features based on a config file.
    
    returns train_X, train_y, test_X, test_y
    """

    print("Feature selection: Starting.")
    result = {}
    if dfs is None:
        dfs = {}
    used_dfs = dfs

    for task_name, task_config in config.task_configs.items():

        # load if None
        if dfs.get(task_name, None) is None:
            if not task_config.input_hdf_path:
                print(f"Feature selection: Task {task_name}: No data from previous step and input hdf path not defined.")
                continue
            print(f"Feature selection: Task {task_name}: Loading data.")
            used_dfs[task_name] = (pd.read_hdf(task_config.input_hdf_path, key="X_train"),
                                   pd.read_hdf(task_config.input_hdf_path, key="y_train"),
                                   pd.read_hdf(task_config.input_hdf_path, key="X_test"),
                                   pd.read_hdf(task_config.input_hdf_path, key="y_test"))

        X_train, y_train, X_test, y_test = used_dfs[task_name]
        
        # temporarily remove the columns
        X_train, ignored_train, columns_ignored_train = ignore_cols(X_train, config.columns_to_ignore)
        X_test, ignored_test, columns_ignored_test = ignore_cols(X_test, config.columns_to_ignore)
        
        # create Transformer and select
        print(f"Feature selection: Task {task_name}: Selecting features.")
        transformer = FeatureSelectionTransformer(
            n_estimators=config.n_trees,
            max_iter=config.max_iterations,
            use_weak_boruta_support=config.use_weak_features_boruta
        )
        X_train = transformer.fit_transform(X_train, y_train)
        X_test = transformer.transform(X_test)

        # put the ignored features back
        if columns_ignored_train:
            X_train[columns_ignored_train] = ignored_train
        if columns_ignored_test:
            X_test[columns_ignored_test] = ignored_test
        
        # optionally save for each of the tasks
        if task_config.output_hdf_path:
            print(f"Feature selection: Task {task_name}: Saving output.")
            save_to_hdf(task_config.output_hdf_path, X_train, y_train, X_test, y_test)

        result[task_name] = (X_train, y_train, X_test, y_test)

    return result
    