#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Jun 29 20:29:36 2022

@author: tanu
"""
import sys, os
import pandas as pd
import numpy as np
import re
#import prettyprint as pp 
###############################################################################
homedir = os.path.expanduser("~")
sys.path.append(homedir + '/git/LSHTM_analysis/scripts/ml/ml_functions')
sys.path
###############################################################################
outdir = homedir + '/git/LSHTM_ML/output/fs/'

#====================
# Import ML functions
#====================

from MultClfs import *
from GetMLData import *
from SplitTTS import *
from FS import *
# param dict for getmldata()
combined_model_paramD = {'data_combined_model'   : False
        , 'use_or'                   : False
        , 'omit_all_genomic_features': False
        , 'write_maskfile'           : False
        , 'write_outfile'            : False }
###############################################################################
#ml_genes = ["pncA", "embB", "katG", "rpoB", "gid"]
# outdir = homedir + '/git/Data/ml_combined/fs/'

ml_gene_drugD = {
        'pncA' : 'pyrazinamide', # NOTE: may need re-run for 80_20 and sl
        #'embB' : 'ethambutol',
        #'katG' : 'isoniazid', #NOTE: RF only for all split-types actual
        #'rpoB' : 'rifampicin',
        #'gid'  : 'streptomycin' # NOTE: for gid, run 'actual' on 80/20 and sl only
        }
gene_dataD={}

#split_types = ['70_30', '80_20', 'sl']
#split_data_types = ['actual', 'complete']

split_types = ['70_30']
#split_data_types = ['actual', 'complete']
split_data_types = ['actual']

fs_models = [
        #('Ridge Classifier'          , RidgeClassifier(**rs)  ),
        #('Ridge ClassifierCV'        , RidgeClassifierCV(cv = 3) ),
        #('Logistic Regression'       , LogisticRegression(**rs, **njobs) ),
        #('AdaBoost Classifier'       , AdaBoostClassifier(**rs) ),
        #('Gradient Boosting'        , GradientBoostingClassifier(**rs) ),
        #('Stochastic GDescent'      , SGDClassifier(**rs, **njobs) ),
        #('Decision Tree'            , DecisionTreeClassifier(**rs) ),
        #('Extra Trees'              , ExtraTreesClassifier(**rs, **njobs) ),
        #('Extra Tree'               , ExtraTreeClassifier(**rs) ),
        #('LDA'                      , LinearDiscriminantAnalysis() ),
        #('Logistic RegressionCV'    , LogisticRegressionCV(cv = 3, **rs, **njobs) ),
        #('Passive Aggresive'        , PassiveAggressiveClassifier(**rs, **njobs) )
        #('Random Forest'            , RandomForestClassifier(n_estimators = 1000, verbose=3, **rs, **njobs ) )
        ('XGBoost'                   , XGBClassifier(verbosity=3, use_label_encoder=False, **rs, **njobs) )
        ]

for gene, drug in ml_gene_drugD.items():
    #print ('\nGene:', gene
    #       , '\nDrug:', drug)
    gene_low = gene.lower()
    gene_dataD[gene_low] = getmldata(gene, drug
            , data_combined_model = False # this means it doesn't include 'gene_name' as a feauture as a single gene-target shouldn't have it.
            , use_or = False
            , omit_all_genomic_features = False
            , write_maskfile = False
            , write_outfile = False)

    for split_type in split_types:
        for data_type in split_data_types:
            # unused per-split outfile
            #out_filename = outdir + gene.lower() + '_'+split_type+'_' + data_type + '.json'
            tempD=split_tts(gene_dataD[gene_low]
                    , data_type = data_type
                    , split_type = split_type
                    , oversampling = True # TURN IT ON TO RUN THE OTHERS BIS
                    , dst_colname = 'dst'
                    , target_colname = 'dst_mode'
                    , include_gene_name = True
                    )
            paramD = {
                    'baseline_paramD': { 'input_df'        : tempD['X']
                        , 'target'         : tempD['y']
                        , 'var_type'       : 'mixed'
                        , 'resampling_type': 'none'}
                    #, 'smnc_paramD'  : { 'input_df'          : tempD['X_smnc']
                    #                    , 'target'           : tempD['y_smnc']
                    #                    , 'var_type'         : 'mixed'
                    #                    , 'resampling_type'  : 'smnc'}
                    #, 'ros_paramD'   : { 'input_df'           : tempD['X_ros']
                    #                    , 'target'             : tempD['y_ros']
                    #                    , 'var_type'           : 'mixed'
                    #                    , 'resampling_type'    : 'ros'}
                    #, 'rus_paramD'   : { 'input_df'          : tempD['X_rus']
                    #                    , 'target'           : tempD['y_rus']
                    #                    , 'var_type'         : 'mixed'
                    #                    , 'resampling_type'  : 'rus'}
                    #, 'rouC_paramD'  : { 'input_df'         : tempD['X_rouC']
                    #                    , 'target'         : tempD['y_rouC']
                    #                    , 'var_type'       : 'mixed'
                    #                    , 'resampling_type': 'rouC'}
                    }

            out_fsD = {}
            index = 1
            for model_name, model_fn in fs_models:
                print('\nRunning classifier with FS:', index
                        , '\nModel_name:'               , model_name
                        , '\nModel func:'               , model_fn)
                #, '\nList of models:', models)
                index = index+1
                #out_fsD[model_name] = {}
                current_model = {}
                model_name_clean = model_name.replace(' ','-')

                for k, v in paramD.items():
                    out_filename = outdir + gene.lower() + '_' + split_type + '_' + data_type + '_' + model_name_clean + '_' + k + '.json'
                    fsD_params=paramD[k]

                    #out_fsD[model_name][k] = fsgs_rfecv(
                    #current_model[k] = v

                    # NOTE: this will silently fail with a syntax error if you don't have all the necessary libraries installed.
                    # Python will NOT warn you of the missing lib!
                    current_model[k] = fsgs_rfecv(
                            **fsD_params
                            , param_gridLd = [{'fs__min_features_to_select': [1]}]
                            , blind_test_df = tempD['X_bts']
                            , blind_test_target = tempD['y_bts']
                            , estimator = model_fn
                            , 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
                            )

                    # write current model to disk
                    #print(current_model)
                    print("⚠️ ⚠️ ⚠️  WRITING TO FILE: ", out_filename, "⚠️ ⚠️ ⚠️'")
                    out_json = json.dumps(current_model)
                    with open(out_filename, 'w', encoding="utf-8") as file:
                        file.write(out_json)
                    print("⚠️ ⚠️ ⚠️  Finished writing to: ", out_filename, "⚠️ ⚠️ ⚠️'")