from __future__ import print_function
import keras
from keras.datasets import mnist
from keras.layers import Dense, Flatten
from keras.layers import Conv2D, MaxPooling2D
from keras.models import Sequential
import matplotlib.pylab as plt
import numpy as np
from keras.layers.core import Reshape,Dense,Dropout,Activation,Flatten



# In order to correctly run this we are going to need a lot more data! 
seed = 7
np.random.seed(seed)
batch_size = 50
num_classes = 4
# epochs = 30

# input image dimensions
# **********************************************************************
#*****************  INSTRUCTION ****************************************
# the variable img_x changes the number of rows in our matrix, the 
# greater the number the larger a single matrix but the fewer samples 
# we will have to train with. Ideally we should have like 5000 samples 
# with a matrix size of 50, 128 or 133. So we need more data ! 
#***********************************************************************
img_x, img_y = 20, 133

# Load the data into the model
# theData = np.load("cnnData.npy")
# theLabels = np.load("cnnLabels.npy")
# theData = np.reshape(theData, (len(theLabels), img_y))

# # Now format the data so the last element of each array is the label. 
# index = 0 
# tempData = np.array([]) 
# for data in theData: 
#   tempData = np.append(tempData, np.append(data, theLabels[index]))
#   index += 1

# tempData = np.reshape(tempData, (len(theData), img_y+1))

# print(len(tempData))

# # We are going to make matrixes with a height of 10 to do this we are going to ignore some packets...
# diff = len(tempData)%img_x 
# tempData = tempData[:len(tempData)-diff, :]
# print(len(tempData))

# # Now I need to make matrices with the same labels. 

# matrixData = np.array([np.zeros((img_x,img_y+1))])
# tempMatrix = np.array([np.zeros((img_x,img_y+1))]) 
# tempZeros = np.zeros(img_y)

# dataZero = tempData[0] 
# aLabel = dataZero[img_y]
# index = 0
# for data in tempData: 
#   # find the length of the current temporary matrix 
#   if index != img_x: 
#     # Check if the label matches the current labels we are looking at. 
#     if data[img_y] == aLabel: 
#       tempMatrix[0][index] = data
#       index += 1
#     else:
#       distFromTen = len(tempMatrix) % img_x
#       for i in range(0,distFromTen):
#         tempMatrix[0][index] = np.append(tempZeros, aLabel)
#         index += 1
#   elif index == img_x: 
#     matrixData = np.concatenate((matrixData, tempMatrix), axis=0)
#     tempMatrix = np.array([np.zeros((img_x,img_y+1))])
#     aLabel = data[img_y] 
#     index = 0
#     tempMatrix[0][index] = data
#     index += 1


# idxTrain = np.random.choice(len(matrixData), len(matrixData))
# matrixData = matrixData[idxTrain,] 

# # Now I want to grab the labels from the data before running the ConvNet
# randomLabels = matrixData[:, :, img_y:]
# matrixLabels = np.array([]) 

# for subMatrix in randomLabels:
#   matrixLabels = np.append(matrixLabels, subMatrix[0][0])

# #print(matrixLabels)

# # Now i want to convert the trainData to one-hot arrays
# oneHotLabels = keras.utils.to_categorical(matrixLabels, num_classes)


# # Make the test and training sets 
# justData = matrixData[:, :, 0:img_y]
# offset = int(len(justData) - len(justData)*0.2)
# trainData = justData[0:offset]
# trainData = trainData.reshape(trainData.shape[0], img_x, img_y, 1)
# trainLabels = oneHotLabels[0:offset]
# testData = justData[offset:]
# testData = testData.reshape(testData.shape[0], img_x, img_y, 1)
# testLabels = oneHotLabels[offset:]

input_shape = (img_x, img_y,1)

# # np.save("cnnTrainData", trainData)
# # np.save("cnnTestData", testData)
# # np.save("cnnTrainLabels", trainLabels)
# # np.save("cnnTestLabels", testLabels)

# trainData = np.load("cnnTrainData.npy")
# testData = np.load("cnnTestData.npy")
# trainLabels = np.load("cnnTrainLabels.npy")
# testLabels = np.load("cnnTestLabels.npy")
trainData = np.load("cnnTrainDataEqual.npy")
testData = np.load("cnnTestDataEqual.npy")
trainLabels = np.load("cnnTrainLabelsEqual.npy")
testLabels = np.load("cnnTestLabelsEqual.npy")
# # trainLabel = np.reshape(trainLabels (len(trainData), 3))
trainLabels = trainLabels[1:]
testLabels = testLabels[1:]
print(len(trainLabels))
for i in testLabels: 
  print(np.argmax(i))
