import numpy as np
import tensorflow as tf
import os
import csv
import pickle
import pandas
from yadlt.models.autoencoders import deep_autoencoder
from yadlt.utils import datasets, utilities
from unigram_extract import get_binary_headers,get_binary_api,get_unigrams,get_headers
from creatAsm.api_header_get import get_header_api
from collections import defaultdict
from creatAsm.main import get_opcode_csv
from keras.models import Sequential
from keras.layers import Dense, Activation,Dropout
import gc
from sklearn import preprocessing
# #################### #
#   Flags definition   #
# #################### #
flags = tf.app.flags
FLAGS = flags.FLAGS

# Global configuration
flags.DEFINE_string('dataset', 'custom', 'Which dataset to use. ["mnist", "cifar10", "custom"]')
flags.DEFINE_string('train_dataset', '', 'Path to train set data .npy file.')
flags.DEFINE_string('train_ref', '', 'Path to train reference .npy file.')
flags.DEFINE_string('valid_dataset', '', 'Path to valid set .npy file.')
flags.DEFINE_string('valid_ref', '', 'Path to valid reference data .npy file.')
flags.DEFINE_string('test_dataset', '', 'Path to Son set .npy file.')
flags.DEFINE_string('test_ref', '', 'Path to Son reference data .npy file.')
flags.DEFINE_string('cifar_dir', '', 'Path to the cifar 10 dataset directory.')
flags.DEFINE_integer('seed', -1, 'Seed for the random generators (>= 0). Useful for testing hyperparameters.')
flags.DEFINE_boolean('do_pretrain', True, 'Whether or not doing unsupervised pretraining.')
flags.DEFINE_string('save_reconstructions',
                    'E:/Pycharm Projects/Deep-Learning-TensorFlow/model/save_reconstructions/reconstruction_arr',
                    'Path to a .npy file to save the reconstructions of the model.')
flags.DEFINE_string('save_layers_output_test', '',
                    'Path to a .npy file to save Son set output from all the layers of the model.')
flags.DEFINE_string('save_layers_output_train', 'E:/Pycharm Projects/Deep-Learning-TensorFlow/model/',
                    'Path to a .npy file to save train set output from all the layers of the model.')
flags.DEFINE_string('save_model_parameters', '',
                    'Path to a directory to save the parameters of the model. One .npy file per layer.')
flags.DEFINE_string('name', 'un_sdae', 'Name for the model.')
flags.DEFINE_float('momentum', 0.5, 'Momentum parameter.')
flags.DEFINE_boolean('tied_weights', True, 'Whether to use tied weights for the decoders.')

# Supervised fine tuning parameters
flags.DEFINE_string('finetune_loss_func', 'cross_entropy', 'Last Layer Loss function.["cross_entropy", "mse"]')
flags.DEFINE_integer('finetune_num_epochs', 30, 'Number of epochs for the fine-tuning phase.')
flags.DEFINE_float('finetune_learning_rate', 0.001, 'Learning rate for the fine-tuning phase.')
flags.DEFINE_string('finetune_enc_act_func', 'relu,',
                    'Activation function for the encoder fine-tuning phase. ["sigmoid, "tanh", "relu"]')
flags.DEFINE_string('finetune_dec_act_func', 'sigmoid,',
                    'Activation function for the decoder fine-tuning phase. ["sigmoid, "tanh", "relu"]')
flags.DEFINE_float('finetune_dropout', 1, 'Dropout parameter.')
flags.DEFINE_string('finetune_opt', 'sgd', '["sgd", "ada_grad", "momentum"]')
flags.DEFINE_integer('finetune_batch_size', 20, 'Size of each mini-batch for the fine-tuning phase.')

