import pandas as pd
from imblearn.under_sampling import RandomUnderSampler
from sklearn import dummy
from sklearn.feature_selection import RFECV, mutual_info_classif
from sklearn.impute import SimpleImputer
from sklearn.metrics import confusion_matrix, accuracy_score, f1_score, matthews_corrcoef
from sklearn.model_selection import StratifiedKFold, StratifiedShuffleSplit
from sklearn.preprocessing import LabelBinarizer, OneHotEncoder
from sklearn.utils.multiclass import unique_labels
from sklearn.utils.validation import check_X_y, check_array, check_is_fitted
from sklearn_pandas import DataFrameMapper, gen_features

from plots.plots_benchmark_mode import plot_metrics
from plots.plots_machine_learning import *
from preprocessing import utils_tools


class Metapredictions(object):
    """
    Class that uses machine learning techniques to evaluate prediction tools.
    """

    def __init__(self, df: pd.DataFrame,
                 location: str,
                 thresholds: list,
                 out_dir: str,
                 top_tools: dict,
                 missing_rate_allowed: float = 0.3,
                 rank_metric: str = "weighted_accuracy",
                 undersample: bool = False,
                 only_numeric: bool = True):
        """
        :param pd.DataFrame df: Input data
        :param str location: Location of the variants associated
            with the input dataframe
        :param list thresholds: Reference thresholds of the tools
        :param str out_dir: Ouptut dir
        :param dict top_tools: Top tools generated by the performance
            comparison step (per each location), ranked by the
            metric provided by the user.
        :param float missing_rate_allowed: Fraction of missing values
            allowed per class in a tool so that it is not discarded from the
            analysis. Default: `0.3`. Tools with bigger rate will
            not be used. If missing rate is lower than this value,
            tool will be included and missing values will be imputed
            with median.
        :param str rank_metric: Metric to rank classifier performance
            to compare with individual scores provided in `top_tools`
            arg. Default: "weighted_accuracy".
        :param bool undersample: Undersample data based on majority class
        (useful for unbalanced data)
        :param bool only_numeric: Include only numerical features (tool scores).
            Default: `False`. If `True`, categorical and boolean features
            will not be included.
        """
        # Remove useless columns
        useless_cols = [c for c in ['index', 'chr', 'pos', 'ref', 'alt', 'id', 'HGVSc'] if c in df.columns]
        
        df = df.drop(useless_cols, axis=1)

        self.label_fields = ['label']  # 'outcome' would also work
        self.categorical_fields = ['type', 'subtype', 'Gene', 'Consequence']

        self.out_dir = out_dir
        self.location = location
        self.top_tools = top_tools
        self.rank_metric = rank_metric
        self.only_numeric = only_numeric

        # Remove tools with too many missing data in any of the classes
        _patho = df[df.label].copy()
        _benign = df[~df.label].copy()
        
        na_patho = set(list(df)).difference(set(list(_patho.dropna(thresh=_patho.shape[0] * missing_rate_allowed, axis=1))))
        na_benign = set(list(df)).difference(set(list(_benign.dropna(thresh=_benign.shape[0] * missing_rate_allowed, axis=1))))
        to_remove = list(na_patho.union(na_benign))
        logging.info('Features (tool scores) that will not be used because they display high rate of missing predictions (> {}): {}'.format(missing_rate_allowed,
                                                                                                                                           to_remove))
        
        self.df = df.drop(to_remove, axis=1)
 
        # Update threshold list
        self.thresholds = [t for t in thresholds if t[0] in self.df.columns]

        # Remove score transformation in both S-CAP and TrAP,
        # if still present after NA removal
        self.s_cap_pred, self.trap_pred = [], []

        if "S-CAP" in self.df.columns:
            self.s_cap_pred = self.df['S-CAP'].copy().to_numpy()
            self.df.loc[:, 'S-CAP'] = self.df['S-CAP'].apply(lambda x: x - 1 if x > 1 else x)

        if "TraP" in self.df.columns:
            self.trap_pred = self.df['TraP'].copy().to_numpy()
            self.df.loc[:, 'TraP'] = self.df['TraP'].apply(lambda x: x - 1 if x > 1 else x)

        mapper_X, mapper_y = self.build_converter()
        self.X = mapper_X.fit_transform(self.df)
        self.y = mapper_y.fit_transform(self.df).ravel()
        self.features = mapper_X.transformed_names_

        if undersample:
            self.X_resampled, self.y_resampled = self.undersample_data()

        data_to_write = np.append(self.X, self.y.reshape(self.y.shape[0], -1), axis=1)
        pd.DataFrame(data_to_write, columns=self.features + ['label']).to_csv(os.path.join(self.out_dir, 'ml_processed_data_used.tsv'), sep="\t", index=False)       
    
    def build_converter(self):
        """
        Prepares a mapper between Pandas Dataframe and sklearn matrix
        """

        label_encoding = gen_features(columns=self.label_fields,
                                      classes=[LabelBinarizer])

        categorical = gen_features(columns=[[f] for f in self.categorical_fields],
                                   classes=[{'class': SimpleImputer, 'strategy': "most_frequent"},
                                            {'class': OneHotEncoder, 'sparse': False}],
                                   suffix="_cat")

        numeric = gen_features(columns=[[t[0]] for t in self.thresholds],
                               classes=[{'class': SimpleImputer, 'strategy': 'median'}])

        # if any boolean
        boolean = []
        for f in self.df.columns:
            if self.df[f].dtype == bool and f not in self.label_fields:
                self.df[f] = pd.to_numeric(self.df[f])
                boolean.append(([f], None))

        mapper_X = DataFrameMapper(numeric) if self.only_numeric else DataFrameMapper(categorical + boolean + numeric)
        mapper_y = DataFrameMapper(label_encoding)
        return mapper_X, mapper_y

    def undersample_data(self):
        """
        Performs random undersample of the majority class
        """

        rus = RandomUnderSampler(sampling_strategy="majority",
                                 random_state=0)

        X_resampled, y_resampled = rus.fit_sample(self.X, self.y)
        return X_resampled, y_resampled

    def do_classification(self):
        """
        Split input data and train several well known
        classifiers using a gridsearch approach. More
        evaluations are done with test data which is
        also used to compare against individual tools

        :return dict: Dictionary with trained pipelines
            for each of the tested classifiers
        """
        # Available metrics with: sklearn.metrics.SCORERS.keys()
        scoring_metrics = {'AUC': 'roc_auc',
                           'Accuracy': 'accuracy',
                           'F1_score': 'f1'}
        refit_metric = 'F1_score'

        trained_models = {}
        # Split training and test data to have equally
        # balanced classes in both training and test
        sss = StratifiedShuffleSplit(n_splits=1, test_size=0.2 , random_state=0)

        for train_idx, test_idx in sss.split(self.X, self.y):
            train_X, test_X = self.X[train_idx], self.X[test_idx]
            train_y, test_y = self.y[train_idx], self.y[test_idx]

            #############################
            ## Machine Learning class ###
            #############################
            clf = Classifiers(train_X, train_y, scoring_metrics, refit_metric)

            classifier_list = [#'KNN', 
                               'Naive_Bayes',
                               'LDA',
                               'QDA',
                               'Decision_Tree',
                               'Random_Forest',
                               'Ada_Boost',
                               'SVM',
                               'Logistic_Regression']

            # , 'Genetic_Programming']

            out_metrics = {}
            out_best_params = []
            for _c in classifier_list:
                pipeline_out, best_params, cv_score = clf.train(_c)
                out_best_params.append([_c, ';'.join(['{}={}'.format(k,v) for k, v in best_params.items()]) if best_params else None, cv_score])
                y_pred = pipeline_out.predict(test_X)

                # Update dict like in top_tools
                # (repeated metrics because coverage is 1, thus weighted
                # metrics are the same as the original metric value
                out_metrics[_c] = self.make_metrics(test_y, y_pred)
                trained_models[_c] = pipeline_out

            _clf_performance = self._dict_to_df(out_metrics)
            out_file = os.path.join(self.out_dir, "{}_gridSearch_output.tsv".format(self.location))
            pd.DataFrame.from_records(out_best_params, columns=['Classifier', 'Parameters_of_best_estimator', 'Cross_validation_score']).to_csv(out_file,
                                                                                                                                                index=False,
                                                                                                                                                sep='\t')
            
            #####################################
            ### Top tools overall performance ###
            #####################################
            top_tools_loc = self.top_tools[self.location]
            merged = pd.concat([_clf_performance, top_tools_loc], ignore_index=True)

            assert self.rank_metric in _clf_performance.columns, "Ranking metric provided is not available."

            out_file = os.path.join(self.out_dir, "{}_metrics_clf_along_best_tools.pdf".format(self.location))
            plot_metrics(merged, out_file, self.rank_metric)

            ##########################################
            ### Top tools performance on TEST data ###
            ##########################################
            # Train Single tool classifier for fair comparison
            single_tool_classifiers = self.expand_classifiers_list()
            out_metrics = {}
            for _c in single_tool_classifiers:
                _name = _c[0]
                assert _name in self.features, "{} scores must be present in the " \
                                               "training data to predict with single " \
                                               "feature classifier".format(_name)
                _clf = _c[1]
                _clf.fit(train_X, train_y)

                # Recover pathogenic predictions made by multi threshold tools
                # (pathogenic predictions transformed (+1))
                if _name == "TraP":

                    m = np.median(self.trap_pred[~np.isnan(self.trap_pred)])
                    self.trap_pred[np.isnan(self.trap_pred)] = m
                    test_x_single_feature = self.trap_pred[test_idx].reshape(-1, 1)

                elif _name == "S-CAP" or _name == "SCAP":
                    m = np.median(self.s_cap_pred[~np.isnan(self.s_cap_pred)])
                    self.s_cap_pred[np.isnan(self.s_cap_pred)] = m
                    test_x_single_feature = self.s_cap_pred[test_idx].reshape(-1, 1)

                else:
                    i = self.features.index(_name)
                    test_x_single_feature = np.expand_dims(test_X[:, i], axis=1)

                y_pred = _clf.predict(test_x_single_feature)
                out_metrics[_name] = self.make_metrics(test_y, y_pred)

            single_tool_perfor = self._dict_to_df(out_metrics)
            merged = pd.concat([_clf_performance, single_tool_perfor], ignore_index=True)

            out_file = os.path.join(self.out_dir, "{}_metrics_clf_only_on_test_data.pdf".format(self.location))
            merged.to_csv(os.path.join(self.out_dir, "{}_metrics_clf_only_on_test_data_N_variants_equals_to_{}.tsv".format(self.location, test_X.shape[0])), sep="\t", index=False)
            plot_metrics(merged, out_file, self.rank_metric)
            
        return trained_models

    def expand_classifiers_list(self):
        """
        Expand classifiers list considering the top individual
        tools as single classifiers.
        :return:
        """
        extra_classifiers = []
        _df_top = self.top_tools[self.location]
        for tool, direction, threshold, *args in self.thresholds:

            if tool in _df_top.tool.values:
                single_cls = SingleFeatureClassifier(tool, direction, threshold)
                extra_classifiers.append((tool, single_cls))

        return extra_classifiers

    def make_metrics(self, y_true: np.array, y_pred: np.array) -> list:
        """
        Generate common metrics for a given classifier

        :param np.array y_true: True labels
        :param np.array y_pred: Predicted labels

        :return list: List with some metrics
        """

        tn, fp, fn, tp = confusion_matrix(y_true, y_pred).ravel()
        specificity = round(utils_tools.ratio(tn, tn + fp), 3)
        sensitivity = round(utils_tools.ratio(tp, tp + fn), 3)
        precision = round(utils_tools.ratio(tp, tp + fp), 3)
        accuracy = round(accuracy_score(y_true, y_pred), 3)
        f1 = round(f1_score(y_true, y_pred), 3)
        coverage = 1
        mcc = round(matthews_corrcoef(y_true, y_pred), 3)
        normalized_mcc = (mcc + 1) / 2
        weighted_norm_mcc = round(normalized_mcc * coverage, 3)

        
        return [accuracy, accuracy, f1, f1, coverage,
                specificity, sensitivity, precision,
                mcc, normalized_mcc, weighted_norm_mcc]

    def _dict_to_df(self, input_d: dict) -> pd.DataFrame:
        """
        Converts a dict of statistics into a df

        :param dict input_d: Input dict
        :return pd.DataFrame: Output df
        """
        return pd.DataFrame.from_dict(input_d,
                                      orient='index',
                                      columns=['weighted_accuracy',
                                               'accuracy',
                                               'weighted_F1',
                                               'F1',
                                               'coverage',
                                               'specificity',
                                               'sensitivity',
                                               'precision',
                                               'mcc',
                                               'normalized_mcc',
                                               'weighted_norm_mcc']).rename_axis('tool').reset_index()

    def generate_feature_ranking(self, trained_clfs: dict):
        """
        Apply methodologies to rank features
        and identify the most important tools for predictions

        :param dict trained_clfs: Dict with trained classifiers
        """
        self.out_dir = os.path.join(self.out_dir, "feature_ranking")
        os.makedirs(self.out_dir, exist_ok=True)
        self.information_gain()

        try:
            rf_pipeline = trained_clfs['Random_Forest']
            self.feature_importance(rf_pipeline)

        except KeyError:
            pass

        self.recursive_feature_elimination(trained_clfs)

    def information_gain(self):
        """
        Calculates mutual information for the target label.
        Uses all the data
        """
        ig = mutual_info_classif(self.X, self.y, discrete_features=False, random_state=0)
        indices = np.argsort(ig)
        sorted_tools = [self.features[i] for i in indices]
        plot_feature_importance(dict(zip(sorted_tools, ig[indices])),
                                'information_gain',
                                self.location,
                                self.out_dir)

    def feature_importance(self, rf_pipeline: Pipeline):
        """
        Calculates feature importance scores based
        on the random forest pipeline and generates
        corresponding plot

        :param rf_pipeline: A trained Random forest pipeline
        :return:
        """
        _clf = rf_pipeline.named_steps['rf']
        indices = np.argsort(_clf.feature_importances_)
        sorted_tools = [self.features[i] for i in indices]
        importance_map = dict(zip(sorted_tools, _clf.feature_importances_[indices]))
        std = np.std([_clf.feature_importances_ for _tree in _clf.estimators_], axis=0)[indices]
        out_dir = os.path.join(self.out_dir, "random_forest")
        os.makedirs(out_dir, exist_ok=True)
        plot_feature_importance(importance_map,
                                'feature_importance',
                                self.location,
                                out_dir,
                                std=std)

    def recursive_feature_elimination(self, trained_clfs: dict):
        """
        Performs recursive feature elimination by fitting some
        of the classifiers provided and iteratively removing one
        or a subset of features.

        Fitting is done with best parameters found during the
        gridsearch for initial classification

        :param dict trained_clfs: Dict with trained classifiers
        """
        for _name, _clf in trained_clfs.items():
            new_clf, selector = "", ""
            if _name == "Random_Forest":
                out_dir = os.path.join(self.out_dir, 'random_forest')
                os.makedirs(out_dir, exist_ok=True)
                best_params = _clf.named_steps['rf'].get_params()
               
                new_clf = RandomForestClassifier(**best_params)
                selector = RFECV(new_clf,
                                 min_features_to_select=1,
                                 cv=StratifiedKFold(),
                                 n_jobs=-1)

                selector = selector.fit(self.X, self.y)

            elif _name == "Logistic_Regression":
                out_dir = os.path.join(self.out_dir, 'logistic_regression')
                os.makedirs(out_dir, exist_ok=True)
                # Couldn't use a pipeline in the RFECV, therefore
                # data needs to be scaled separately here
                scaler = StandardScaler()
                _X = scaler.fit_transform(self.X)
                best_params = _clf.named_steps['lr'].get_params()
        
                new_clf = LogisticRegression(**best_params)

                selector = RFECV(new_clf,
                                 min_features_to_select=1,
                                 cv=StratifiedKFold(),
                                 n_jobs=-1)

                selector = selector.fit(_X, self.y)

            else:
                continue

            plot_rfecv(selector, 1, self.features,
                       _name, self.location, out_dir)

    def generate_tree_plot(self, trained_clfs: dict):
        """
        Generate plot from a fitted decision tree model

        :param dict trained_clfs: Dict with trained classifiers
        """
        plot_decision_tree(trained_clfs['Decision_Tree'],
                           self.X,
                           self.y,
                           self.features,
                           self.location,
                           self.out_dir)


class SingleFeatureClassifier(dummy.DummyClassifier):
    """
    Class to create a Single Feature classifier
    based on the scores of individual tools
    """

    def __init__(self, tool, direction, threshold):
        super(SingleFeatureClassifier, self).__init__()
        self.tool = tool
        self.direction = direction
        self.threshold = float(threshold)

    def meaning(self, x):
        """
        Dummy decision function

        :param float x: Feature value of a single instance
        :return:
        """
        if self.direction == ">":
            return True if x >= self.threshold else False
        else:
            return True if x < self.threshold else False

    def fit(self, X: np.array, y: np.array, **kwargs):
        """
        Fit a dummy estimator without actually training anything

        :param np.array X: Training features
        :param np.array y: Training labels
        :param kwargs: Extra args
        :return:
        """
        X, y = check_X_y(X, y)
        self.classes_ = unique_labels(y)
        self.X_ = X
        self.y_ = y
        return self

    def predict(self, X: np.array) -> np.array:
        """
        Makes prediction over an array with single feature values
        :param np.array X: Input array
        :return np.array: Array with predictions
        """
        check_is_fitted(self)
        X = check_array(X)
        return np.array([self.meaning(x) for x in X]).astype(int)
