import warnings
warnings.filterwarnings('ignore')
import tensorflow as tf
#tf.compat.v1.disable_eager_execution()
import numpy as np
from matplotlib import pyplot as plt
#from art import *
#from art.art.attacks.evasion import fast_gradient.py
#(x_train, y_train), (x_test, y_test), min_pixel_value, max_pixel_value = load_mnist()

from art.estimators.classification import KerasClassifier
from art.attacks.evasion import FastGradientMethod, CarliniLInfMethod
import sys
import warnings
import os
from math import log
import cv2
import numpy as np
import scipy as sp
#import pydot
#import pydotplus
#from keras.utils import plot_model
#import graphviz
from PIL import Image
Image.MAX_IMAGE_PIXELS = None
warnings.simplefilter('ignore', Image.DecompressionBombWarning)
warnings.simplefilter('error', Image.DecompressionBombWarning)
#tf.set_random_seed(1234)
import matplotlib.pyplot as plt
import keras
#Load dataset
path_root = 'mal_dataset\\'
from keras.preprocessing.image import ImageDataGenerator
batches = ImageDataGenerator().flow_from_directory(directory=path_root, target_size=(224,224),color_mode="grayscale",batch_size=12000)
batches.class_indices
imgs, labels = next(batches)
#Split Data
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(imgs,labels, test_size=0.2)
#X_test =X_test[0:200]
X_test =X_test[0:200]
y_test= y_test[0:200]
X_train, X_test = X_train / 255.0, X_test / 255.0
print(X_train.shape)
print(X_test.shape)
print(y_train.shape)
print(y_test.shape)
#Create model
from keras.models import Sequential
from keras.optimizers import SGD
from keras.utils import np_utils
from keras.models import Model
from keras.layers import Input, Dense
from keras.models import Model
from keras.layers import Dropout, Flatten, Dense, Conv1D, Conv2D, Activation, MaxPooling2D, GlobalMaxPooling1D, GlobalMaxPooling2D, Input, Embedding, Multiply, MaxPooling1D,GlobalAveragePooling2D, GlobalAveragePooling1D

from keras.utils import plot_model
from keras import metrics
from keras.layers.advanced_activations import PReLU
from keras.layers.normalization import BatchNormalization
import keras
from keras import regularizers
from keras import backend
sess = tf.compat.v1.Session()
keras.backend.set_session(sess)
num_classes = 2
def malware_model():
    Malware_model = Sequential()
    Malware_model.add(keras.layers.Conv2D(30, kernel_size=(3, 3),
                     activation='relu',
                     input_shape=(224,224,1)))

    Malware_model.add(keras.layers.MaxPooling2D(pool_size=(2, 2)))
    Malware_model.add(keras.layers.Conv2D(15, (3, 3), activation='relu'))
    Malware_model.add(keras.layers.MaxPooling2D(pool_size=(2, 2)))
    Malware_model.add(keras.layers.Flatten())
    Malware_model.add(keras.layers.Dense(128, activation='relu'))
    #Malware_model.add(keras.layers.Dropout(0.2))
    #Malware_model.add(keras.layers.Dense(50, activation='relu'))
    Malware_model.add(keras.layers.Dense(num_classes, activation='softmax'))
    Malware_model.compile(loss='binary_crossentropy', optimizer=keras.optimizers.SGD(lr=0.01,momentum=0.9,nesterov=True,decay=1e-6), metrics=['accuracy'])
    return Malware_model
Malware_model = malware_model()
Malware_model.summary()
Malware_model.fit(X_train, y_train,epochs=10, validation_data=(X_test, y_test),batch_size=100)
Malware_model.save("final_model.h5")
loss_test, accuracy_t = Malware_model.evaluate(X_test, y_test)
print('Accuracy on test data: {:4.2f}%'.format(accuracy_t * 100))
classifier = KerasClassifier(model=Malware_model, clip_values=(0, 1))
# Fast Gradient Sign Method attack on baseline classifier
attack_fgsm = FastGradientMethod(estimator=classifier, eps=0.015)
x_test_adv_fgsm = attack_fgsm.generate(X_test)
loss_test, accuracy_test = Malware_model.evaluate(x_test_adv_fgsm, y_test)
perturbation = np.mean(np.abs((x_test_adv_fgsm - X_test)))
print('Accuracy on adversarial test data: {:4.2f}%'.format(accuracy_test * 100))
print('Average perturbation: {:4.2f}'.format(perturbation))
#fig = plt.figure(figsize=(50, 10))
#cols = 10

images = list(x_test_adv_fgsm[:200])
images_ori = list(X_test[:200])

for i in range(200):
    #adv_filename = 'adv' + i + '.bmp'
    #ori_filename = 'ori' + i + '.bmp'
    cv2.imwrite('data_eps_0.015/original/original_{:03d}.bmp'.format(i+1), images_ori[i]*255.0)
    cv2.imwrite('data_eps_0.015/adversarial/adversarial_{:03d}.bmp'.format(i+1), images[i]*255.0)

#cv2.imwrite("adv_images/{}/adv_img_{}.png".format("malware", i+1000), x_adv_norm)
#plt.imshow(images[0], cmap='gray')
#plt.show()
#plt.imshow(images_ori[0], cmap='gray')
#plt.show()
#for i in range(0, len(images)):
#    ax = fig.add_subplot((i/5)+1, cols, (i%5)+1)
#    ax.tick_params(axis='both', which='both',
#                   bottom=False, top=False,
#                   right=False, left=False,
#                   labelbottom=False, labelleft=False)
#    plt.imshow(images[i], cmap='gray')
#plt.show()
