#Import all required libraries
import tensorflow as tf
from tensorflow.keras import datasets,layers,models
import matplotlib.pyplot as plt
import numpy as np
import os
import PIL
import PIL.Image


#set a random seed for reproducability
seed= 42
np.random.seed(seed)

#splitting the image data in to train and test split
import splitfolders
splitfolders.ratio("/home/sowmya/malware_classification/data/images",output="/home/sowmya/malware_classification/data/images/split_images",seed=294,ratio=(0.8,0.2))

Path for datasets
train_dataset="/malware_classification/data/images/split_images/train/"
test_dataset="/malware_classification/data/images/split_images/val"

image_width = 256
image_height = None
if tf.keras.backend.image_data_format() == 'channels_first':
    input_shape = (1,image_width,None)
else:
    input_shape = (image_width,None,1)

print(input_shape)

#Model Architecture
model = tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(16,(3,3),activation='relu',padding="same",input_shape=(256,None,1)),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Conv2D(32,(3,3),activation='relu',padding="same"),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Conv2D(64,(3,3),activation='relu'),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Conv2D(64,(3,3),activation='relu'),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Conv2D(64,(3,3),activation='relu'),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.GlobalMaxPool2D(),
    tf.keras.layers.Dense(128,activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(1,activation='sigmoid')
])


model.summary()

model.comfrom tensorflow.keras.optimizers import RMSprop
model.compile(optimizer = RMSprop(lr=0.001),loss='binary_crossentropy',metrics=['accuracy'])
print("CNN model complied")
#rescaling the data to feed the images from directories
from tensorflow.keras.preprocessing.image import ImageDataGenerator
train_datagen=ImageDataGenerator(rescale=1./255)
test_datagen=ImageDataGenerator(rescale=1./255)

#creating the data generators for traing and testing datasets
train_generator=train_datagen.flow_from_directory(
    train_dataset,
    target_size=(256,256),
    color_mode="grayscale",
    class_mode='binary',
    batch_size=128
)

test_generator=test_datagen.flow_from_directory(
    test_dataset,
    target_size=(256,256),
    color_mode="grayscale",
    class_mode='binary',
    batch_size=32
)

import time
start_time=time.time()
history=model.fit(train_generator,steps_per_epoch=8,epochs=100,verbose=1,validation_data = test_generator,validation_steps=8)
end_time=time.time()-start_time
print(end_time)

import matplotlib.pyplot as plt
plt.plot(history.history['accuracy'], label='train')
plt.plot(history.history['val_accuracy'], label='test')
plt.title('model accuracy')
plt.ylabel('accuracy')
plt.xlabel('epoch')
plt.legend()
# # Save the accuracy plot
#plt.savefig("CNN_model_2-acc.png")
plt.show() # Display  the plot

plt.plot(history.history['loss'], label='train')
plt.plot(history.history['val_loss'], label='test')
plt.title('model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend()
# # Save the accuracy plot
#plt.savefig("CNN_model_2-acc.png")
plt.show() # Display  the plot


