import numpy as np
import sklearn
import subprocess
from sklearn import model_selection, tree

import data
import feature_selection
import model_sel
import os

import matplotlib.pyplot as plt
import seaborn as sns

def main():

    #parameter space
    list_test_size = [0.1,0.15,0.2] # decide this
    list_ftsel_method = ['chi2','mutlinfo','pca','dt']
    list_num_features = [10,15,20] # decide this
    list_Kfold = [3,5]
    list_corr_threshold = [1,0.5,0.6,0.7] # decide this

    param_space = {
        'criterion': ['gini', 'entropy'],
        'max_depth': [2, 3, 4, 5, 7],
        'min_samples_split': [2, 5, 10],
        'min_samples_leaf': [2, 5, 10],
        'max_leaf_nodes': [2, 4, 6, 8, 10, 12, 15],
    }

    repeat = 1

    #output dictrionary list
    list_output_dict = []

    # output directory path
    outdir = "../results/run1/"

    if(not os.path.isdir(outdir)):
        os.mkdir(outdir)

    o_models_file = open(outdir+"models.csv","w")
    o_models_file.write("test size,run num,ftsel method,Kfold,number of features,correlation threshold,best features,criterion,max_depth,max_leaf_nodes,min_samples_leaf,min_samples_split,training accuracy,test accuracy\n")

    #splitting data and weights into train, test (refer to optimal_params.py)
    poll_data = data.PollDataProxy(remove_nan=False, convert_to_float=False)

    acc = []

    '''refer to optimal_params.py. Functions from this python scripts are transferred here. (get_bad_questions() and separate_weights().)'''

    for ts in list_test_size:
        for run_num in range(repeat):
            all_data, all_data_questions = poll_data.all_data_except(get_bad_questions())
            X = all_data[:, :-1]
            y = all_data[:, -1]
            X_train, X_test, y_train, y_test = model_selection.train_test_split(X, y,
                                                                                test_size=ts,
                                                                                shuffle=True)
            X_train, weights_train, questions = separate_weights(X_train, all_data_questions[:-1])
            X_test, weights_test, _ = separate_weights(X_test, all_data_questions[:-1])

            print("Number of Training Samples:", len(X_train))
            print("Number of Testing Samples:", len(X_test))


            data_dict = {
                'X_train': X_train,
                'X_test': X_test,
                'y_train': y_train,
                'y_test': y_test
            }
            weights_dict = {
                'weights_train': weights_train,
                'weights_test': weights_test}

            for meth in list_ftsel_method:
                '''Create class objects of the current selection method'''
                for thres in list_corr_threshold:
                    data_ranked_dict, ranked_questions = {}, []
                    ftsel_obj =None
                    if(meth=='chi2'):
                        ftsel_obj = feature_selection.FeatureSelection(
                            necess_que_file="../extern/manage_data/list_all_questions.txt",
                            unnecess_que_file="../extern/manage_data/list_unnecessary_columns.txt",
                            bool_necess_que=False,
                            run_name="test_chi2"
                        )
                        data_ranked_dict, ranked_questions = ftsel_obj.ftsel_chi2(data_dict, thres)
                    elif(meth=='mutlinfo'):
                        ftsel_obj = feature_selection.FeatureSelection(
                            necess_que_file="../extern/manage_data/list_all_questions.txt",
                            unnecess_que_file="../extern/manage_data/list_unnecessary_columns.txt",
                            bool_necess_que=False,
                            run_name="test_mutlinfo"
                        )
                        data_ranked_dict, ranked_questions = ftsel_obj.ftsel_mutlinfo(data_dict, thres)
                    elif(meth=='pca'):
                        ftsel_obj = feature_selection.FeatureSelection(
                            necess_que_file="../extern/manage_data/list_all_questions.txt",
                            unnecess_que_file="../extern/manage_data/list_unnecessary_columns.txt",
                            bool_necess_que=False,
                            run_name="test_pca"
                        )
                        data_ranked_dict,_ = ftsel_obj.ftsel_pca(data_dict)
                        fts = data_sel_dict['X_train'].shape[1]
                        questions_int = list(map(str, list(range(1,fts+1,1))))
                        ranked_questions = ["ft_"+x for x in questions_int]
                    elif(meth=='dt'):
                        ftsel_obj = feature_selection.FeatureSelection(
                            necess_que_file="../extern/manage_data/list_all_questions.txt",
                            unnecess_que_file="../extern/manage_data/list_unnecessary_columns.txt",
                            bool_necess_que=False,
                            run_name="test_dt"
                        )
                        data_ranked_dict, ranked_questions = ftsel_obj.ftsel_decision_tree_method(data_dict, thres)
                    for num in list_num_features:
                        data_sel_dict, sel_questions = ftsel_obj.select_num_features(data_ranked_dict, num, ranked_questions)
                        ftsel_obj.plot_heatmap(data_sel_dict['X_train'], sel_questions)
                        for K in list_Kfold:
                            '''Here create a class onject of "model_sel" and output all the best parameters and values into "list_output_dict". Then, can create a .csv file to list all the models and accuracies.'''
                            model_obj = model_sel.model_sel(ts, run_num, meth, param_space, K, num, thres, data_sel_dict ,weights_dict, sel_questions, outdir).select_model()
                            #   intermediate = model_obj.select_model()
                            acc.append(model_obj['test_acc'])

                            o_models_file.write(str(ts)+",")
                            o_models_file.write(str(run_num)+",")
                            o_models_file.write(meth+",")
                            o_models_file.write(str(K)+",")
                            o_models_file.write(str(num)+",")
                            o_models_file.write(str(thres)+",")
                            for ii in range(len(model_obj['best_features'])):
                                o_models_file.write(model_obj['best_features'][ii]+" ")
                            o_models_file.write(",")
                            o_models_file.write(model_obj['best_params']['criterion']+",")
                            o_models_file.write(str(model_obj['best_params']['max_depth'])+",")
                            o_models_file.write(str(model_obj['best_params']['max_leaf_nodes'])+",")
                            o_models_file.write(str(model_obj['best_params']['min_samples_leaf'])+",")
                            o_models_file.write(str(model_obj['best_params']['min_samples_split'])+",")
                            o_models_file.write(str(model_obj['train_acc'])+",")
                            o_models_file.write(str(model_obj['test_acc'])+",")
                            o_models_file.write("\n")

                            list_output_dict.append(model_obj)
    '''Once all the models are run, select the model with best test accuracy and return the output dict for that model.'''
    o_models_file.close()

    best_index = np.argmax(acc)
    best_model_dict = list_output_dict[best_index]

    print("The best model parameters:")
    print(best_model_dict)


