#!/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'
                 # , 'embB' : 'ethambutol'
                 # , 'katG' : 'isoniazid'
                 # , 'rpoB' : 'rifampicin'
                 # , 'gid'  : 'streptomycin'
                 }
gene_dataD={}
# NOTE: for gid, run 'actual' on 80/20 and sl only
split_types = ['70_30', '80_20', 'sl']
split_data_types = ['actual', 'complete']
#split_types = ['70_30']
#split_data_types = ['actual', 'complete']

#fs_models  = [('Logistic Regression'       , LogisticRegression(**rs) )]

fs_models = [
        ('Logistic Regression'       , LogisticRegression(**rs) )
        , ('Ridge Classifier'          , RidgeClassifier(**rs)  )
        #, ('AdaBoost Classifier'     , AdaBoostClassifier(**rs) )
        #, ('Decision Tree'             , DecisionTreeClassifier(**rs) ) 
        #, ('Extra Tree'                , ExtraTreeClassifier(**rs) )
        #, ('Extra Trees'               , ExtraTreesClassifier(**rs) ) 
        #, ('Gradient Boosting'         , GradientBoostingClassifier(**rs) )
        #, ('LDA'                       , LinearDiscriminantAnalysis() )
        #, ('Logistic RegressionCV'     , LogisticRegressionCV(cv = 3, **rs))
        #, ('Passive Aggresive'         , PassiveAggressiveClassifier(**rs, **njobs) )
        #, ('Random Forest'             , RandomForestClassifier(**rs, n_estimators = 1000 ) ) 
        #, ('Ridge ClassifierCV'        , RidgeClassifierCV(cv = 3) )          
        #, ('Stochastic GDescent'       , SGDClassifier(**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 = {}

                for k, v in paramD.items():
                    out_filename = gene.lower() + '_' + split_type + '_' + data_type + '_' + model_name + '_' + 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)
                    out_json = json.dumps(current_model)
                    with open(out_filename, 'w', encoding="utf-8") as file:
                        file.write(out_json)