#I love Alex

# To run the code type the following into the command prompt 
# in the EC_521_ML folder type python test1.py. 
import keras as keras
import time
import datetime
from keras.models import Sequential
from keras.layers import Dense
import numpy as np 
import random 
import matplotlib.pyplot as plt 
import matplotlib

seed = 7
np.random.seed(seed)

# *************************************************************************
#**************** This part of the code formats the data ******************
#**************************************************************************
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+ "02" + '.png')
# This is where we will load the data-set 
# Assuming we will load from csv... 
# dataset = np.loadtxt(...)
arrayDataLength = 133

theData = np.load("nnData.npy")
theLabels = np.load("nnLabels.npy")


zeusD = np.array([np.zeros(133)])
confickerD = np.array([np.zeros(133)])
dridexD = np.array([np.zeros(133)])
zeus = np.array([])
conficker = np.array([])
dridex = np.array([])
index = 0
for i in theLabels: 
	if i == 0:
		zeusD = np.concatenate((zeusD, [theData[index]]), axis=0)
		# print(i)
		zeus = np.concatenate((zeus, [i]), axis=0)
	if i == 1: 
		confickerD = np.concatenate((confickerD, [theData[index]]), axis=0)
		conficker = np.concatenate((conficker, [i]), axis=0)
	if i == 2: 
		# print(i)
		dridexD = np.concatenate((dridexD, [theData[index]]), axis=0)
		dridex = np.concatenate((dridex, [i]), axis=0) 
	index += 1


zeus = zeus[:len(dridex)]
conficker = conficker[:len(dridex)] 


zeusD = zeusD[:len(dridex)]  
confickerD = confickerD[:len(dridex)] 
print("***************************************************************************")
print(len(zeusD))
print(len(confickerD))
print(len(dridexD))

indices = random.sample(range(0,len(dridex)), int(len(dridex)*0.5))

trainZeus = zeusD[indices]
trainConficker = confickerD[indices]
trainDridex = dridexD[indices]
print("***************************************************************************")
print(len(trainZeus))
print(len(trainConficker))
print(len(trainDridex))

tLZ = zeus[indices]
tLC = conficker[indices]
tLD = dridex[indices]

trainData = np.concatenate((trainZeus, trainConficker, trainDridex), axis=0)
trainLabels = np.concatenate((tLZ, tLC, tLD), axis=0)

classType = [0,0,0]
for i in trainLabels:
	if i == 0:
		classType[0] += 1
	if i == 1:
		classType[1] += 1
	if i == 2:
		classType[2] += 1
print(classType[0])
print(classType[1])
print(classType[2])

tempArr = list(range(len(tLZ)))
for i in indices: 
	if i in tempArr:
		tempArr[i] = 0

tI = list(filter((0).__ne__, tempArr))

testZeus = zeusD[tI]
testConficker = confickerD[tI]
testDridex = dridexD[tI] 

tLZ = zeus[tI]
tLC = conficker[tI]
tLD = dridex[tI] 



testData = np.concatenate((testZeus, testConficker, testDridex), axis=0)
testLabels = np.concatenate((tLZ, tLC, tLD), axis=0)

classType = [0,0,0]
for i in testLabels:
	if i == 0:
		classType[0] += 1
	if i == 1:
		classType[1] += 1
	if i == 2:
		classType[2] += 1
print(classType[0])
print(classType[1])
print(classType[2])

trainLabels = keras.utils.to_categorical(trainLabels, num_classes=3)
testLabels = keras.utils.to_categorical(testLabels, num_classes=3)

# randomize 
# indices = random.sample(range(0,len(trainData)), int(len(dridex))

# theData = np.reshape(theData, (len(theLabels), arrayDataLength))
# Split the data set into training and testing here: 
# [[packet, packet, packet, packet, packet, packet label]
#  [ packet packet packet packet packet,    packet label]]
# trainData = theData[0:len(theData)-1000] # np.random.choice(theData, len(theData) -1000)
# count = 0
# trainData = np.array([])

# for data in theData:
#   trainData = np.append(trainData, np.append(data, theLabels[count]))
#   count += 1

# trainData = np.reshape(trainData, (len(theLabels), arrayDataLength+1))

# idxTrain = np.random.choice(len(trainData), len(trainData)-1000)
# idxTest = np.random.choice(len(trainData), 1000)


