import numpy as np

import pandas as pd

import matplotlib

matplotlib.use('Agg')

import matplotlib.pyplot as plt

from matplotlib.backends.backend_pdf import PdfPages

import seaborn as sns

sns.set()

from sklearn.base import BaseEstimator, ClassifierMixin

from sklearn.linear_model import LogisticRegression

from sklearn.ensemble import RandomForestClassifier

from sklearn.svm import LinearSVC

from sklearn.metrics import roc_auc_score, accuracy_score, roc_curve

from sklearn.preprocessing import StandardScaler, RobustScaler, MinMaxScaler, MaxAbsScaler

from sklearn.utils.class_weight import compute_sample_weight

from sklearn.model_selection import GridSearchCV

from sklearn.feature_selection import RFE

import os

from tqdm import tqdm

import pickle

import json

from scipy.cluster import hierarchy

from scipy.spatial.distance import pdist, squareform

from tqdm import tqdm

import logging

logging.basicConfig(level=logging.INFO, format='[%(asctime)s] [%(levelname)s] %(name)s: %(message)s')

def bootstrap(*arrays):

    '''Perform bootstrap resampling

    Parameters

    ----------

    arrays: list of array-like objects

        Data to resample. Bootstrap resampling is performed on the first dimension of each array

    Returns

    -------

    arrays: tuple of array-like objects

        Resampled data

    '''

    n_samples = arrays[0].shape[0]

    if not all(n_samples == a.shape[0] for a in arrays):

        raise ValueError('the first dimension of all arrays should be equal')

    indices = np.random.randint(n_samples, size=n_samples)

    return (np.take(a, indices, axis=0) for a in arrays)

def jackknife(*arrays, remove=1, indices=None):

    '''Perform jackknife resampling

    Parameters

    ----------

    arrays: list of array-like objects

        Data to resample. Bootstrap resampling is performed on the first dimension of each array

    remove: int or float

        If `remove` is a integer, remove that number of samples.

        If `remove` is a float, remove that fraction of samples.

    indices: array-like

        Indices of samples to remove. The `remove` parameter will be ignored.

    Returns

    -------

    arrays: tuple of array-like objects

        Resampled data

    '''

    n_samples = arrays[0].shape[0]

    if not all(n_samples == a.shape[0] for a in arrays):

        raise ValueError('the first dimension of all arrays should be equal')

    if indices is None:

        if remove < 1:

            n_remove = round(remove*n_samples)

        else:

            n_remove = remove

        indices = np.random.choice(n_samples, replace=False, size=n_remove)

    return (np.delete(a, indices, axis=0) for a in arrays)

