#Importing Libraries
import sys
import warnings
import os
from math import log
import numpy as np
import scipy as sp
from PIL import Image
Image.MAX_IMAGE_PIXELS = None
warnings.simplefilter('ignore', Image.DecompressionBombWarning)
warnings.simplefilter('error', Image.DecompressionBombWarning)
import matplotlib.pyplot as plt

#Load dataset
path_root = 'mal_dataset\\'
from keras.preprocessing.image import ImageDataGenerator
batches = ImageDataGenerator().flow_from_directory(directory=path_root, target_size=(224,224),batch_size=12000)
batches.class_indices
imgs, labels = next(batches)
print(imgs.shape)
print(labels.shape)

#Split Data
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(imgs/255.,labels, test_size=0.2)
print(X_train.shape)
print(X_test.shape)
print(y_train.shape)
print(y_test.shape)

import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True  # dynamically grow the memory used on the GPU

sess = tf.Session(config=config)
set_session(sess)  # set this TensorFlow session as the default session for Keras
from keras.models import Sequential
from keras.optimizers import SGD
from keras.utils import np_utils

from keras.layers import Input, Dense
from keras.layers import Dense, Dropout, Flatten
from keras.models import Model
from keras.layers import Dropout, Dense, Conv1D, Flatten,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 import regularizers
from keras import backend as K
K.set_image_dim_ordering('tf')

#Create model

num_classes = 2
def malware_model():
    Malware_model = Sequential()
    Malware_model.add(Conv2D(30, kernel_size=(3, 3),
                     activation='relu',
                     input_shape=(224,224,3)))

    Malware_model.add(MaxPooling2D(pool_size=(2, 2)))
    Malware_model.add(Conv2D(15, (3, 3), activation='relu'))
    Malware_model.add(MaxPooling2D(pool_size=(2, 2)))
    Malware_model.add(Dropout(0.25))
    Malware_model.add(Flatten())
    Malware_model.add(Dense(128, activation='relu'))
    Malware_model.add(Dropout(0.5))
    Malware_model.add(Dense(50, activation='relu'))
    Malware_model.add(Dense(num_classes, activation='softmax'))
    Malware_model.compile(loss='binary_crossentropy', optimizer = SGD(lr=0.01,momentum=0.9,nesterov=True,decay=1e-6), metrics=['accuracy'])
    return Malware_model
Malware_model = malware_model()
Malware_model.summary()

seed = 7
np.random.seed(seed)

nb_epoch = 50
from keras.callbacks import EarlyStopping
es = EarlyStopping(monitor='val_loss', mode='min', verbose=1, patience=10)
K.set_learning_phase(1)
Malware_model.fit(X_train, y_train, epochs=50, validation_data=(X_test, y_test), batch_size=100,callbacks=[es])
#Evaluate model
from sklearn import metrics
from sklearn.metrics import accuracy_score
test_acc = Malware_model.evaluate(X_test, y_test)
train_acc = Malware_model.evaluate(X_train, y_train)
Malware_model.save("detection_model_1.h5")
print('Final CNN test_accuracy: ', test_acc[1])
print('Final CNN train_accuracy: ', train_acc[1])
#y_pred = Malware_model.predict_classes(X_test, verbose=0)
y_pred = np.argmax(Malware_model.predict(X_test), axis=-1)
y_test2 = np.argmax(y_test, axis=1)
c_matrix = metrics.confusion_matrix(y_test2, y_pred)
report = metrics.classification_report(y_test2, y_pred)
acc = metrics.accuracy_score(y_test2, y_pred)
print(acc)
print(report)
print(c_matrix)