# testData = trainData[idxTest,:]
# trainData = trainData[idxTrain,:] 
# trainLabels = trainData[:, arrayDataLength:]
# trainLabels = keras.utils.to_categorical(trainLabels, num_classes=3)
# trainData = trainData[:,0:arrayDataLength]
# testLabels = testData[:,arrayDataLength:]
# testData = testData[:,0:arrayDataLength]

# #print(trainLabels)
# #test = dataset[:]

# #************************************************************************
# #************************************************************************

# #************************************************************************
# #************** This part of the code makes the CNN *********************
# #************************************************************************

# # Create the model here: 

# 	pass
neuralNetModel = Sequential()
numInputs = arrayDataLength 
# # neuralNetModel = load_model('weight1.h5')
# # This should create a sequential model that accepts ? inputs and
# # has a signle hidden layer with 1000 nodes. The dense command makes it a 
# # fully connected graph. We have 1 output node one for our first model
# # that does binary classification. 


# # ******************* Instructions ***************************
# # The following are the things you can tune within the model: 
# # nH1 is equal to the nodes in the hidden layer 1.
# #testNumber = 0 
# #activationList = ["sigmoid","relu","softmax"]
# nH1 = 0
# # logFile = open("Test_result2s.txt","a+")
# progStartTime = str(datetime.datetime.now())
# # logFile.write("Test started at %s \r\n" % progStartTime)
# # for activation_type in activationList:
# # 	logFile.write("\r\n")
# # 	logFile.write("------------------------------------------------------ \r\n")
# # 	logFile.write("Changing Activation to %s \r\n" % activation_type)
# # 	logFile.write("------------------------------------------------------ \r\n")
# # 	logFile.write("\r\n")
startTime = time.time()
# # Too many nodes will overfit the training set. Too little nodes and you wont 
# # be capturing all the features. 
# # The more nodes the longer it will take to run. So be forewarned. 
# # ****************** End Instruction ****************************
# neuralNetModel.add(Dense(nH1, input_dim=numInputs, activation='relu')) # Tested different activation function
neuralNetModel = Sequential()
numInputs = arrayDataLength 
activation_func = "sigmoid"
neuralNetModel.add(Dense(1000,input_dim=numInputs, activation=activation_func))
# # neuralNetModel.add(Dense(nH1, input_dim=numInputs, activation='softmax')) # Tested different activation function
# # 

neuralNetModel.add(Dense(3, activation=activation_func)) 
# 	# This command compiles the model
neuralNetModel.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])


# # ***************** Instruction ************************** 
# # The following are the things you can change in the neural net
# # The number of Epochs of training. The more epochs the more you overfit the 
# # the training set. Too little epochs and you will not capture enough features. 
# # sizeBatch is how many training values does the model train on at once. 
# # the batch size is always less then the actual number of data. The smaller the 
# # batchSize the fewer features you are training against at once. Speeds up or slows down 
# # your model 
numEpochs = 10
sizeBatch = 100
# # ***************** End Instruction **************************
# # The following command will train the model on some data: 
neuralNetModel.fit(trainData, trainLabels, epochs=numEpochs, batch_size=sizeBatch)

# # Find predictions on some test set 
# # predictions = model.predict(test) 
endTime = time.time()
# # ******************* Instruction **************************************
# # The metric you are going to try and increase is accuracy. 
# # If you can get it to 60 or 70 then we have succeeded and just save the weights. 
# # ******************** End Instruction ********************************
scores = neuralNetModel.evaluate(testData, testLabels)

print("\n%s: %.2f%%" % (neuralNetModel.metrics_names[1], scores[1]*100))
now = datetime.datetime.now()
filename ="Weights_"+ activation_func+(now.strftime("_%Y_%m_%d_%H_%M"))+".h5"
print("Saving Weights: ")
neuralNetModel.save_weights(filename)
print("Saved Weights")



progEndTime = (datetime.datetime.now())
progRunTime = endTime-startTime
print("Test Run Time was %s \r\n" %progRunTime) 

test_Y_hat = neuralNetModel.predict(testData, batch_size=100)
conf = np.zeros([3,3])
confnorm = np.zeros([3,3])
for i in range(0,testData.shape[0]):
    j = list(testLabels[i,:]).index(1)
    k = int(np.argmax(test_Y_hat[i,:]))
    conf[j,k] = conf[j,k] + 1
for i in range(0,3):
    confnorm[i,:] = conf[i,:] / np.sum(conf[i,:])
the_classes = ['zeus', 'conficker', 'dridex']
plot_confusion_matrix(confnorm, labels=the_classes)

	