import numpy as np
import pandas as pd
import joblib
import argparse
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.neural_network import MLPClassifier
from IPython.display import display

SIZE_OF_SAMPLE = 5000

def main():
    fields = ['sha256','strings', 'imports', 'size', 'signed', 'good_file', 'entropy']

    parser = argparse.ArgumentParser(description="File Indexer for Blue Sentry")
    parser.add_argument("--data_file", help="File that the data is output to in csv form", required=True)
    parser.add_argument("--imports_file", help="File corresponding to the imports of the data being classified",
                        required=True)
    parser.add_argument("--strings_file", help="File corresponding to the imports of the data being classified",
                        required=True)
    parser.add_argument("--classifier_imports", help="The file containing the saved imports classifier", required=True)
    parser.add_argument("--classifier_strings", help="The file containing the saved strings classifier", required=True)
    parser.add_argument("--data_out_file", help="File to add the estimations to", required=True)
    args = parser.parse_args()

    args = parser.parse_args()
    # import some data to play with
    to_test = range(50000)
    df_parser = pd.read_csv(args.data_file, chunksize=SIZE_OF_SAMPLE, iterator=True, skipinitialspace=True)

    text_encoding_imports = open(args.imports_file, 'r').read()
    text_encoding_strings = open(args.strings_file, 'r').read()
    encoder_imports = CountVectorizer().fit(text_encoding_imports.split())
    encoder_strings = CountVectorizer().fit(text_encoding_strings.split())
    clf_imports = joblib.load(args.classifier_imports)
    clf_strings = joblib.load(args.classifier_strings)

    output_file = open(args.data_out_file, 'a')

    first_file = True
    current_index = 0
    for df in df_parser:
        df.columns = fields
        data_imports_encoded = encoder_imports.transform(df['imports'])
        data_strings_encoded = encoder_strings.transform(df['strings'])
        pred_prob_imports = clf_imports.predict_proba(data_imports_encoded)
        pred_prob_strings = clf_strings.predict_proba(data_strings_encoded)
        pred_prob_rounded_imports = []
        for x in range(0, data_imports_encoded.shape[0]):
            pred_prob_rounded_imports.append(round(pred_prob_imports[x][1], 8))

        df['imports_processed'] = pred_prob_rounded_imports

        pred_prob_rounded_strings = []
        for x in range(0, data_strings_encoded.shape[0]):
            pred_prob_rounded_strings.append(round(pred_prob_strings[x][1], 8))

        df['strings_processed'] = pred_prob_rounded_strings
        df = df.drop(['strings'], axis=1)
        df = df.drop(['imports'], axis=1)
        print df

        if first_file:
            df.to_csv(output_file, header=True, index=False, index_label=False)
            first_file = False
        else:
            df.to_csv(output_file, header=False, index=False, index_label=False)
        current_index += 5000
        print current_index
    output_file.close()
'''
    data[args.data_column_name] = pred_prob_rounded
    data = data[['sha256', args.data_column_name, 'imports', 'signed', 'entropy', 'size', 'good_file']]

    data.to_csv(args.data_out_file, index=False, index_label=False)
    display(data)
'''


if __name__ == "__main__":
    main()
