#!/usr/bin/env python
# coding: utf-8

# # scRFE

# In[ ]:





# In[84]:


# MENTION ONE VS ALL CLASSIFICATION in description


# In[117]:


# Imports
import numpy as np
import pandas as pd
import scanpy as sc
import random
from anndata import read_h5ad
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.feature_selection import SelectFromModel
from sklearn.metrics import accuracy_score
from sklearn.feature_selection import RFE
from sklearn.feature_selection import RFECV
import seaborn as sns
import matplotlib.pyplot as plt


# In[119]:


def filterNormalize (dataMatrix, classOfInterest):
    np.random.seed(644685)
    sc.logging.print_versions()
    sc.settings.verbosity = 3
    sc.logging.print_versions()
    tiss = dataMatrix
    tiss.obs['n_counts'] = tiss.X.sum(axis=1).A1
    sc.pp.filter_cells(tiss, min_genes=250)
    sc.pp.filter_genes(tiss, min_cells=3)
    tiss = tiss[tiss.obs['n_counts'] > 1500, :]
    sc.pp.normalize_per_cell(tiss, counts_per_cell_after=1e5)
    sc.pp.log1p(tiss)
    tiss.raw = tiss
    tiss = tiss[tiss.obs[classOfInterest]!='nan']
    return tiss


# In[120]:


# goal: get labels on a per class basis that will go into randomForest function for y
def getLabels (dataMatrix, classOfInterest):
    """
    Gets labels on a per class basis that will inputted to the randomForest function

    Parameters
    ----------
    dataMatrix : anndata object
        The data file of interest
    classOfInterest : str
        The class you will split the data by in the set of dataMatrix.obs

    Returns
    -------
    labelsDict : dict
        Dictionary with labels for each class
    """
    dataMatrix = filterNormalize (dataMatrix, classOfInterest)
    labelsDict = {}
    for label in np.unique(dataMatrix.obs[classOfInterest]):
        lists = []
        for obs in dataMatrix.obs[classOfInterest]:
            if obs == label:
                lists.append('A')
            else:
                lists.append('B')
        labelsDict[label] = lists #this is usually in line w if and else
    return labelsDict


# In[121]:


def makeOneForest (dataMatrix, classOfInterest, labelOfInterest, nEstimators = 5000,
                   randomState = 0,  nJobs = -1, oobScore = True, Step = 0.2, Cv = 5):
    """
    Builds and runs a random forest for one label in a class of interest

    Parameters
    ----------
    dataMatrix : anndata object
        The data file of interest
    classOfInterest : str
        The class you will split the data by in the set of dataMatrix.obs
    labelOfInterest : str
        The specific label within the class that the random forezt will run a
        "one vs all" classification on
    nEstimators : int
        The number of trees in the forest
    randomState : int
        Controls random number being used
    nJobs : int
        The number of jobs to run in parallel
    oobScore : bool
        Whether to use out-of-bag samples to estimate the generalization accuracy
    Step : float
        Corresponds to percentage of features to remove at each iteration
    Cv : int
        Determines the cross-validation splitting strategy

    Returns
    -------
    feature_selected : list
        list of top features from random forest
    selector.estimator_.feature_importances_ : list
        list of top ginis corresponding to to features

    """
    dataMatrix = filterNormalize (dataMatrix, classOfInterest)

    # print('makeOneForest' + labelOfInterest)
    labelsDict = getLabels(dataMatrix, classOfInterest)

    feat_labels = dataMatrix.var_names #this is equivalent of the genes
    X = dataMatrix.X
    y = labelsDict[labelOfInterest]
    # print('Y')
    clf = RandomForestClassifier(n_estimators = nEstimators, random_state = randomState,
                                 n_jobs = nJobs, oob_score = oobScore)
    selector = RFECV(clf, step = Step, cv = Cv)

    # print('training...')
    clf.fit(X, y)
    selector.fit(X, y)
    feature_selected = feat_labels[selector.support_]

    return feature_selected, selector.estimator_.feature_importances_


# In[122]:


def resultWrite (classOfInterest, results_df, labelOfInterest,
                feature_selected, feature_importance):
    # print ('result writing')
    # print(results_df)

    column_headings = []
    column_headings.append(labelOfInterest)
    column_headings.append(labelOfInterest + '_gini')
    resaux = pd.DataFrame(columns = column_headings)
    resaux[labelOfInterest] = feature_selected
    resaux[labelOfInterest + '_gini'] = feature_importance
    resaux = resaux.sort_values(by = [labelOfInterest + '_gini'], ascending = False)
    resaux.reset_index(drop = True, inplace = True)

    results_df = pd.concat([results_df, resaux], axis=1)
    return results_df


# In[123]:


def scRFE(dataMatrix, classOfInterest, nEstimators = 5000, randomState = 0,
                  nJobs = -1, oobScore = True, Step = 0.2, Cv = 5):
    """
    Builds and runs a random forest with one vs all classification for each label
    for one class of interest

    Parameters
    ----------
    dataMatrix : anndata object
        The data file of interest
    classOfInterest : str
        The class you will split the data by in the set of dataMatrix.obs
    labelOfInterest : str
        The specific label within the class that the random forezt will run a
        "one vs all" classification on
    nEstimators : int
        The number of trees in the forest
    randomState : int
        Controls random number being used
    nJobs : int
        The number of jobs to run in parallel
    oobScore : bool
        Whether to use out-of-bag samples to estimate the generalization accuracy
    Step : float
        Corresponds to percentage of features to remove at each iteration
    Cv : int
        Determines the cross-validation splitting strategy

    Returns
    -------
    results_df : pd.DataFrame
        Dataframe with results for each label in the class, formatted as
        "label" for one column, then "label + gini" for the corresponding column

    """

    dataMatrix = filterNormalize (dataMatrix, classOfInterest)
    results_df = pd.DataFrame()
    for labelOfInterest in np.unique(dataMatrix.obs[classOfInterest]): #for timeliness
        # print( 'scRFE' + labelOfInterest)

        feature_selected, feature_importance = makeOneForest(dataMatrix,
                                                             classOfInterest,
                          labelOfInterest = labelOfInterest)

        results_df = resultWrite (classOfInterest, results_df,
                            labelOfInterest = labelOfInterest,
                    feature_selected = feature_selected,
                    feature_importance = feature_importance)
        # print(results_df.shape)
    return results_df