class RobustEstimator(BaseEstimator, ClassifierMixin):

    '''A wrapper to select robust features based on feature reccurence

    Parameters:

    ----------

    estimator: sklearn.base.BaseEstimator object

        Base estimator for feature selection

    n_select: int

        Number of features to select

    resample_method: str, choices: ('jackknife', 'bootstrap')

        Resampling method

    remove: float or int

        Fraction (float, 0.0-1.0) or number (int, >= 1) of samples to remove for each round of resampling

    max_runs: int

        Maximum rounds of jackknife resampling

    recurring_fraction: float

        Minimum required fraction of reccuring samples to select features

    grid_search: dict or None

        Parameter grid for optimizing hyper-parameters using GridSearchCV

    rfe: bool

        Whether to use RFE to select features rather than select features with higest importance during each run

    Attributes:

    ----------

    feature_recurrence_: array-like, shape (n_features,)

        Number of resampling rounds in which each feature is retained

    features_: array-like, shape (n_features,)

        Indices of selected features

    feature_selection_matrix_: array-like, (max_runs, n_select), type int

        Indicator matrix for selected features during each run

    feature_importance_matrix_: array-like, (max_runs, n_features)

        Feature importance during each resampling run

    feature_importances_mean_: array-like, (n_features,)

        Feature importances averaged across resampling runs

    feature_importance_std_: array-like, (n_features,)

        Standard deviation of feature importances across resampling runs

    estimator_: sklearn.base.BaseEstimator object

        Estimator fitted on all data using selected features

    feature_importances_: array-like, (n_select,)

        Feature importances trained on all data using selected features

    '''

    def __init__(self, estimator, n_select=None, resample_method = 'jackknife',

            remove=1, max_runs=50, recurring_fraction=0.5,

            grid_search=None, rfe=False):

        self.estimator = estimator

        self.remove = remove

        self.max_runs = max_runs

        self.n_select = n_select

        self.resample_method = resample_method

        self.recurring_fraction = recurring_fraction

        self.grid_search = grid_search

        self.rfe = rfe

    def fit(self, X, y, sample_weight=None):

        n_samples, n_features = X.shape

        feature_rank_matrix = np.zeros(n_samples)

        if self.remove == 1:

            max_runs = n_samples

        else:

            max_runs = self.max_runs

        n_select = self.n_select

        if n_select is None:

            n_select = n_features

        feature_rank_matrix = np.zeros((max_runs, n_features))

        feature_importances_matrix = np.zeros((max_runs, n_features))

        if sample_weight is None:

            sample_weight = np.ones(n_samples)

        for i_run in range(max_runs):

            if self.resample_method == 'bootstrap':

                X_, y_, sample_weight_ = bootstrap(X, y, sample_weight)

            elif self.resample_method == 'jackknife':  

                indices_remove = None

                if self.remove == 1:

                    indices_remove = i_run

                X_, y_, sample_weight_ = jackknife(X, y, sample_weight,

                    remove=self.remove, indices=indices_remove)

            if self.grid_search is not None:

                cv = GridSearchCV(self.estimator, param_grid=self.grid_search, cv=3)

                cv.fit(X_, y_, sample_weight=sample_weight_)

                self.estimator = cv.best_estimator_

            else:

                self.estimator.fit(X_, y_, sample_weight=sample_weight_)

            if self.rfe:

                rfe = RFE(self.estimator, n_select, step=0.1)

                rfe.fit(X_, y_)

                feature_importances_matrix[i_run] = rfe.support_.astype('float')

                feature_rank_matrix[i_run] = rfe.ranking_

            else:

                if hasattr(self.estimator, 'coef_'):

                    feature_importances = np.square(self.estimator.coef_.flatten())

                else:

                    feature_importances = self.estimator.feature_importances_

                feature_importances_matrix[i_run] = feature_importances

                feature_orders = np.argsort(-feature_importances)

                feature_rank_matrix[i_run, feature_orders] = np.arange(n_features)

        if self.rfe:

            feature_selection_matrix = (feature_rank_matrix == 1).astype(np.int32)

        else:

            feature_selection_matrix = (feature_rank_matrix < n_select).astype(np.int32)

        feature_recurrence = np.sum(feature_selection_matrix, axis=0)

        #robust_features = np.nonzero(feature_recurrence > round(n_samples)*self.recurring_fraction)[0]

        #robust_features = robust_features[np.argsort(feature_recurrence[robust_features])][::-1]

        robust_features = np.argsort(-feature_recurrence)[:n_select]

        feature_selection_matrix = feature_selection_matrix[:, robust_features]

        # refit the estimator

        if self.grid_search is not None:

            cv = GridSearchCV(self.estimator, param_grid=self.grid_search, cv=3)

            cv.fit(X[:, robust_features], y, sample_weight=sample_weight)

            self.estimator = cv.best_estimator_

        else:

            self.estimator.fit(X[:, robust_features], y, sample_weight=sample_weight)

        # save attributes

        self.feature_recurrence_ = feature_recurrence

        self.features_ = robust_features

        self.feature_selection_matrix_ = feature_selection_matrix

        self.feature_importances_matrix_ = feature_importances_matrix

        self.feature_importances_mean_ = np.mean(feature_importances_matrix, axis=0)

        self.feature_importances_std_ = np.std(feature_importances_matrix, axis=0)

        self.estimator_ = self.estimator

        if hasattr(self.estimator, 'coef_'):

            self.feature_importances_ = np.square(self.estimator.coef_.flatten())

        else:

            self.feature_importances_ = self.estimator.feature_importances_

        return self

    def predict(self, X):

        return self.estimator_.fit(X)

    def score(self, X):

        return self.estimator_.score(X)