# Autoencoder layers specific parameters
flags.DEFINE_string('dae_layers', '256,', 'Comma-separated values for the layers in the sdae.')
flags.DEFINE_string('dae_regcoef', '5e-4,', 'Regularization parameter for the autoencoders. If 0, no regularization.')
flags.DEFINE_string('dae_enc_act_func', 'sigmoid,', 'Activation function for the encoder. ["sigmoid", "tanh"]')
flags.DEFINE_string('dae_dec_act_func', 'none,', 'Activation function for the decoder. ["sigmoid", "tanh", "none"]')
flags.DEFINE_string('dae_loss_func', 'mse,', 'Loss function. ["mse" or "cross_entropy"]')
flags.DEFINE_string('dae_opt', 'sgd,', '["sgd", "ada_grad", "momentum", "adam"]')
flags.DEFINE_string('dae_learning_rate', '0.01,', 'Initial learning rate.')
flags.DEFINE_string('dae_num_epochs', '10,', 'Number of epochs.')
flags.DEFINE_string('dae_batch_size', '10,', 'Size of each mini-batch.')
flags.DEFINE_string('dae_corr_type', 'none,', 'Type of input corruption. ["none", "masking", "salt_and_pepper"]')
flags.DEFINE_string('dae_corr_frac', '0.0,', 'Fraction of the input to corrupt.')

# Conversion of Autoencoder layers parameters from string to their specific type
dae_layers = utilities.flag_to_list(FLAGS.dae_layers, 'int')
dae_enc_act_func = utilities.flag_to_list(FLAGS.dae_enc_act_func, 'str')
dae_dec_act_func = utilities.flag_to_list(FLAGS.dae_dec_act_func, 'str')
dae_opt = utilities.flag_to_list(FLAGS.dae_opt, 'str')
dae_loss_func = utilities.flag_to_list(FLAGS.dae_loss_func, 'str')
dae_learning_rate = utilities.flag_to_list(FLAGS.dae_learning_rate, 'float')
dae_regcoef = utilities.flag_to_list(FLAGS.dae_regcoef, 'float')
dae_corr_type = utilities.flag_to_list(FLAGS.dae_corr_type, 'str')
dae_corr_frac = utilities.flag_to_list(FLAGS.dae_corr_frac, 'float')
dae_num_epochs = utilities.flag_to_list(FLAGS.dae_num_epochs, 'int')
dae_batch_size = utilities.flag_to_list(FLAGS.dae_batch_size, 'int')

finetune_enc_act_func = utilities.flag_to_list(FLAGS.finetune_enc_act_func, 'str')
finetune_dec_act_func = utilities.flag_to_list(FLAGS.finetune_dec_act_func, 'str')

# Parameters validation
assert all([0. <= cf <= 1. for cf in dae_corr_frac])
assert all([ct in ['masking', 'salt_and_pepper', 'none'] for ct in dae_corr_type])
assert FLAGS.dataset in ['mnist', 'cifar10', 'custom']
assert len(dae_layers) > 0
assert all([af in ['sigmoid', 'tanh'] for af in dae_enc_act_func])
assert all([af in ['sigmoid', 'tanh', 'none'] for af in dae_dec_act_func])


# 命令行
# python  Deep-Learning-TensorFlow/cmd_line/autoencoders/run_stacked_autoencoder_unsupervised.py --dae_layers 512,256,128 --dae_batch_size 25 --dae_num_epochs 5 --verb
# ose 1 --dae_corr_type masking --dae_corr_frac 0.0 --finetune_num_epochs 25 --finetune_batch_size 32 --finetune_opt gradient_descent --finetune_learning_rate 0.05 --dae_enc_act_func sigm
# oid --dae_dec_act_func sigmoid --dae_loss_func cross_entropy --finetune_enc_act_func tanh --finetune_dec_act_func sigmoid --finetune_loss_func cross_entropy --dropout 0.7
def load_keras_model(test_arr):
    model = Sequential()
    model.add(Dense(30, input_dim=30, activation='relu'))
    model.add(Dropout(0.5))
    model.add(Dense(30, activation='relu'))
    model.add(Dropout(0.5))
    model.add(Dense(30, activation='relu'))
    model.add(Dropout(0.5))
    model.add(Dense(1, activation='sigmoid'))
    # adam rmsprop adagrad sgd
    model.compile(optimizer='rmsprop', loss='binary_crossentropy',
                  metrics=['accuracy'])
    model.load_weights('E:\\malware_detector\\web\\DAE_signature_weights_10000_dae5_finetune25_batchsize30_header_.h5')
    return model.predict_classes(test_arr)

