#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Mar  4 15:25:33 2022

@author: tanu
"""
#%%
import os, sys
import pandas as pd
import numpy as np
import pprint as pp
from copy import deepcopy
from sklearn import linear_model
from sklearn import datasets
from collections import Counter

from sklearn.linear_model import LogisticRegression, LogisticRegressionCV
from sklearn.linear_model import RidgeClassifier, RidgeClassifierCV, SGDClassifier, PassiveAggressiveClassifier

from sklearn.naive_bayes import BernoulliNB
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier, ExtraTreeClassifier
from sklearn.ensemble import RandomForestClassifier, ExtraTreesClassifier, AdaBoostClassifier, GradientBoostingClassifier, BaggingClassifier
from sklearn.naive_bayes import GaussianNB
from sklearn.gaussian_process import GaussianProcessClassifier, kernels
from sklearn.gaussian_process.kernels import RBF, DotProduct, Matern, RationalQuadratic, WhiteKernel

from sklearn.discriminant_analysis import LinearDiscriminantAnalysis, QuadraticDiscriminantAnalysis
from sklearn.neural_network import MLPClassifier

from sklearn.svm import SVC
from xgboost import XGBClassifier
from sklearn.naive_bayes import MultinomialNB
from sklearn.preprocessing import StandardScaler, MinMaxScaler, OneHotEncoder

from sklearn.compose import ColumnTransformer
from sklearn.compose import make_column_transformer

from sklearn.metrics import make_scorer, confusion_matrix, accuracy_score, balanced_accuracy_score, precision_score, average_precision_score, recall_score
from sklearn.metrics import roc_auc_score, roc_curve, f1_score, matthews_corrcoef, jaccard_score, classification_report

# added
from sklearn.model_selection import train_test_split, cross_validate, cross_val_score, LeaveOneOut, KFold, RepeatedKFold, cross_val_predict

from sklearn.model_selection import train_test_split, cross_validate, cross_val_score
from sklearn.model_selection import StratifiedKFold,RepeatedStratifiedKFold, RepeatedKFold

from sklearn.pipeline import Pipeline, make_pipeline

from sklearn.feature_selection import RFE, RFECV

import itertools
import seaborn as sns
import matplotlib.pyplot as plt

from statistics import mean, stdev, median, mode

from imblearn.over_sampling import RandomOverSampler
from imblearn.under_sampling import RandomUnderSampler
from imblearn.over_sampling import SMOTE
from sklearn.datasets import make_classification
from imblearn.combine import SMOTEENN
from imblearn.combine import SMOTETomek

from imblearn.over_sampling import SMOTENC
from imblearn.under_sampling import EditedNearestNeighbours
from imblearn.under_sampling import RepeatedEditedNearestNeighbours

from sklearn.model_selection import GridSearchCV
from sklearn.base import BaseEstimator
from sklearn.impute import KNNImputer as KNN
import json
import argparse
import re
import itertools
from sklearn.model_selection import LeaveOneGroupOut
from sklearn.decomposition import PCA

#%% GLOBALS
#rs = {'random_state': 42}
#njobs = {'n_jobs': os.cpu_count() } # the number of jobs should equal the number of CPU cores

scoring_fn =  ({ 'mcc'        : make_scorer(matthews_corrcoef)
                , 'fscore'    : make_scorer(f1_score)
                , 'precision' : make_scorer(precision_score)
                , 'recall'    : make_scorer(recall_score)
                , 'accuracy'  : make_scorer(accuracy_score)
                , 'roc_auc'   : make_scorer(roc_auc_score)
                , 'jcc'       : make_scorer(jaccard_score)
            }) 
 
# for sel_cv
#skf_cv = StratifiedKFold(n_splits = 10
#                          #, shuffle = False, random_state= None)
#                           , shuffle = True, **rs)

#rskf_cv = RepeatedStratifiedKFold(n_splits = 10
#                                  , n_repeats = 3
#                                  , **rs)

mcc_score_fn = {'mcc': make_scorer(matthews_corrcoef)}
jacc_score_fn = {'jcc': make_scorer(jaccard_score)}

###############################################################################
score_type_ordermapD = { 'mcc'      : 1
                   , 'fscore'       : 2
                   , 'jcc'          : 3
                   , 'precision'    : 4
                   , 'recall'       : 5      
                   , 'accuracy'     : 6  
                   , 'roc_auc'      : 7
                   , 'TN'           : 8
                   , 'FP'           : 9
                   , 'FN'           : 10
                   , 'TP'           : 11  
                   , 'trainingY_neg': 12  
                   , 'trainingY_pos': 13    
                   , 'blindY_neg'   : 14
                   , 'blindY_pos'   : 15
                   , 'fit_time'     : 16
                   , 'score_time'   : 17
                   }

scoreCV_mapD = {'test_mcc'         : 'MCC'
                , 'test_fscore'    : 'F1'
                , 'test_precision' : 'Precision'
                , 'test_recall'    : 'Recall'
                , 'test_accuracy'  : 'Accuracy'
                , 'test_roc_auc'   : 'ROC_AUC'
                , 'test_jcc'       : 'JCC'
                }

scoreBT_mapD = {'bts_mcc'          : 'MCC'
                , 'bts_fscore'     : 'F1'
                , 'bts_precision'  : 'Precision'
                , 'bts_recall'     : 'Recall'
                , 'bts_accuracy'   : 'Accuracy'
                , 'bts_roc_auc'    : 'ROC_AUC'
                , 'bts_jcc'        : 'JCC'
               }

#gene_group = 'gene_name'
#%%############################################################################
############################
# MultModelsCl()
# Run Multiple Classifiers
############################
# Multiple Classification - Model Pipeline
def MultModelsCl_logo_skf(input_df
                       , target
                       , sel_cv
                       , tts_split_type
                       , resampling_type
                       #, group = None
                       
                       , add_cm = True # adds confusion matrix based on cross_val_predict
                       , add_yn = True  # adds target var class numbers
                       , var_type = ['numerical', 'categorical','mixed']
                       , scale_numeric = ['min_max', 'std', 'min_max_neg', 'none'] 

                       , run_blind_test = True
                       , blind_test_df = pd.DataFrame()
                       , blind_test_target = pd.Series(dtype = int)
                       , return_formatted_output = True

                       , random_state = 42
                       , n_jobs = os.cpu_count() # the number of jobs should equal the number of CPU cores
                       ):

    '''
    @ param input_df: input features 
    @ type: df with input features WITHOUT the target variable
    
    @param target: target (or output) feature
    @type: df or np.array or Series
    
    @param skv_cv: stratifiedK fold int or object to allow shuffle and random state to pass
    @type: int or StratifiedKfold()
    
    @var_type: numerical, categorical and mixed to determine what col_transform to apply (MinMaxScalar and/or one-hot encoder)
    @type: list

    returns
    Dict containing multiple classification scores for each model and mean of each Stratified Kfold including training
    '''
    
#%% Func globals        
    rs = {'random_state': random_state}
    njobs = {'n_jobs': n_jobs}


    # select CV type:           
    # if group == None:
    #     sel_cv = skf_cv
    # else: 
    #     sel_cv = logo
    #======================================================
    # Determine categorical and numerical features
    #======================================================
    numerical_ix = input_df.select_dtypes(include=['int64', 'float64']).columns
    numerical_ix
    categorical_ix = input_df.select_dtypes(include=['object', 'bool']).columns
    categorical_ix    

    #======================================================
    # Determine preprocessing steps ~ var_type
    #======================================================
   
    # if var_type == 'numerical':
    #     t = [('num', MinMaxScaler(), numerical_ix)]

    # if var_type == 'categorical':
    #     t = [('cat', OneHotEncoder(), categorical_ix)]
    
    # # if var_type == 'mixed':
    # #     t = [('num', MinMaxScaler(), numerical_ix)
    # #         , ('cat', OneHotEncoder(), categorical_ix) ]

    # col_transform = ColumnTransformer(transformers = t
    #                                    , remainder='passthrough')
    
    if type(var_type) == list: 
        var_type = str(var_type[0])
    else:
        var_type = var_type
    
    if var_type in ['numerical','mixed']:
        if scale_numeric == ['none']:
            t = [('cat', OneHotEncoder(), categorical_ix)]
        if scale_numeric != ['none']:
            if scale_numeric == ['min_max']:
                scaler = MinMaxScaler()
            if scale_numeric == ['min_max_neg']:
                scaler = MinMaxScaler(feature_range=(-1, 1))
            if scale_numeric == ['std']:
                scaler = StandardScaler()
            
            t = [('num', scaler, numerical_ix)
              , ('cat', OneHotEncoder(), categorical_ix)]
        
                
    if var_type == 'categorical':
        t = [('cat', OneHotEncoder(), categorical_ix)]

        
    col_transform = ColumnTransformer(transformers = t
                                       , remainder='passthrough')
    
    
    #======================================================
    # Specify multiple Classification Models  
    #======================================================
    models = [('AdaBoost Classifier'            , AdaBoostClassifier(**rs) )
               , ('Bagging Classifier'          , BaggingClassifier(**rs, **njobs, bootstrap = True, oob_score = True, verbose = 3, n_estimators = 100) )
               , ('Decision Tree'             , DecisionTreeClassifier(**rs) ) 
               , ('Extra Tree'                , ExtraTreeClassifier(**rs) )
               , ('Extra Trees'               , ExtraTreesClassifier(**rs) ) 
               , ('Gradient Boosting'         , GradientBoostingClassifier(**rs) )
               , ('Gaussian NB'               , GaussianNB() )
               , ('Gaussian Process'          , GaussianProcessClassifier(**rs) )
               , ('K-Nearest Neighbors'       , KNeighborsClassifier() ) 
               , ('LDA'                       , LinearDiscriminantAnalysis() )
               , ('Logistic Regression'       , LogisticRegression(**rs) )
               , ('Logistic RegressionCV'     , LogisticRegressionCV(cv = 3, **rs))
               , ('MLP'                       , MLPClassifier(max_iter = 500, **rs) ) 
               , ('Multinomial'               , MultinomialNB() )
               , ('Naive Bayes'               , BernoulliNB() )
               , ('Passive Aggresive'         , PassiveAggressiveClassifier(**rs, **njobs) )
               , ('QDA'                       , QuadraticDiscriminantAnalysis() )
               , ('Random Forest'             , RandomForestClassifier(**rs, n_estimators = 1000, **njobs ) ) 
               , ('Random Forest2'            , RandomForestClassifier(min_samples_leaf = 5
                                                                       , n_estimators     = 1000
                                                                       , bootstrap        = True
                                                                       , oob_score        = True
                                                                       , **njobs
                                                                       , **rs
                                                                       , max_features     = 'auto') ) 
                , ('Ridge Classifier'          , RidgeClassifier(**rs)  )
                , ('Ridge ClassifierCV'        , RidgeClassifierCV(cv = 3) )          
                , ('SVC'                       , SVC(**rs) ) 
                , ('Stochastic GDescent'       , SGDClassifier(**rs, **njobs) )
                , ('XGBoost'                   , XGBClassifier(**rs, verbosity = 0, use_label_encoder = False, **njobs) )
             ]
                
    mm_skf_scoresD = {}
    
    print('\n==============================================================\n'
          , '\nRunning several classification models (n):', len(models)
          ,'\nList of models:')
    for m in models:
        print(m)
    print('\n================================================================\n')
    
    index = 1
    for model_name, model_fn in models:
        print('\nRunning classifier:', index
              , '\nModel_name:'               , model_name
              , '\nModel func:'               , model_fn)
        index = index+1
        
        model_pipeline = Pipeline([
            ('prep'     , col_transform)
            , ('model'  , model_fn)])
        
        # model_pipeline = Pipeline([
        #     ('prep'     , col_transform)
        #     ,  ('pca'   , PCA(n_components = 2))
        #     , ('model'  , model_fn)])
            
            
        print('\nRunning model pipeline:', model_pipeline)
        cv_modD = cross_validate(model_pipeline
                              , input_df
                              , target
                              , cv = sel_cv
                              #, groups = group
                              , scoring = scoring_fn
                              , return_train_score = True)
        #==============================
        # Extract mean values for CV 
        #==============================
        mm_skf_scoresD[model_name] = {}
        
        for key, value in cv_modD.items():
            print('\nkey:', key, '\nvalue:', value)
            print('\nmean value:', np.mean(value))
            mm_skf_scoresD[model_name][key] = round(np.mean(value),2)
            
        # ADD more info: meta data related to input df
        mm_skf_scoresD[model_name]['resampling']        = resampling_type
        mm_skf_scoresD[model_name]['n_training_size']   = len(input_df)
        mm_skf_scoresD[model_name]['n_trainingY_ratio'] = round(Counter(target)[0]/Counter(target)[1], 2)
        mm_skf_scoresD[model_name]['n_features']        = len(input_df.columns)
        mm_skf_scoresD[model_name]['tts_split']         = tts_split_type
        
        #######################################################################
        #======================================================
        # Option: Add confusion matrix from cross_val_predict
        # Understand and USE with caution
        #======================================================
        if add_cm:  
           cmD = {}

            # Calculate cm         
           y_pred   = cross_val_predict(model_pipeline
                                        , input_df
                                        , target
                                        , cv = sel_cv
                                        #, groups = group
                                        , **njobs)
            #_tn, _fp, _fn, _tp = confusion_matrix(y_pred, y).ravel() # internally
           tn, fp, fn, tp = confusion_matrix(y_pred, target).ravel()
    
            # Build cm dict
           cmD = {'TN'  : tn
                   , 'FP': fp
                   , 'FN': fn
                   , 'TP': tp}
            
            # Update cv dict cmD           
           mm_skf_scoresD[model_name].update(cmD)
            
        #=============================================
        # Option: Add targety numbers for data
        #=============================================
        if add_yn:   
            tnD = {}
                       
            # Build tn numbers dict
            tnD = {'n_trainingY_neg'    : Counter(target)[0]
                   , 'n_trainingY_pos'  : Counter(target)[1] }
            
            # Update cv dict with cmD and tnD
            mm_skf_scoresD[model_name].update(tnD)

#%%
        #=========================
        # Option: Blind test (bts)
        #=========================
        if run_blind_test:
           btD = {}
           
           # Build bts numbers dict
           btD = {'n_blindY_neg'    : Counter(blind_test_target)[0]
                  , 'n_blindY_pos'  : Counter(blind_test_target)[1]
                  , 'n_testY_ratio' : round(Counter(blind_test_target)[0]/Counter(blind_test_target)[1], 2)
                  , 'n_test_size'   : len(blind_test_df) }
           
           # Update cmD+tnD dicts with btD
           mm_skf_scoresD[model_name].update(btD)
                    
           #--------------------------------------------------------
           # Build the final results with all scores for the model
           #--------------------------------------------------------
           #bts_predict = gscv_fs.predict(blind_test_df)
           model_pipeline.fit(input_df, target)
           bts_predict = model_pipeline.predict(blind_test_df)
           
           bts_mcc_score = round(matthews_corrcoef(blind_test_target, bts_predict),2)
           print('\nMCC on Blind test:'      , bts_mcc_score)
           #print('\nAccuracy on Blind test:', round(accuracy_score(blind_test_target, bts_predict),2))
           print('\nMCC on Training:'        , mm_skf_scoresD[model_name]['test_mcc'] )

           mm_skf_scoresD[model_name]['bts_mcc']       = bts_mcc_score
           mm_skf_scoresD[model_name]['bts_fscore']    = round(f1_score(blind_test_target, bts_predict),2)
           mm_skf_scoresD[model_name]['bts_precision'] = round(precision_score(blind_test_target, bts_predict),2)
           mm_skf_scoresD[model_name]['bts_recall']    = round(recall_score(blind_test_target, bts_predict),2)
           mm_skf_scoresD[model_name]['bts_accuracy']  = round(accuracy_score(blind_test_target, bts_predict),2)
           mm_skf_scoresD[model_name]['bts_roc_auc']   = round(roc_auc_score(blind_test_target, bts_predict),2)
           mm_skf_scoresD[model_name]['bts_jcc']       = round(jaccard_score(blind_test_target, bts_predict),2)
           #mm_skf_scoresD[model_name]['diff_mcc']      = train_test_diff_MCC


    #return(mm_skf_scoresD)
    #============================
    # Process the dict to have WF
    #============================
    if return_formatted_output:
        CV_BT_metaDF = ProcessMultModelsCl(mm_skf_scoresD)
        return(CV_BT_metaDF)
    else:
        return(mm_skf_scoresD)

#%% Process output function ###################################################
############################
# ProcessMultModelsCl() 
############################
#Processes the dict from above if use_formatted_output = True 

def ProcessMultModelsCl(inputD = {}, blind_test_data = True):
    
    scoresDF = pd.DataFrame(inputD)
    
    #------------------------
    #  Extracting split_name
    #-----------------------
    tts_split_nameL = []
    for k,v in inputD.items():
        tts_split_nameL = tts_split_nameL + [v['tts_split']]
    
    if len(set(tts_split_nameL)) == 1:
        tts_split_name = str(list(set(tts_split_nameL))[0])
        print('\nExtracting tts_split_name:', tts_split_name)
    
    #----------------------
    #  WF: CV results
    #----------------------
    scoresDFT = scoresDF.T
    
    scoresDF_CV = scoresDFT.filter(regex='^test_.*$', axis = 1); scoresDF_CV.columns
    # map colnames for consistency to allow concatenting
    scoresDF_CV.columns = scoresDF_CV.columns.map(scoreCV_mapD); scoresDF_CV.columns
    scoresDF_CV['source_data'] = 'CV'
    
    #----------------------
    #  WF: Meta data 
    #----------------------
    metaDF = scoresDFT.filter(regex='^(?!test_.*$|bts_.*$|train_.*$).*'); metaDF.columns
    
    print('\nTotal cols in each df:'
          , '\nCV df:', len(scoresDF_CV.columns)
          , '\nmetaDF:', len(metaDF.columns))
    
    #-------------------------------------
    # Combine WF: CV + Metadata
    #-------------------------------------

    combDF = pd.merge(scoresDF_CV, metaDF, left_index = True, right_index = True)  
    print('\nAdding column: Model_name')
    combDF['Model_name'] = combDF.index

    #----------------------
    #  WF: BTS results
    #----------------------      
    if blind_test_data:
          
        scoresDF_BT = scoresDFT.filter(regex='^bts_.*$', axis = 1); scoresDF_BT.columns
        # map colnames for consistency to allow concatenting
        scoresDF_BT.columns = scoresDF_BT.columns.map(scoreBT_mapD); scoresDF_BT.columns
        scoresDF_BT['source_data'] = 'BT'
    
    
        print('\nTotal cols in bts df:'
              , '\nBT_df:', len(scoresDF_BT.columns))
        
        if  len(scoresDF_CV.columns) == len(scoresDF_BT.columns):
            print('\nFirst proceeding to rowbind CV and BT dfs:')
            expected_ncols_out = len(scoresDF_BT.columns) + len(metaDF.columns)
            print('\nFinal output should have:', expected_ncols_out, 'columns' )

        #-----------------
        # Combine WF
        #-----------------
        dfs_combine_wf = [scoresDF_CV, scoresDF_BT]
    
        print('\nCombinig', len(dfs_combine_wf), 'using pd.concat by row ~ rowbind'
              , '\nChecking Dims of df to combine:'
              , '\nDim of CV:', scoresDF_CV.shape
              , '\nDim of BT:', scoresDF_BT.shape)
        #print(scoresDF_CV)
        #print(scoresDF_BT)
    
        dfs_nrows_wf = []
        for df in dfs_combine_wf:
            dfs_nrows_wf = dfs_nrows_wf + [len(df)]
        dfs_nrows_wf = max(dfs_nrows_wf)
            
        dfs_ncols_wf = []
        for df in dfs_combine_wf:
            dfs_ncols_wf = dfs_ncols_wf + [len(df.columns)]
        dfs_ncols_wf = max(dfs_ncols_wf)
        print(dfs_ncols_wf)
        
        expected_nrows_wf = len(dfs_combine_wf) * dfs_nrows_wf
        expected_ncols_wf = dfs_ncols_wf
        
        common_cols_wf = list(set.intersection(*(set(df.columns) for df in dfs_combine_wf)))
        print('\nNumber of Common columns:', dfs_ncols_wf
              , '\nThese are:', common_cols_wf)
    
        if len(common_cols_wf) == dfs_ncols_wf :
            combined_baseline_wf = pd.concat([df[common_cols_wf] for df in dfs_combine_wf], ignore_index=False)
            print('\nConcatenating dfs with different resampling methods [WF]:'
                  , '\nSplit type:', tts_split_name
                  , '\nNo. of dfs combining:', len(dfs_combine_wf))
            #print('\n================================================^^^^^^^^^^^^')
            if len(combined_baseline_wf) == expected_nrows_wf  and len(combined_baseline_wf.columns) == expected_ncols_wf:
                #print('\n================================================^^^^^^^^^^^^')
    
                print('\nPASS:', len(dfs_combine_wf), 'dfs successfully combined'
                      , '\nnrows in combined_df_wf:', len(combined_baseline_wf)
                      , '\nncols in combined_df_wf:', len(combined_baseline_wf.columns))
            else:
                print('\nFAIL: concatenating failed'
                      , '\nExpected nrows:', expected_nrows_wf
                      , '\nGot:', len(combined_baseline_wf)
                      , '\nExpected ncols:', expected_ncols_wf
                      , '\nGot:', len(combined_baseline_wf.columns))
                sys.exit('\nFIRST IF FAILS')
            ##
            c1L = list(set(combined_baseline_wf.index))
            c2L = list(metaDF.index)

            #if set(c1L) == set(c2L):
            if set(c1L) == set(c2L) and all(x in c2L for x in c1L) and all(x in c1L for x in c2L):
                print('\nPASS: proceeding to merge metadata with CV and BT dfs')
                combDF = pd.merge(combined_baseline_wf, metaDF, left_index = True, right_index = True)
                print('\nAdding column: Model_name')
                combDF['Model_name'] = combDF.index
                
            else:
                sys.exit('\nFAIL: Could not merge metadata with CV and BT dfs')

        else:
        #    print('\nConcatenting dfs not possible [WF],check numbers ')    
            print('\nOnly combining CV and metadata')    

        #-------------------------------------
        # Combine WF+Metadata: Final output
        #-------------------------------------
       
        # if len(combDF.columns) == expected_ncols_out:
        #     print('\nPASS: Combined df has expected ncols')
        # else:
        #     sys.exit('\nFAIL: Length mismatch for combined_df')
            
        # print('\nAdding column: Model_name')
        # combDF['Model_name'] = combDF.index
        
        print('\n========================================================='
              , '\nSUCCESS: Ran multiple classifiers'
              , '\n=======================================================')

    #resampling_methods_wf = combined_baseline_wf[['resampling']]
    #resampling_methods_wf = resampling_methods_wf.drop_duplicates()
              #, '\n', resampling_methods_wf)

    return combDF

###############################################################################