# # print(trainData.shape)
print(trainLabels.shape)
print(trainLabels[0, :])
# classType = [0,0,0] 
# index = 0
# zeusD = np.array([np.zeros((20,133))])
# confickerD = np.array([np.zeros((20,133))])
# dridexD = np.array([np.zeros((20,133))])
# confickerD = confickerD[:,:,:,np.newaxis]
# zeusD = zeusD[:,:,:,np.newaxis]
# dridexD = dridexD[:,:,:,np.newaxis]

# zeus = np.array([np.zeros(4)])
# conficker = np.array([np.zeros(4)])
# dridex = np.array([np.zeros(4)])

# for label in trainLabels: 
#   if np.argmax(label) == 0:
#     classType[0] += 1
#     # print(label)
#     # print(zeus.shape)
#     zeus = np.concatenate((zeus, [label]), axis=0)
#     # print(trainData[index].shape)
#     zeusD = np.concatenate((zeusD, [trainData[index]]), axis=0)
#       # zeusD = np.concatenate((zeusD, [trainData[index]]), axis=0)

#   if np.argmax(label) == 1:
#     classType[1] += 1
#     conficker = np.concatenate((conficker, [label]), axis=0)
#     confickerD = np.concatenate((confickerD, [trainData[index]]), axis=0)
#   if np.argmax(label) == 2: 
#     classType[2] += 1
#     dridex = np.concatenate((dridex, [label]), axis=0)
#     dridexD = np.concatenate((dridexD, [trainData[index]]), axis=0)
#   # if np.argmax(label) == 3: 
#   #   classType[3] += 1
#   #   normal = np.concatenate((normal, label), axis=0)
#   #   normalD = np.concatenate((normalD, [trainData[index]]), axis=0)
#   index += 1

# theMin = np.argmin(classType)

# print(theMin) 
# print(min(classType))
# print("*******************************")
# zeusD = zeusD[0:min(classType)]
# confickerD = confickerD[0:min(classType)]
# dridexD = dridexD[0:min(classType)]

# zeus = zeus[0:min(classType)]
# conficker = conficker[0:min(classType)] 

# # # Prepared the training set. 
# trainData = np.concatenate((zeusD, confickerD, dridexD), axis=0)
# trainLabels = np.concatenate((zeus, conficker, dridex), axis=0)
# np.save("cnnTrainDataEqual", trainData)
# np.save("cnnTrainLabelsEqual", trainLabels)

# # Prepare the test set. 
# classType = [0,0,0] 
# index = 0
# zeusD = np.array([np.zeros((20,133))])
# confickerD = np.array([np.zeros((20,133))])
# dridexD = np.array([np.zeros((20,133))])
# confickerD = confickerD[:,:,:,np.newaxis]
# zeusD = zeusD[:,:,:,np.newaxis]
# dridexD = dridexD[:,:,:,np.newaxis]

# zeus = np.array([np.zeros(4)])
# conficker = np.array([np.zeros(4)])
# dridex = np.array([np.zeros(4)])

# for label in testLabels: 
#   if np.argmax(label) == 0:
#     classType[0] += 1
#     zeus = np.concatenate((zeus, [label]), axis=0)
#     # print(trainData[index].shape)
#     zeusD = np.concatenate((zeusD, [testData[index]]), axis=0)
#       # zeusD = np.concatenate((zeusD, [trainData[index]]), axis=0)

#   if np.argmax(label) == 1:
#     classType[1] += 1
#     conficker = np.concatenate((conficker, [label]), axis=0)
#     confickerD = np.concatenate((confickerD, [testData[index]]), axis=0)
#   if np.argmax(label) == 2: 
#     classType[2] += 1
#     dridex = np.concatenate((dridex, [label]), axis=0)
#     dridexD = np.concatenate((dridexD, [testData[index]]), axis=0)
# #   # if np.argmax(label) == 3: 
# #   #   classType[3] += 1
# #   #   normal = np.concatenate((normal, label), axis=0)
# #   #   normalD = np.concatenate((normalD, [trainData[index]]), axis=0)
# #   index += 1

# theMin = np.argmin(classType)

# print(theMin) 
# print(min(classType))
# print("*******************************")
# zeusD = zeusD[0:min(classType)]
# confickerD = confickerD[0:min(classType)]
# dridexD = dridexD[0:min(classType)]

# zeus = zeus[0:min(classType)]
# conficker = conficker[0:min(classType)] 

# # #*****************************************************************************
# testData = np.concatenate((zeusD, confickerD, dridexD), axis=0)
# testLabels = np.concatenate((zeus, conficker, dridex), axis=0)
# np.save("cnnTestDataEqual", testData)
# np.save("cnnTestLabelsEqual", testLabels)

#******************************************************************************
#******************************************************************************
# ************************ The following is the Model *************************
#******************************************************************************

model = Sequential()

