#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon May 23 23:25:26 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
#####################################
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)
            }) 
  

mcc_score_fn = {'mcc': make_scorer(matthews_corrcoef)}
jacc_score_fn = {'jcc': make_scorer(jaccard_score)}
skf_cv = StratifiedKFold(n_splits = 10
                          #, shuffle = False, random_state= None)
                           , shuffle = True,**rs)

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

###############################################################################
def fsgs_rfecv(input_df
         , target
         , param_gridLd = [{'fs__min_features_to_select' : [1]}]
         , blind_test_df = pd.DataFrame()
         , blind_test_target = pd.Series(dtype = 'int64')
         , estimator = LogisticRegression(**rs) # placeholder
         , use_fs = False # uses estimator as the RFECV parameter for fs. Set to TRUE if you want to supply custom_fs as shown below
         , custom_fs = RFECV(DecisionTreeClassifier(**rs) , cv =  skf_cv, scoring = 'matthews_corrcoef')
         , cv_method =  skf_cv
         , var_type = ['numerical', 'categorical' , 'mixed']
         , resampling_type = 'none'
         , verbose = 3
         , random_state = 42
         , n_jobs = os.cpu_count()
         ):
    '''
    returns
    Dict containing results from FS and hyperparam tuning for a given estiamtor
    
    >>> ADD MORE <<<
    
    optimised/selected based on mcc
    
    '''
    rs = {'random_state': random_state}
    njobs = {'n_jobs': n_jobs}

    
    ###########################################################################
    #================================================
    # 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 = [('cat', OneHotEncoder(), categorical_ix)
              , ('num', MinMaxScaler(), numerical_ix)]
        
    col_transform = ColumnTransformer(transformers = t
                                        , remainder='passthrough')
    
    ###########################################################################
    #==================================================
    # Create var_type ~ column names
    # using one hot encoder with RFECV means 
    # the names internally are lost. Hence
    # fit col_transformeer to my input_df and get 
    # all the column names out and stored in a var
    # to allow the 'selected features' to be subsetted
    # from the numpy boolean array
    #=================================================
    col_transform.fit(input_df)
    col_transform.get_feature_names_out()
    
    var_type_colnames = col_transform.get_feature_names_out()
    var_type_colnames = pd.Index(var_type_colnames)
    
    if var_type == 'mixed':
        print('\nVariable type is:', var_type
              , '\nNo. of columns in input_df:', len(input_df.columns)
              , '\nNo. of columns post one hot encoder:', len(var_type_colnames))
    else:
        print('\nNo. of columns in input_df:', len(input_df.columns))
        
    #==================================
    # Build FS with supplied estimator
    #==================================
    if use_fs:
        fs = custom_fs
    else:
        fs = RFECV(estimator, cv = skf_cv, scoring = 'matthews_corrcoef')
    
    #==================================
    # Build basic param grid
    #==================================
    # param_gridD = [
    #     {'fs__min_features_to_select' : [1] 
    #       }]
        
    ############################################################################   
    # Create Pipeline object
    pipe = Pipeline([
        ('pre', col_transform),
        ('fs', fs),
        ('clf', estimator)])
    ############################################################################   
    # Define GridSearchCV
    gscv_fs = GridSearchCV(pipe
                           #, param_gridLd = param_gridD
                           , param_gridLd
                           , cv = cv_method
                           , scoring = scoring_fn
                           , refit = 'mcc'
                           , verbose = 3
                           , return_train_score = True
                           , **njobs)
    
    gscv_fs.fit(input_df, target)
    
    ###########################################################################
    # Get best param and scores out
    gscv_fs.best_params_
    gscv_fs.best_score_
    
    # Training best score corresponds to the max of the mean_test<score>
    train_bscore = round(gscv_fs.best_score_, 2); train_bscore
    print('\nTraining best score (MCC):', train_bscore)
    gscv_fs.cv_results_['mean_test_mcc']
    round(gscv_fs.cv_results_['mean_test_mcc'].max(),2)
    round(np.nanmax(gscv_fs.cv_results_['mean_test_mcc']),2)
    
    check_train_score = [round(gscv_fs.cv_results_['mean_test_mcc'].max(),2)
                        , round(np.nanmax(gscv_fs.cv_results_['mean_test_mcc']),2)]
    
    check_train_score = np.nanmax(check_train_score)
    
    # Training results
    gscv_tr_resD = gscv_fs.cv_results_
    mod_refit_param =  gscv_fs.refit
    
    # sanity check
    if train_bscore == check_train_score:
        print('\nVerified training score (MCC):', train_bscore )
    else:
        sys.exit('\nTraining score could not be internatlly verified. Please check training results dict')
        
    #-------------------------
    # Dict of CV results
    #-------------------------
    cv_allD = gscv_fs.cv_results_
    cvdf0   = pd.DataFrame(cv_allD)
    cvdf    = cvdf0.filter(regex='mean_test', axis = 1)
    cvdfT   = cvdf.T
    cvdfT.columns = ['cv_score']
    cvdfTr = cvdfT.loc[:,'cv_score'].round(decimals = 2) # round values
    cvD     = cvdfTr.to_dict()
    print('\n CV results dict generated for:', len(scoring_fn), 'scores'
          , '\nThese are:', scoring_fn.keys())
        
    #-------------------------
    # Blind test: REAL check!
    #-------------------------
    #tp = gscv_fs.predict(X_bts)
    tp = gscv_fs.predict(blind_test_df)

    print('\nMCC on Blind test:'     , round(matthews_corrcoef(blind_test_target, tp),2))
    print('\nAccuracy on Blind test:', round(accuracy_score(blind_test_target, tp),2))
    
    #=================
    # info extraction
    #=================
    # gives input vals??
    gscv_fs._check_n_features
    
    # gives gscv params used
    gscv_fs._get_param_names()
    
    # gives ??
    gscv_fs.best_estimator_
    gscv_fs.best_params_ # gives best estimator params as a dict 
    gscv_fs.best_estimator_._final_estimator # similar to above, doesn't contain max_iter
    gscv_fs.best_estimator_.named_steps['fs'].get_support()
    gscv_fs.best_estimator_.named_steps['fs'].ranking_ # array of ranks for the features
    
    gscv_fs.best_estimator_.named_steps['fs'].grid_scores_.mean()
    gscv_fs.best_estimator_.named_steps['fs'].grid_scores_.max()
    #gscv_fs.best_estimator_.named_steps['fs'].grid_scores_
    
    estimator_mask = gscv_fs.best_estimator_.named_steps['fs'].get_support()

    
    ############################################################################
    #============
    # FS results
    #============
    # Now get the features out
    
    #--------------
    # All features 
    #--------------
    all_features = gscv_fs.feature_names_in_
    n_all_features =  gscv_fs.n_features_in_
    #all_features = gsfit.feature_names_in_
    
    #--------------
    # Selected features by the classifier
    # Important to have var_type_colnames here
    #----------------
    #sel_features = X.columns[gscv_fs.best_estimator_.named_steps['fs'].get_support()] 3 only for numerical df    
    sel_features = var_type_colnames[gscv_fs.best_estimator_.named_steps['fs'].get_support()]
    n_sf = gscv_fs.best_estimator_.named_steps['fs'].n_features_
    
    #--------------
    # Get model name
    #--------------
    model_name     = gscv_fs.best_estimator_.named_steps['clf']
    b_model_params = gscv_fs.best_params_
    
    print('\n========================================'
          , '\nRunning model:'
          , '\nModel name:', model_name
          , '\n==============================================='
          , '\nRunning feature selection with RFECV for model'
          , '\nTotal no. of features in model:', len(all_features)
          , '\nThese are:\n',  all_features, '\n\n'
          , '\nNo of features for best model: ', n_sf
          , '\nThese are:', sel_features, '\n\n'
          , '\nBest Model hyperparams:', b_model_params
          )
    
    ###########################################################################
    ############################## OUTPUT #####################################
    ###########################################################################
    #=========================
    # Blind test: BTS results
    #=========================
    # Build the final results with all scores for a feature selected model
    #bts_predict = gscv_fs.predict(X_bts)
    bts_predict = gscv_fs.predict(blind_test_df)

    print('\nMCC on Blind test:'     , round(matthews_corrcoef(blind_test_target, bts_predict),2))
    print('\nAccuracy on Blind test:', round(accuracy_score(blind_test_target, bts_predict),2))
    bts_mcc_score = round(matthews_corrcoef(blind_test_target, bts_predict),2)
    
    # Diff b/w train and bts test scores
    train_test_diff = train_bscore - bts_mcc_score
    print('\nDiff b/w train and blind test score (MCC):', train_test_diff)
    
    lr_btsD ={}
    #lr_btsD['bts_mcc']       = bts_mcc_score
    lr_btsD['bts_fscore']    = round(f1_score(blind_test_target, bts_predict),2)
    lr_btsD['bts_precision'] = round(precision_score(blind_test_target, bts_predict),2)
    lr_btsD['bts_recall']    = round(recall_score(blind_test_target, bts_predict),2)
    lr_btsD['bts_accuracy']  = round(accuracy_score(blind_test_target, bts_predict),2)
    lr_btsD['bts_roc_auc']   = round(roc_auc_score(blind_test_target, bts_predict),2)
    lr_btsD['bts_jcc']       = round(jaccard_score(blind_test_target, bts_predict),2)
    lr_btsD
    
    #===========================
    # Add FS related model info
    #===========================
    model_namef = str(model_name)
    # FIXME: doesn't tell you which it has chosen
    fs_methodf = str(gscv_fs.best_estimator_.named_steps['fs'])
    all_featuresL = list(all_features)
    fs_res_arrayf = str(list( gscv_fs.best_estimator_.named_steps['fs'].get_support()))
    fs_res_array_rankf = str(list( gscv_fs.best_estimator_.named_steps['fs'].ranking_))
    sel_featuresf = list(sel_features)
    n_sf = int(n_sf)
    
    output_modelD = {'model_name': model_namef
                     , 'model_refit_param': mod_refit_param
                     , 'Best_model_params': b_model_params 
                     , 'n_all_features': n_all_features
                     , 'fs_method': fs_methodf
                     , 'fs_res_array': fs_res_arrayf
                     , 'fs_res_array_rank': fs_res_array_rankf
                     , 'all_feature_names': all_featuresL
                     , 'n_sel_features': n_sf
                     , 'sel_features_names': sel_featuresf}
    #output_modelD
    
    #========================================
    # Update output_modelD with bts_results
    #========================================
    output_modelD.update(lr_btsD)
    output_modelD
    
    output_modelD['train_score (MCC)'] = train_bscore
    output_modelD['bts_mcc'] = bts_mcc_score
    output_modelD['train_bts_diff'] = round(train_test_diff,2)
    output_modelD['resampling'] = resampling_type

    print(output_modelD)
    
    nlen = len(output_modelD)
    
    #========================================
    # Update output_modelD with cv_results
    #========================================
    output_modelD.update(cvD)
    
    if (len(output_modelD) == nlen + len(cvD)):
        print('\nFS run complete for model:', estimator
              , '\nFS using:', fs
              , '\nOutput dict size:', len(output_modelD))
        return(output_modelD)
    else:
        sys.exit('\nFAIL:numbers mismatch output dict length not as expected. Please check')
    
