diff --git a/UQ_ML_data.py b/UQ_ML_data.py
new file mode 100644
index 0000000..27a6bd2
--- /dev/null
+++ b/UQ_ML_data.py
@@ -0,0 +1,424 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Sun Mar  6 13:41:54 2022
+
+@author: tanu
+"""
+def setvars(gene,drug):
+    #https://stackoverflow.com/questions/51695322/compare-multiple-algorithms-with-sklearn-pipeline
+    import os, sys
+    import pandas as pd
+    import numpy as np
+    print(np.__version__)
+    print(pd.__version__)
+    import pprint as pp
+    from copy import deepcopy
+    from collections import Counter
+    from sklearn.impute import KNNImputer as KNN
+    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
+    
+    #%% REMOVE once config is set up
+    #from UQ_MultModelsCl import MultModelsCl
+    rs = {'random_state': 42}
+    njobs = {'n_jobs': 10}
+    
+    #%%
+    homedir = os.path.expanduser("~")
+    
+    #==============
+    # directories
+    #==============
+    datadir = homedir + '/git/Data/'
+    indir   = datadir + drug + '/input/'
+    outdir  = datadir + drug + '/output/'
+    
+    #=======
+    # input
+    #=======
+    infile_ml1 = outdir + gene.lower() + '_merged_df3.csv' 
+    #infile_ml2 = outdir + gene.lower() + '_merged_df2.csv'
+    
+    my_df = pd.read_csv(infile_ml1, index_col = 0)
+    my_df.dtypes
+    my_df_cols = my_df.columns
+    
+    geneL_basic = ['pnca']
+    
+    # -- CHECK script -- imports.py
+    #%% get cols
+    mycols = my_df.columns
+    mycols
+    
+    # change from numberic to 
+    num_type = ['int64', 'float64']
+    cat_type = ['object', 'bool']
+    
+    #TODO:
+    # #Treat active site aa pos as category and not numerical: This needs to be part of merged_df3!
+    # if my_df['active_aa_pos'].dtype in num_type:
+    #     my_df['active_aa_pos'] = my_df['active_aa_pos'].astype(object)
+    #     my_df['active_aa_pos'].dtype
+    
+    # -- CHECK script -- imports.py
+    #%%============================================================================
+    #%% IMPUTE values for OR [check script for exploration: UQ_or_imputer]
+    #or_cols = ['or_mychisq', 'log10_or_mychisq', 'or_fisher']
+    sel_cols = ['mutationinformation', 'or_mychisq', 'log10_or_mychisq']
+    or_cols = ['or_mychisq', 'log10_or_mychisq']
+    
+    print("count of NULL values before imputation\n")
+    my_df[or_cols].isnull().sum()
+    
+    my_dfI = pd.DataFrame(index = my_df['mutationinformation'] )
+    
+        
+    my_dfI = pd.DataFrame(KNN(n_neighbors= 5, weights="uniform").fit_transform(my_df[or_cols])
+                          , index =  my_df['mutationinformation']
+                          , columns = or_cols )
+    my_dfI.columns = ['or_rawI', 'logorI']
+    my_dfI.columns
+    my_dfI = my_dfI.reset_index(drop = False) # prevents old index from being added as a column
+    my_dfI.head()
+    
+    # merge with original based on index
+    my_df['index_bm'] = my_df.index
+    mydf_imputed = pd.merge(my_df
+                        , my_dfI
+                        , on = 'mutationinformation')
+    mydf_imputed = mydf_imputed.set_index(['index_bm'])
+    
+    my_df['log10_or_mychisq'].isna().sum()
+    mydf_imputed['log10_or_mychisq'].isna().sum()
+    mydf_imputed['logorI'].isna().sum()
+    
+    len(my_df.columns)
+    len(mydf_imputed.columns)  
+    
+    #%% AA property change
+    
+    # Water prop change
+    my_df['water_change'] = my_df['wt_prop_water'] + str('_to_') + my_df['mut_prop_water']
+    my_df['water_change'].value_counts()
+    water_prop_changeD = {
+        'hydrophobic_to_neutral'          : 'change'
+        , 'hydrophobic_to_hydrophobic'    : 'no_change'
+        , 'neutral_to_neutral'            : 'no_change'
+        , 'neutral_to_hydrophobic'        : 'change'
+        , 'hydrophobic_to_hydrophilic'    : 'change'
+        , 'neutral_to_hydrophilic'        : 'change'
+        , 'hydrophilic_to_neutral'        : 'change'
+        , 'hydrophilic_to_hydrophobic'    : 'change'
+        , 'hydrophilic_to_hydrophilic'    : 'no_change'
+    }
+
+    my_df['water_change'] = my_df['water_change'].map(water_prop_changeD)
+    my_df['water_change'].value_counts()
+
+    # Polarity change
+    my_df['polarity_change'] = my_df['wt_prop_polarity'] + str('_to_') + my_df['mut_prop_polarity']
+    my_df['polarity_change'].value_counts()
+
+    polarity_prop_changeD = {
+        'non-polar_to_non-polar'     : 'no_change'
+        , 'non-polar_to_neutral'     : 'change'  
+        , 'neutral_to_non-polar'     : 'change'  
+        , 'neutral_to_neutral'       : 'no_change'  
+        , 'non-polar_to_basic'       : 'change'  
+        , 'acidic_to_neutral'        : 'change'  
+        , 'basic_to_neutral'         : 'change'  
+        , 'non-polar_to_acidic'      : 'change'  
+        , 'neutral_to_basic'         : 'change'  
+        , 'acidic_to_non-polar'      : 'change'  
+        , 'basic_to_non-polar'       : 'change'
+        , 'neutral_to_acidic'        : 'change'
+        , 'acidic_to_acidic'         : 'no_change'
+        , 'basic_to_acidic'          : 'change'
+        , 'basic_to_basic'           : 'no_change'
+        , 'acidic_to_basic'          : 'change'}
+
+    my_df['polarity_change'] = my_df['polarity_change'].map(polarity_prop_changeD)
+    my_df['polarity_change'].value_counts()
+
+    # Electrostatics change
+    my_df['electrostatics_change'] = my_df['wt_calcprop'] + str('_to_') + my_df['mut_calcprop']
+    my_df['electrostatics_change'].value_counts()
+
+    calc_prop_changeD = {
+            'non-polar_to_non-polar'     : 'no_change'
+            , 'non-polar_to_polar'       : 'change'
+            , 'polar_to_non-polar'       : 'change'
+            , 'non-polar_to_pos'         : 'change'
+            , 'neg_to_non-polar'         : 'change'
+            , 'non-polar_to_neg'         : 'change'
+            , 'pos_to_polar'             : 'change'
+            , 'pos_to_non-polar'         : 'change'
+            , 'polar_to_polar'           : 'no_change'
+            , 'neg_to_neg'               : 'no_change'
+            , 'polar_to_neg'             : 'change'
+            , 'pos_to_neg'               : 'change'
+            , 'pos_to_pos'               : 'no_change'
+            , 'polar_to_pos'             : 'change'
+            , 'neg_to_polar'             : 'change'
+            , 'neg_to_pos'               : 'change'
+    }
+
+    my_df['electrostatics_change'] = my_df['electrostatics_change'].map(calc_prop_changeD)
+    my_df['electrostatics_change'].value_counts()
+    
+    # Create a combined column summarising these three cols
+    detect_change = 'change'
+    check_prop_cols = ['water_change', 'polarity_change', 'electrostatics_change']
+    #my_df['aa_prop_change'] = (my_df.values == detect_change).any(1).astype(int)
+    my_df['aa_prop_change'] = (my_df[check_prop_cols].values == detect_change).any(1).astype(int)
+    my_df['aa_prop_change'].value_counts()
+    my_df['aa_prop_change'].dtype
+    
+    my_df['aa_prop_change'] = my_df['aa_prop_change'].map({1:'change'
+                                                           , 0: 'no_change'})
+
+    my_df['aa_prop_change'].value_counts()
+    my_df['aa_prop_change'].dtype
+    #%% Add lineage calc
+    total_mtblineage_u = 8
+
+    lineage_colnames = ['lineage_list_all', 'lineage_count_all', 'lineage_count_unique', 'lineage_list_unique', 'lineage_multimode']
+    #bar = my_df[lineage_colnames]
+    my_df['lineage_proportion']      = my_df['lineage_count_unique']/my_df['lineage_count_all']
+    my_df['dist_lineage_proportion'] = my_df['lineage_count_unique']/total_mtblineage_u
+
+    #%% Combine mmCSM_lig Data: DONE
+    #%% Combine PROVEAN data: DONE
+    #%% Combine ED logo data
+
+    #%% Masking columns (mCSM-lig, mCSM-NA, mCSM-ppi2) values for lig_dist >10
+    
+    # get logic from upstream!
+    #my_df_ml = my_df.copy()
+    my_df_ml = mydf_imputed.copy()
+    
+    my_df_ml['mutationinformation'][my_df['ligand_distance']>10].value_counts()
+    my_df_ml.groupby('mutationinformation')['ligand_distance'].apply(lambda x: (x>10)).value_counts()
+    my_df_ml.groupby(['mutationinformation'])['ligand_distance'].apply(lambda x: (x>10)).value_counts()
+    
+    my_df_ml.loc[(my_df_ml['ligand_distance'] > 10), 'ligand_affinity_change'] = 0
+    (my_df_ml['ligand_affinity_change'] == 0).sum()
+    
+    #%%============================================================================
+    # Separate blind test set
+    my_df_ml[drug].isna().sum()
+    
+    blind_test_df = my_df_ml[my_df_ml[drug].isna()]
+    blind_test_df.shape
+    
+    training_df =  my_df_ml[my_df_ml[drug].notna()]
+    training_df.shape
+    
+    # Target1: dst
+    training_df[drug].value_counts()
+    training_df['dst_mode'].value_counts()
+    
+    #%% Build X
+    common_cols_stabiltyN = ['ligand_distance'
+               , 'ligand_affinity_change'
+               , 'duet_stability_change'
+               , 'ddg_foldx'
+               , 'deepddg'
+               , 'ddg_dynamut2']
+    
+    foldX_cols = ['contacts'
+    , 'electro_rr', 'electro_mm', 'electro_sm', 'electro_ss'
+    , 'disulfide_rr', 'disulfide_mm', 'disulfide_sm', 'disulfide_ss'
+    , 'hbonds_rr', 'hbonds_mm', 'hbonds_sm', 'hbonds_ss'
+    , 'partcov_rr', 'partcov_mm', 'partcov_sm', 'partcov_ss'
+    , 'vdwclashes_rr', 'vdwclashes_mm', 'vdwclashes_sm', 'vdwclashes_ss'
+    , 'volumetric_rr', 'volumetric_mm', 'volumetric_ss'
+    ]
+    
+    X_strFN =  ['rsa'
+               #, 'asa'
+               , 'kd_values'
+               , 'rd_values']    
+    
+    X_evolFN =  ['consurf_score'
+               , 'snap2_score']
+    
+    # quick inspection which lineage to use:
+    #foo = my_df_ml[['lineage', 'lineage_count_all', 'lineage_count_unique']]
+        
+    X_genomicFN =  ['maf'
+    #            , 'or_mychisq'
+    #            , 'or_logistic'
+    #            , 'or_fisher'
+    #            , 'pval_fisher'
+                 #, 'lineage'
+                 #, 'lineage_count_all'
+                 #, 'lineage_count_unique'
+    ]
+    
+    #%% Construct numerical and categorical column names
+    
+    # numerical feature names
+    numerical_FN = common_cols_stabiltyN + foldX_cols + X_strFN + X_evolFN + X_genomicFN 
+    
+    #categorical feature names
+    categorical_FN = ['ss_class'
+    #             , 'wt_prop_water'
+                # , 'lineage_labels' # misleading if using merged_df3
+    #             , 'mut_prop_water'
+    #             , 'wt_prop_polarity'
+    #             , 'mut_prop_polarity'
+    #             , 'wt_calcprop'
+    #             , 'mut_calcprop'
+                 #, 'active_aa_pos'
+                 ]
+    
+    #%% extracting dfs based on numerical, categorical column names
+    #----------------------------------
+    # WITHOUT the target var included
+    #----------------------------------
+    num_df = training_df[numerical_FN]
+    num_df.shape
+    
+    cat_df = training_df[categorical_FN]
+    cat_df.shape
+    
+    all_df = training_df[numerical_FN + categorical_FN]
+    all_df.shape
+    
+    #------------------------------
+    # WITH the target var included:
+        #'wtgt': with target
+    #------------------------------
+    # drug and dst_mode should be the same thing
+    num_df_wtgt = training_df[numerical_FN + ['dst_mode']]
+    num_df_wtgt.shape
+    
+    cat_df_wtgt = training_df[categorical_FN + ['dst_mode']]
+    cat_df_wtgt.shape
+    
+    all_df_wtgt = training_df[numerical_FN + categorical_FN + ['dst_mode']]
+    all_df_wtgt.shape
+    #%%================================================================
+    #%% Apply ML
+    
+    #%% Data
+    #------
+    # X
+    #------
+    X     = all_df_wtgt[numerical_FN + categorical_FN] # training data ALL
+    X_bts = blind_test_df[numerical_FN + categorical_FN] # blind test data ALL
+    #X = all_df_wtgt[numerical_FN] # training numerical only
+    #X_bts = blind_test_df[numerical_FN] # blind test data numerical
+    
+    #------
+    # y
+    #------
+    y = all_df_wtgt['dst_mode'] # training data y
+    y_bts = blind_test_df['dst_mode'] # blind data test y
+    
+    #Blind test data {same format}
+    #X_bts = blind_test_df[numerical_FN]
+    #X_bts = blind_test_df[numerical_FN + categorical_FN]
+    #y_bts = blind_test_df['dst_mode']
+    
+    X_bts_wt = blind_test_df[numerical_FN + ['dst_mode']] 
+    
+    # Quick check
+    (X['ligand_affinity_change']==0).sum() == (X['ligand_distance']>10).sum()
+    ##############################################################################
+    print('Original Data\n', Counter(y)
+          , 'Data dim:', X.shape)
+    
+    ###############################################################################
+    #%% 
+    ############################################################################
+    #                               RESAMPLING
+    ###############################################################################
+    #------------------------------
+    # Simple Random oversampling
+    # [Numerical + catgeorical]
+    #------------------------------
+    oversample = RandomOverSampler(sampling_strategy='minority')
+    X_ros, y_ros = oversample.fit_resample(X, y)
+    print('Simple Random OverSampling\n', Counter(y_ros))
+    print(X_ros.shape)
+    
+    #------------------------------
+    # Simple Random Undersampling
+    # [Numerical + catgeorical]
+    #------------------------------
+    undersample = RandomUnderSampler(sampling_strategy='majority')
+    X_rus, y_rus = undersample.fit_resample(X, y)
+    print('Simple Random UnderSampling\n', Counter(y_rus))
+    print(X_rus.shape)
+    
+    #------------------------------
+    # Simple combine ROS and RUS
+    # [Numerical + catgeorical]
+    #------------------------------
+    oversample = RandomOverSampler(sampling_strategy='minority')
+    X_ros, y_ros = oversample.fit_resample(X, y)
+    undersample = RandomUnderSampler(sampling_strategy='majority')
+    X_rouC, y_rouC = undersample.fit_resample(X_ros, y_ros)
+    print('Simple Combined Over and UnderSampling\n',  Counter(y_rouC))
+    print(X_rouC.shape)
+    
+    #------------------------------
+    # SMOTE_NC: oversampling 
+    # [numerical + categorical]
+    #https://stackoverflow.com/questions/47655813/oversampling-smote-for-binary-and-categorical-data-in-python
+    #------------------------------
+    # Determine categorical and numerical features
+    numerical_ix = X.select_dtypes(include=['int64', 'float64']).columns
+    numerical_ix
+    num_featuresL = list(numerical_ix)
+    numerical_colind = X.columns.get_indexer(list(numerical_ix) )
+    numerical_colind
+    
+    categorical_ix = X.select_dtypes(include=['object', 'bool']).columns
+    categorical_ix    
+    categorical_colind = X.columns.get_indexer(list(categorical_ix))
+    categorical_colind
+    
+    k_sm = 5 # 5 is deafult
+    sm_nc = SMOTENC(categorical_features=categorical_colind, k_neighbors = k_sm, **rs, **njobs)
+    X_smnc, y_smnc = sm_nc.fit_resample(X, y)
+    print('SMOTE_NC OverSampling\n', Counter(y_smnc))
+    print(X_smnc.shape)
+    globals().update(locals()) # TROLOLOLOLOLOLS
+    #print("i did a horrible hack :-)")
+    ###############################################################################
+    #%% SMOTE RESAMPLING for NUMERICAL ONLY*
+    # #------------------------------
+    # # SMOTE: Oversampling
+    # # [Numerical ONLY]
+    # #------------------------------
+    # k_sm = 1
+    # sm = SMOTE(sampling_strategy = 'auto', k_neighbors = k_sm, **rs)
+    # X_sm, y_sm = sm.fit_resample(X, y)
+    # print(X_sm.shape)
+    # print('SMOTE OverSampling\n', Counter(y_sm))
+    # y_sm_df = y_sm.to_frame()
+    # y_sm_df.value_counts().plot(kind = 'bar')
+    
+    # #------------------------------
+    # # SMOTE: Over + Undersampling COMBINED
+    # # [Numerical ONLY]
+    # #-----------------------------
+    # sm_enn = SMOTEENN(enn=EditedNearestNeighbours(sampling_strategy='all', **rs, **njobs ))
+    # X_enn, y_enn = sm_enn.fit_resample(X, y)
+    # print(X_enn.shape)
+    # print('SMOTE Over+Under Sampling combined\n', Counter(y_enn))
+    
+    ###############################################################################
+    # TODO: Find over and undersampling JUST for categorical data