#******************************************************************************
#*************** This makes the convolution layer of the model ****************
# Model Parameters: 
# Kernel Size: The kernel size divides the size of the matrix... So if your 
# matrix dimensions are 10,128 and you have a kernel of 2 then the outcome will
# be 5,64. 
# Strides: how quickly the kernel filter moves around the original image. 1,1
# is the best. The stride should never be more than the kernel_size dimensions.
# numK1 = the number of kernels in this convolution layer
#******************************************************************************
#***************************** Instruction ************************************
# If you change the kernel size or the maxpool size you have to make sure that 
# they don't go past the dimensions of the actuall matrix. Look at explanation 
# above ^^^^^ 
epochs = 50
dr = 0.9
numK1 = 32
model.add(Conv2D(numK1, kernel_size=2, strides=(1, 1),
                 activation='relu',
                 input_shape=input_shape,
                 data_format="channels_last"))
#******************************************************************************
#***************** This makes the pooling layer********************************
#pool_size: chooses the size of the pooling filter. 2,2 is the smallest, whil 
#anything larger just filter the data even more. 
#strides: Should always be equal to the pool_size. 
#******************************************************************************
model.add(Conv2D(numK1, kernel_size=(2, 2), strides=(2, 2), activation='relu'))
#******************************************************************************
#***************** This makes the second convolution layer ********************
# We don't need this other convolution layer unless we intend to make the 
# matrix dimensions bigger, Please look at explanation above regarding 
# adding convolution layer. For each one of these layers you are making the 
# Matrix smaller.... So you will have to worry about dimensionality issues. 
#******************************************************************************
numK2 = 64
model.add(Conv2D(numK2, kernel_size=2, strides=(1,1), activation='relu'))
model.add(Conv2D(numK2, kernel_size=2, strides=(1,1), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))
#******************************************************************************
#This command will flatten the model so that it is prepared for the fully
#connected layer. 
#******************************************************************************
model.add(Flatten())

#******************************************************************************
#*********************** This trains the fully connected layer ****************
# fNodes: is the number of nodes in the fully connected layer 
# *****************************************************************************
#***************** INSTRUCTIONS ***********************************************
# You can change the number of nodes in the fully connected layer. 
fNodes = 1000
model.add(Dense(fNodes, activation='relu'))
fNodes2 = 2000 
model.add(Dense(fNodes2, activation='sigmoid'))
model.add(Dropout(dr))
fNodes3 = 2000 
model.add(Dense(fNodes3, activation='sigmoid'))
fNodes4 = 1000
model.add(Dense(fNodes4, activation='relu'))


#******************************************************************************
#*********************** Trains the output layer ******************************
model.add(Dense(num_classes, activation='softmax'))

#*****************************************************************************
#*****************************************************************************
#*****************************************************************************
#************************** Runs the Model ***********************************

model.compile(loss=keras.losses.categorical_crossentropy,
              optimizer=keras.optimizers.Adam(),
              metrics=['accuracy'])


class AccuracyHistory(keras.callbacks.Callback):
    def on_train_begin(self, logs={}):
        self.acc = []

    def on_epoch_end(self, batch, logs={}):
        self.acc.append(logs.get('acc'))

history = AccuracyHistory()

model.fit(trainData, trainLabels,
          batch_size=batch_size,
          epochs=epochs,
          verbose=1,
          validation_data=(testData, testLabels),
          callbacks=[history])
score = model.evaluate(testData, testLabels, verbose=0)
print('Test loss:', score[0])
print('Test accuracy:', score[1])
plt.plot(range(1, epochs+1), history.acc)
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
# plt.show()
plt.savefig('plot_numK1_'+str(numK1)+'_epochs_'+str(epochs)+'.png')
# # del plt
def plot_confusion_matrix(cm, title='Confusion matrix', cmap=plt.cm.Blues, labels=[]):
  # matplotlib.use('QT4Agg')
  plt.imshow(cm, interpolation='nearest', cmap=cmap)
  plt.title(title)
  plt.colorbar()
  tick_marks = np.arange(len(labels))
  plt.xticks(tick_marks, labels, rotation=45)
  plt.yticks(tick_marks, labels)
  plt.tight_layout()
  plt.ylabel('True label')
  plt.xlabel('Predicted label')
  # plt.safefig(title, dpi=NONe)
  plt.savefig(title+"conv"+ " " + "05" + '.png')
test_Y_hat = model.predict(testData, batch_size=10)
conf = np.zeros([3,3])
confnorm = np.zeros([3,3])
for i in range(0,testLabels.shape[0]):
    # print(trainLabels[i,:])
    if(1 in testLabels[i,:]):
      j = list(trainLabels[i,:]).index(1)
      k = int(np.argmax(test_Y_hat[i,:]))
      conf[j,k] = conf[j,k] + 1
for i in range(0,3):
    print(conf[i,:])
    confnorm[i,:] = conf[i,:] / np.sum(conf[i,:])
the_classes = ['zeus', 'conficker', 'dridex']
plot_confusion_matrix(confnorm, labels=the_classes)