def load_from_csv(dataset_path):
    # data_arr=np.array()
    data = pandas.read_csv(dataset_path,header=None)
    data=data.drop(data.columns[0],axis=1).values
    # data=data.drop(data.columns=[0], axis=1)
    # with open(dataset_path) as fin:
    #     reader = csv.reader(fin)
    #     for row in reader:
    #         data_arr = np.append(data_arr, row[1:])
    # data.dtype='float'
    #print("csv_data.shape:",data.shape)
    return data
def dicvalue_to_numpy(dic_values):
    list=[]
    for row in dic_values:
        list.append(row)
    return list
#删除为0的列
def del_zero_col(data):
    zero_col = []
    for i in range(data.shape[1]):
        if sum(data[:, i]) == 0:
            zero_col.append(i)
    data = np.delete(data, zero_col, 1)
    return data
if __name__ == '__main__':

    binary_benign_arr=load_from_csv("E:\malware_detector\web\\apiHeaderOpcodeBenign.csv")

    binary_malware_arr = load_from_csv("E:\malware_detector\web\\apiHeaderOpcodeMalware.csv")

    # 测试生成单个文件的opcode的特征序列
    print('正在生成opcode特征序列...\n')
    exeFilePath="E:\\malware_detector\\web\\creatAsm\\exeFile\\2018-9-8"
    # get_opcode_csv('E:\Pycharm Projects\creatAsm\exeFile')
    get_opcode_csv(exeFilePath)
    #opcode_arr=load_from_csv('E:\\malware_detector\\web\\creatAsm\\malfeature.csv')
    #测试提取单个样本api、header特征
    # header_list,api_list=get_header_api('E:\Pycharm Projects\creatAsm\exeFile')
    print("正在生成api、header特征序列...\n")
    header_list, api_list = get_header_api(exeFilePath)
    #测试生成单个文件的api的0,1序列
    unigramsOfeachFile_dic = defaultdict(list)
    for line in api_list:
        if line[0]:
            file_name=line[0]
        for col in line[2:]:
            if file_name not in unigramsOfeachFile_dic.keys():
                unigramsOfeachFile_dic[file_name] = []
            if col not in unigramsOfeachFile_dic[file_name] and col != '' and col:
                col = str(col)
                if col.isupper():
                    col = '.' + col.lower()
                unigramsOfeachFile_dic[file_name].append(col)
    f=open('E:\\malware_detector\\web\\apiOfall','rb')
    apiOfall=pickle.load(f)
    f.close()
    binary_api=get_binary_api(unigramsOfeachFile_dic,apiOfall)
    #print('binary_api:',binary_api)
    api_arr=np.array(dicvalue_to_numpy(binary_api))
    #print("api_arr:",api_arr.shape)
    #测试生成单个文件的header的特征序列
    unigram_dic1=defaultdict(list)
    unigram_dic1[header_list[0]]=header_list[1:]
    f = open('E:\\malware_detector\\web\\headerOfall', 'rb')
    f2 = open('E:\\malware_detector\\web\\headerColName', 'rb')
    headerOfall = pickle.load(f)
    headerColName=pickle.load(f2)
    print("生成api、header特征序列成功!\n")
    f.close()
    binary_header=get_binary_headers(unigram_dic1, headerOfall, headerColName)
    #print('binary header:',binary_header)
    header_arr=np.array(dicvalue_to_numpy(binary_header))
    #print("header_arr:",header_arr.shape)

    # test_single_sampel_arr=np.concatenate((opcode_arr,api_arr,header_arr),axis=1)
    test_single_sampel_arr=header_arr
    #print('Son sample shape',test_single_sampel_arr.shape)
    #删除无用变量



    #print('benign shape', binary_benign_arr.shape)
    #print('malware shape', binary_malware_arr.shape)
    #归一化处理
    binary_arr=np.concatenate((binary_benign_arr,binary_malware_arr),axis=0)
    # binary_arr=preprocessing.scale(binary_arr)
    #删除全为0的列
    # binary_arr=del_zero_col(binary_arr)
    #分解出去0后的恶意非恶意数组
    # binary_benign_arr=binary_arr[:binary_benign_arr.shape[0],:]
    # binary_malware_arr=binary_arr[binary_benign_arr.shape[0]:,:]
    #按行随机打乱
    np.random.shuffle(binary_arr)
    # binary_arr = binary_benign_arr
    encode_dir = "E:\\malware_detector\\web\\encode_test_arr.npy"
    utilities.random_seed_np_tf(FLAGS.seed)
    #print(FLAGS.dataset)
    if FLAGS.dataset == 'mnist':

        # ################# #
        #   MNIST Dataset   #
        # ################# #

        trX, vlX, teX = datasets.load_mnist_dataset(mode='unsupervised')
        trRef = trX
        print(trX.shape)
        print(trX.shape[1])
        # print(trX[0])
        vlRef = vlX
        teRef = teX

    elif FLAGS.dataset == 'cifar10':

        # ################### #
        #   Cifar10 Dataset   #
        # ################### #

        trX, teX = datasets.load_cifar10_dataset(FLAGS.cifar_dir, mode='unsupervised')
        # Validation set is the first half of the Son set
        vlX = teX[:5000]
        trRef = trX
        vlRef = vlX
        teRef = teX

    elif FLAGS.dataset == 'custom':

        # ################## #
        #   Custom Dataset   #
        # ################## #

        def load_from_np(dataset_path, set_type):
            if dataset_path != '':
                binary_arr_len = len(binary_arr)
                train_len = int(binary_arr_len * 0.6 + 0.5)
                varify_len = int(binary_arr_len * 0.8 + 0.5)
                if set_type == 'train':
                    return binary_arr[:train_len, :], None
                elif set_type == 'varify':
                    return binary_arr[train_len:varify_len, :], None
                elif set_type == 'Son':
                    return binary_arr[varify_len:, :], None
                else:
                    return binary_arr, None
                    # return np.load(dataset_path)
            else:
                return None


        #trX,trRef=load_from_np("E:\Pycharm Projects\\benign-headers.csv",'train')
        #vlX,vlRef=load_from_np("E:\Pycharm Projects\\benign-headers.csv",'varify')
        #teX,teRef=load_from_np("E:\Pycharm Projects\\benign-headers.csv",'Son')



        #Son the model
        #不需要训练测试集
        trX,trRef=binary_arr[:1,35486:],None
        vlX,vlRef=trX,None
        teX,teRef=trX,None




        #删除无用变量
        del binary_arr,binary_benign_arr,binary_malware_arr,header_arr,unigram_dic1,header_list,api_list,unigramsOfeachFile_dic,headerColName,binary_header,binary_api
        gc.collect()
        # vlX = trX[:200, :]
        # teX = trX
        # vlRef = None
        # teRef = None
        #print('train shape', trX.shape)
        #print('varify shape', vlX.shape)
        #print('Son shape', teX.shape)
        # trX, trRef = load_from_np(FLAGS.train_dataset), load_from_np(FLAGS.train_ref)
        # vlX, vlRef = load_from_np(FLAGS.valid_dataset), load_from_np(FLAGS.valid_ref)
        # teX, teRef = load_from_np(FLAGS.test_dataset), load_from_np(FLAGS.test_ref)

        if not trRef:
            trRef = trX
        if not vlRef:
            vlRef = vlX
        if not teRef:
            teRef = teX

    else:
        trX = None
        trRef = None
        vlX = None
        vlRef = None
        teX = None
        teRef = None

    # Create the object
    sdae = None

    dae_enc_act_func = [utilities.str2actfunc(af) for af in dae_enc_act_func]
    dae_dec_act_func = [utilities.str2actfunc(af) for af in dae_dec_act_func]
    finetune_enc_act_func = [utilities.str2actfunc(af) for af in finetune_enc_act_func]
    finetune_dec_act_func = [utilities.str2actfunc(af) for af in finetune_dec_act_func]

    sdae = deep_autoencoder.DeepAutoencoder(
        do_pretrain=FLAGS.do_pretrain, name=FLAGS.name,
        layers=dae_layers, finetune_loss_func=FLAGS.finetune_loss_func,
        finetune_learning_rate=FLAGS.finetune_learning_rate, finetune_num_epochs=FLAGS.finetune_num_epochs,
        finetune_opt=FLAGS.finetune_opt, finetune_batch_size=FLAGS.finetune_batch_size,
        finetune_dropout=FLAGS.finetune_dropout,
        enc_act_func=dae_enc_act_func, dec_act_func=dae_dec_act_func,
        corr_type=dae_corr_type, corr_frac=dae_corr_frac, regcoef=dae_regcoef,
        loss_func=dae_loss_func,
        opt=dae_opt, tied_weights=FLAGS.tied_weights,
        learning_rate=dae_learning_rate, momentum=FLAGS.momentum,
        num_epochs=dae_num_epochs, batch_size=dae_batch_size,
        finetune_enc_act_func=finetune_enc_act_func, finetune_dec_act_func=finetune_dec_act_func)


    def load_params_npz(npzfilepath):
        params = []
        npzfile = np.load(npzfilepath)
        for f in npzfile.files:
            params.append(npzfile[f])
        return params


    encodingw = None
    encodingb = None

    # Fit the model (unsupervised pretraining)
    if FLAGS.do_pretrain:
        encoded_X, encoded_vX = sdae.pretrain(trX, vlX)
    # print(trX.shape)
    # Supervised finetuning
    # sdae.fit(trX, trRef, vlX, vlRef)
    # load the model
    # recons_arr=sdae.load_reconstruction_arr(trX)
    encode_arr = sdae.load_encode_arr(test_single_sampel_arr)
    # Compute the reconstruction loss of the model

    #print('Test set reconstruction loss: {}'.format(sdae.score(teX, teRef)));


    # print('recons_arr shape', recons_arr.shape)
    # print('encode_arr shape', encode_arr.shape)
    # for row in test_arr:
    #     print(row)encode/encode_train_malware_arr', encode_arr)
        # np.save(encode_dir,sdae.transform(teX))
    # Save the predictions of the model
    # print('Saving the encode for the benign data')
    # np.save('E:/Pycharm Projects/Deep-Learning-TensorFlow/model/save_encode/10000_dae10_finetune20_api_noise0.1_noscale_rate0.05_malware_arr', sdae.transform(binary_malware_arr))
    # np.save('E:/Pycharm Projects/Deep-Learning-TensorFlow/model/save_encode/10000_dae10_finetune20_api_noise0.1_noscale_rate0.05_benign_arr',sdae.transform(binary_benign_arr))

    #save the Son result
    if encode_dir:
        #print('Saving the signature for the Son samples...')
        np.save(encode_dir,encode_arr)
        test_arr=np.load(encode_dir)
        print('load keras model...')
        print('predict class:',load_keras_model(test_arr))
        exit()
    #     np.save('E:/Pycharm Projects/Deep-Learning-TensorFlow/model/save_encode/encode_train_arr',sdae.transform(trX))
    #     np.save('E:/Pycharm Projects/Deep-Learning-TensorFlow/model/save_encode/encode_varify_arr',sdae.transform(vlX))

    if not FLAGS.save_reconstructions:
        print('Saving the reconstructions for the Son set...')
        np.save(FLAGS.save_reconstructions, sdae.reconstruct(teX))


    def save_layers_output(which_set):
        if which_set == 'train':
            trout = sdae.get_layers_output(trX)
            signature = trout[2]
            for i, o in enumerate(trout):
                np.save(FLAGS.save_layers_output_train + '-layer-' + str(i + 1) + '-train', o)

        elif which_set == 'Son':
            teout = sdae.get_layers_output(teX)
            for i, o in enumerate(teout):
                np.save(FLAGS.save_layers_output_test + '-layer-' + str(i + 1) + '-Son', o)


    # Save output from each layer of the model
    if FLAGS.save_layers_output_test:
        print('Saving the output of each layer for the Son set')
        save_layers_output('Son')

    # Save output from each layer of the model
    if not FLAGS.save_layers_output_train:
        print('Saving the output of each layer for the train set')
        save_layers_output('train')

    # Save parameters of the model
    if FLAGS.save_model_parameters:
        print('Saving the parameters of the model')
        param_dir = FLAGS.save_model_parameters
        model_params = sdae.get_parameters(
            {
                'enc-weights-layer': sdae.encoding_w_,
                'enc-biases-layer': sdae.encoding_b_,
                'dec-weights-layer': sdae.decoding_w,
                'dec-biases-layer': sdae.decoding_b
            })
        for p in model_params:
            np.save(os.path.join(param_dir, p), model_params[p])