def define_step(inputs=None, outputs=None):

    inputs = inputs if inputs is not None else ()

    outputs = outputs if outputs is not None else ()

    def wrapper_generator(f):

        def wrapper(self, *args, **kwargs):

            for input in inputs:

                if not hasattr(self, input):

                    raise ValueError('missing input "{}" for step f.__name__'.format(input))

            r = f(self, *args, **kwargs)

            for output in outputs:

                if not hasattr(self, output):

                    raise ValueError('missing output "{}" for step f.__name__'.format(output))

            return r

        return wrapper

    return wrapper_generator

class FeatureSelection(object):

    def __init__(self, matrix, 

            sample_classes,

            output_dir,

            remove_zero_features=None,

            top_features_by_median=None,

            transpose=False,

            positive_class=None,

            negative_class=None,

            use_log=False,

            scaler=None,

            compute_sample_weight=False,

            method='logistic_regression',

            rfe=False,

            n_select=None,

            resample_method='jackknife',

            jackknife_max_runs=100,

            jackknife_remove=0.2,

            bootstrap_max_runs=100,

            **kwargs):

        self.matrix_file = matrix

        self.sample_classes_file = sample_classes

        self.output_dir = output_dir

        self.remove_zero_features = remove_zero_features

        self.top_features_by_median = top_features_by_median

        self.transpose = transpose

        self.positive_class = positive_class

        self.negative_class = negative_class

        self.use_log = use_log

        self.scaler = scaler

        self.compute_sample_weight = compute_sample_weight

        self.rfe = rfe

        self.method = method

        self.n_select = n_select

        self.resample_method = resample_method

        self.jackknife_max_runs = jackknife_max_runs

        self.jackknife_remove = jackknife_remove

        self.bootstrap_max_runs = bootstrap_max_runs

        self.logger = logging.getLogger('FeatureSelection')

        if not os.path.exists(self.output_dir):

            self.logger.info('create output directory: ' + self.output_dir)

            os.makedirs(self.output_dir)

    @define_step(inputs=['matrix_file'], outputs=['X', 'sample_classes'])

    def read_data(self):

        self.X = pd.read_table(self.matrix_file, index_col=0)

        if self.transpose:

            self.logger.info('transpose feature matrix')

            self.X = self.X.T

        self.logger.info('{} samples, {} features'.format(*self.X.shape))

        self.logger.info('read sample classes: ' + self.sample_classes_file)

        self.sample_classes = pd.read_table(self.sample_classes_file, header=None, 

            names=['sample_id', 'sample_class'], index_col=0)

        self.sample_classes = self.sample_classes.iloc[:, 0]

    @define_step(inputs=['X'], outputs=['feature_names', 'n_features'])

    def filter_features(self):

        if self.remove_zero_features is not None:

            self.X = self.X.loc[:, np.isclose(m, 0).sum(axis=0) >= (m.shape[0]*self.remove_zero_features)]

        if self.top_features_by_median is not None:

            nonzero_samples = (self.X > 0).sum(axis=0)

            counts_geomean = np.exp(np.sum(np.log(np.maximum(self.X, 1)), axis=0)/nonzero_samples)

            self.X = self.X.loc[:, counts_geomean.sort_values(ascending=False)[:self.top_features_by_median].index.values]

        self.feature_names = self.X.columns.values

        self.n_features = self.X.shape[1]

        self.logger.info('{} features after filtering'.format(self.n_features))

        self.logger.info('features: {} ...'.format(str(self.feature_names[:3])))

    @define_step(inputs=['X', 'positive_class', 'negative_class'],

        outputs=['X', 'y', 'n_samples', 'sample_ids', 'X_raw'])

    def select_samples(self):

        if (self.positive_class is not None) and (self.negative_class is not None):

            self.positive_class = self.positive_class.split(',')

            self.negative_class = self.negative_class.split(',')

        else:

            unique_classes = np.unique(self.sample_classes.values)

            if len(unique_classes) != 2:

                raise ValueError('expect 2 classes but {} classes found'.format(len(unique_classes)))

            self.positive_class, self.negative_class = unique_classes

        self.positive_class = np.atleast_1d(self.positive_class)

        self.negative_class = np.atleast_1d(self.negative_class)

        self.logger.info('positive class: {}, negative class: {}'.format(self.positive_class, self.negative_class))

        X_pos = self.X.loc[self.sample_classes[self.sample_classes.isin(self.positive_class)].index.values]

        X_neg = self.X.loc[self.sample_classes[self.sample_classes.isin(self.negative_class)].index.values]

        self.logger.info('number of positive samples: {}, negative samples: {}, class ratio: {}'.format(

            X_pos.shape[0], X_neg.shape[0], float(X_pos.shape[0])/X_neg.shape[0]))

        self.X = pd.concat([X_pos, X_neg], axis=0)

        self.y = np.zeros(self.X.shape[0], dtype=np.int32)

        self.y[X_pos.shape[0]:] = 1

        self.X_raw = self.X

        self.n_samples = self.X.shape

        self.sample_ids = self.X.index.values

    @define_step(inputs=['X', 'use_log', 'scaler'], outputs=['X'])

    def scale_features(self):

        if self.use_log:

            self.logger.info('apply log2 to feature matrix')

            self.X = np.log2(self.X + 1)

        if self.scaler == 'zscore':

            scaler = StandardScaler()

        elif self.scaler == 'robust':

            scaler = RobustScaler()

        elif self.scaler == 'min_max':

            scaler = MinMaxScaler()

        elif self.scaler == 'max_abs':

            scaler = MaxAbsScaler()

        self.logger.info('scale features using {}'.format(self.scaler))

        self.X = scaler.fit_transform(self.X)

    @define_step(inputs=['X', 'method', 'resample_method'], 

        outputs=['estimator', 'robust_estimator', 'compute_sample_weight'])

    def train_model(self):

        if np.any(np.isnan(self.X)):

            self.logger.info('nan values found in features')

        self.estimator = None

        grid_search = None

        self.logger.info('use {} to select features'.format(self.method))

        if self.method == 'r_test':

            self.estimator = TTestEstimator()

        elif self.method == 'logistic_regression':

            self.estimator = LogisticRegression()

            grid_search = {'C': [1e-5, 1e-4, 1e-3, 1e-2, 1e-1, 1, 1e2, 1e3, 1e4, 1e5]}

        elif self.method == 'random_forest':

            self.estimator = RandomForestClassifier()

            grid_search = {'max_depth': list(range(2, 10))}

        elif self.method == 'linear_svm':

            self.estimator = LinearSVC()

            grid_search = {'C': [1e-5, 1e-4, 1e-3, 1e-2, 1e-1, 1, 1e2, 1e3, 1e4, 1e5]}

        else:

            raise ValueError('unknown feature selection method: {}'.format(self.method))

        resampler_args = {}

        if self.resample_method == 'jackknife':

            resampler_args = {'max_runs': self.jackknife_max_runs,

                'remove': self.jackknife_remove

            }

        elif self.resample_method == 'bootstrap':

            resampler_args = {'max_runs': self.bootstrap_max_runs}

        self.robust_estimator = RobustEstimator(self.estimator, n_select=self.n_select, 

            grid_search=grid_search, resample_method='bootstrap',

            rfe=self.rfe, **resampler_args)

        sample_weight = None

        if self.compute_sample_weight:

            sample_weight = compute_sample_weight('balanced', self.y)

        self.robust_estimator.fit(self.X, self.y, sample_weight=sample_weight)

        self.estimator = self.robust_estimator.estimator_

    @define_step(inputs=['estimator', 'X', 'y'])

    def evaluate_model(self):

        self.logger.info('evaluate the model')

        features = pd.read_table(os.path.join(self.output_dir, 'features.txt'),

            header=None).iloc[:, 0].values

        feature_index = pd.Series(np.arange(self.n_features), 

            index=self.feature_names).loc[features]

        y_pred = self.estimator.predict(self.X[:, feature_index])

        y_score = self.estimator.predict_proba(self.X[:, feature_index])[:, 1]

        roc_auc = roc_auc_score(self.y, y_score)

        fpr, tpr, thresholds = roc_curve(self.y, y_score)

        sns.set_style('whitegrid')

        fig, ax = plt.subplots(figsize=(7, 7))

        ax.plot(fpr, tpr, label='AUC = {:.4f}'.format(roc_auc))

        ax.plot([0, 1], [0, 1], linestyle='dashed', color='gray', linewidth=0.8)

        ax.set_xlabel('False positive rate')

        ax.set_ylabel('True positive rate')

        ax.legend()

        plt.savefig(os.path.join(self.output_dir, 'roc_curve.pdf'))

        plt.close()

    @define_step(inputs=['estimator'])

    def save_model(self):

        self.logger.info('save model')

        with open(os.path.join(self.output_dir, 'model.pkl'), 'wb') as f:

            pickle.dump(self.estimator, f)

    @define_step(outputs=['estimator'])

    def load_model(self):

        self.logger.info('load model')

        with open(os.path.join(self.output_dir, 'model.pkl'), 'rb') as f:

            self.estimator = pickle.load(f)

    @define_step(inputs=['estimator'])

    def save_params(self):

        self.logger.info('save model parameters')

        with open(os.path.join(self.output_dir, 'params.json'), 'w') as f:

            json.dump(self.estimator.get_params(), f, indent=2)

    def save_matrix(self):

        self.logger.info('save matrix')

        df = pd.DataFrame(self.X, columns=self.feature_names, index=self.sample_ids)

        df.index.name = 'sample'

        df.to_csv(os.path.join(self.output_dir, 'matrix.txt'),

            sep='\t', index=True, header=True)

        data = pd.Series(self.y, index=self.sample_ids)

        data.to_csv(os.path.join(self.output_dir, 'labels.txt'),

            sep='\t', index=True, header=False)

    @define_step(inputs=['output_dir', 'n_features', 'feature_names', 'X', 'y'])

    def single_feature_metrics(self):

        self.logger.info('compute metrics for selected features independently')

        features = pd.read_table(os.path.join(self.output_dir, 'features.txt'),

            header=None).iloc[:, 0].values

        feature_index = pd.Series(np.arange(self.n_features), 

            index=self.feature_names).loc[features]

        metrics = {}

        scorers = {'roc_auc': roc_auc_score}

        for metric in ['roc_auc']:

            metrics[metric] = np.zeros(len(features))

            for i, i_feature in enumerate(feature_index):

                metrics[metric][i] = scorers[metric](self.y, self.X[:, i_feature])

        metrics = pd.DataFrame(metrics, index=features)

        metrics.index.name = 'feature'

        metrics.to_csv(os.path.join(self.output_dir, 'single_feature_metrics.txt'),

            sep='\t', index=True, header=True)

    @define_step(inputs=['estimator'])

    def plot_feature_importance(self):

        self.logger.info('plot feature importance')

        features = pd.read_table(os.path.join(self.output_dir, 'features.txt'),

            header=None).iloc[:, 0].values

        feature_importance = pd.read_table(os.path.join(self.output_dir, 'feature_importances.txt'),

            names=['feature', 'feature_importance'], header=None)

        feature_importance = feature_importance.iloc[np.argsort(-feature_importance['feature_importance'].values), :]

        print(feature_importance.head())

        sns.set_style('whitegrid')

        fig, ax = plt.subplots(figsize=(15, 20))

        sns.barplot('feature_importance', 'feature', color='gray',

            data=feature_importance, ax=ax)

        plt.subplots_adjust(left=0.3)

        plt.savefig(os.path.join(self.output_dir, 'feature_importances_refitted.pdf'))

        plt.close()

        feature_importance_matrix = pd.read_table(os.path.join(self.output_dir, 'feature_importance_matrix.txt'))

        feature_importance_matrix = feature_importance_matrix.loc[:, features]

        feature_importance_matrix = feature_importance_matrix.iloc[:, np.argsort(-feature_importance_matrix.mean(axis=0).values)]

        data = pd.melt(feature_importance_matrix, var_name='feature', value_name='feature_importance')

        sns.set_style('whitegrid')

        fig, ax = plt.subplots(figsize=(15, 25))

        sns.barplot('feature_importance', 'feature', color='gray', ci='sd',

            data=data, ax=ax, errwidth=1, capsize=0.2)

        plt.subplots_adjust(left=0.2)

        plt.savefig(os.path.join(self.output_dir, 'feature_importances_estimate.pdf'))

        plt.close()