def get_bad_questions():
    f = open("../extern/manage_data/list_unnecessary_columns.txt", 'r')
    bad_questions = f.readline().split(',')
    bad_questions[-1] = bad_questions[-1][:-1]  # chop the \n off the end
    bad_questions.remove('weight')  # need weight for training
    return bad_questions


def separate_weights(X_train, column_names):
    """
    Removes the column containing weights from X_train, and returns it as
    a separate array.
    """
    weight_column_idx = column_names.index('weight')
    weights = X_train[:, weight_column_idx]
    new_X_train = np.delete(X_train, weight_column_idx, axis=1)
    new_questions = column_names
    new_questions.remove('weight')
    return new_X_train, weights, new_questions


if __name__ == "__main__":
    main()

[(46, 'data.PollDataProxy', 'data.PollDataProxy', False, 'import data\n'), (148, 'numpy.argmax', 'np.argmax', True, 'import numpy as np\n'), (170, 'numpy.delete', 'np.delete', True, 'import numpy as np\n'), (39, 'os.path.isdir', 'os.path.isdir', False, 'import os\n'), (40, 'os.mkdir', 'os.mkdir', False, 'import os\n'), (57, 'sklearn.model_selection.train_test_split', 'model_selection.train_test_split', False, 'from sklearn import model_selection, tree\n'), (83, 'feature_selection.FeatureSelection', 'feature_selection.FeatureSelection', False, 'import feature_selection\n'), (91, 'feature_selection.FeatureSelection', 'feature_selection.FeatureSelection', False, 'import feature_selection\n'), (99, 'feature_selection.FeatureSelection', 'feature_selection.FeatureSelection', False, 'import feature_selection\n'), (110, 'feature_selection.FeatureSelection', 'feature_selection.FeatureSelection', False, 'import feature_selection\n'), (122, 'model_sel.model_sel', 'model_sel.model_sel', False, 'import model_sel\n')]