import os

import arff
import numpy as np
from matplotlib import pyplot as plt
from keras.layers import Dense, Dropout
from keras.models import Sequential
from keras.callbacks import TensorBoard
from keras.utils import to_categorical
from keras.utils import plot_model

# set the directory of the dataset
file = open("data/final-dataset.arff", 'r')


def generate_model(shape):

    # define the model
    model = Sequential()

    model.add(Dense(22, input_dim=shape, kernel_initializer='uniform', activation='relu'))
    model.add(Dropout(0.4))
    model.add(Dense(64, activation='relu'))
    model.add(Dropout(0.4))
    model.add(Dense(128, activation='relu'))
    model.add(Dropout(0.4))
    model.add(Dense(64, activation='relu'))
    model.add(Dropout(0.4))
    model.add(Dense(5, activation='softmax'))
    print(model.summary())

    return model


def scrape_data():
    # decode the .arff data and change text labels into numerical
    decoder = arff.ArffDecoder()
    data = decoder.decode(file, encode_nominal=True)

    # split the raw data into data and labels
    vals = [val[0: -1] for val in data['data']]
    labels = [label[-1] for label in data['data']]

    # split the labels and data into traning and validation sets
    training_data = vals[0: int(.9 * len(vals))]
    training_labels = labels[0: int(.9 * len(vals))]
    validation_data = vals[int(.9 * len(vals)):]
    validation_labels = labels[int(.9 * len(vals)):]

    # flatten labels with one hot encoding
    training_labels = to_categorical(training_labels, 5)
    validation_labels = to_categorical(validation_labels, 5)

    # save all arrays with numpy
    np.save('saved-files/vals', np.asarray(vals))
    np.save('saved-files/labels', np.asarray(labels))
    np.save('saved-files/training_data', np.asarray(training_data))
    np.save('saved-files/validation_data', np.asarray(validation_data))
    np.save('saved-files/training_labels', np.asarray(training_labels))
    np.save('saved-files/validation_labels', np.asarray(validation_labels))


# check to see if saved data exists, if not then create the data
if not os.path.exists('saved-files/training_data.npy') or not os.path.exists(
        'saved-files/training_labels.npy') or not os.path.exists(
        'saved-files/validation_data.npy') or not os.path.exists('saved-files/validation_labels.npy'):
    print('creating')
    if not os.path.exists('saved-files'):
        os.mkdir('saved-files')
    scrape_data()

# load the saved data
data_train = np.load('saved-files/training_data.npy')
label_train = np.load('saved-files/training_labels.npy')
data_eval = np.load('saved-files/validation_data.npy')
label_eval = np.load('saved-files/validation_labels.npy')

# generate and compile the model
model = generate_model(len(data_train[0]))
model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

# initialize tensorboard
tensorboard = TensorBoard(log_dir='logs/', histogram_freq=0, write_graph=True, write_images=True)

# only using 2 epochs otherwise the model would overfit to the data
history = model.fit(data_train, label_train, validation_data=(data_eval, label_eval), epochs=2, callbacks=[tensorboard])

# evaluating the model's performace
print(model.evaluate(data_eval, label_eval))

plot_model(model, to_file='model.png', show_shapes=True)

plt.figure(1)

# summarize history for accuracy
plt.subplot(211)
plt.plot(history.history['acc'])
plt.plot(history.history['val_acc'])
plt.title('model accuracy')
plt.ylabel('accuracy')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='upper left')

# summarize history for loss
plt.subplot(212)
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='upper left')
plt.show()

# save the model for later so no retraining is needed
model.save('saved-files/model.h5')

# play sound when done with code to alert me
os.system('afplay /System/Library/Sounds/Ping.aiff')
os.system('afplay /System/Library/Sounds/Ping.aiff')

@relation 10-7-dataset

@attribute SRC_ADD numeric
@attribute DES_ADD numeric
@attribute PKT_ID numeric
@attribute FROM_NODE numeric
@attribute TO_NODE numeric
@attribute PKT_TYPE {tcp,ack,cbr,ping}
@attribute PKT_SIZE numeric
@attribute FLAGS {-------,---A---}
@attribute FID numeric
@attribute SEQ_NUMBER numeric
@attribute NUMBER_OF_PKT numeric
@attribute NUMBER_OF_BYTE numeric
@attribute NODE_NAME_FROM {Switch1,Router,server1,router,clien-4,client-2,Switch2,client-5,clien-9,clien-2,clien-1,clien-14,clien-5,clien-11,clien-13,clien-0,switch1,client-4,clienthttp,clien-7,clien-19,client-14,clien-12,clien-8,clien-15,webserverlistin,client-18,client-1,switch2,clien-6,client-10,client-7,webcache,clien-10,client-15,clien-3,client-17,client-16,clien-17,clien-18,client-12,client-8,client-0,clien-16,client-13,client-11,client-6,client-3,client-9,client-19,http_client}
@attribute NODE_NAME_TO {Router,server1,Switch2,Switch1,clien-1,clien-5,clien-7,switch1,clien-11,clien-15,clien-13,clien-3,clien-9,clien-6,router,clien-4,clien-14,switch2,clien-8,clienthttp,webcache,clien-10,clien-12,webserverlistin,clien-0,clien-2,http_client,client-13,client-9,client-1,client-19,client-4,client-17,client-7,client-3,client-12,client-2,clien-18,client-16,clien-17,client-0,clien-16,client-18,client-5,client-11,client-14,client-8,client-6,client-10,clien-19,client-15}
@attribute PKT_IN numeric
@attribute PKT_OUT numeric
@attribute PKT_R numeric
@attribute PKT_DELAY_NODE numeric
@attribute PKT_RATE numeric
@attribute BYTE_RATE numeric
@attribute PKT_AVG_SIZE numeric
@attribute UTILIZATION numeric
@attribute PKT_DELAY numeric
@attribute PKT_SEND_TIME numeric
@attribute PKT_RESEVED_TIME numeric
@attribute FIRST_PKT_SENT numeric
@attribute LAST_PKT_RESEVED numeric
@attribute PKT_CLASS {Normal,UDP-Flood,Smurf,SIDDOS,HTTP-FLOOD}

@data
3,24.3,389693,21,23,tcp,1540,-------,4,11339,16091,24780100,Switch1,Router,35.529786,35.529786,35.539909,0,328.240918,505490,1540,0.236321,0,35.519662,35.550032,1,50.02192,Normal
15,24.15,201196,23,24,tcp,1540,-------,16,6274,16092,24781700,Router,server1,20.176725,20.176725,20.186848,0,328.205808,505437,1540,0.236337,0,20.156478,20.186848,1,50.030211,Normal
24.15,15,61905,23,22,ack,55,-------,16,1930,16092,885060,Router,Switch2,7.049955,7.049955,7.059958,0,328.206042,18051.3,55,0.008441,0,7.039952,7.069962,1.030045,50.060221,UDP-Flood
24.9,9,443135,23,21,ack,55,-------,10,12670,16085,884675,Router,Switch1,39.62797,39.62797,39.637973,0,328.064183,18043.5,55,0.008437,0,39.617967,39.647976,1.030058,50.060098,Normal
24.8,8,157335,23,21,ack,55,-------,9,4901,16088,884840,Router,Switch1,16.039806,16.039806,16.04981,0,328.113525,18046.2,55,0.008438,0,16.029803,16.059813,1.030054,50.061864,Normal
24.1,1,219350,21,1,ack,55,-------,2,6837,16091,885005,Switch1,clien-1,21.885768,21.885768,21.895771,0,328.297902,18056.4,55,0.00844,0,21.865762,21.895771,1.030016,50.043427,Normal
24.13,13,480053,24,23,ack,55,-------,14,13609,16103,885665,server1,Router,42.45032,42.45032,42.460323,0,328.460278,18065.3,55,0.008446,0,42.45032,42.48033,1.030032,50.055747,Normal
2.1,24.22,599411,23,24,cbr,1000,-------,23,4156,6250,6250000,Router,server1,58.26832,58.26848,58.27856,0.00016,124.943625,124944,1000,0.059605,0.00032,58.248,58.27856,25,75.02256,Normal
24.2,2,551227,24,23,ack,55,-------,3,15392,16091,885005,server1,Router,47.910078,47.910078,47.920081,0,328.26412,18054.5,55,0.00844,0,47.910078,47.940088,1.030022,50.048477,Normal
2,24.2,399941,21,23,tcp,1540,-------,3,11595,16091,24780100,Switch1,Router,36.314926,36.31548,36.325603,0.000554,328.26404,505526,1540,0.236321,0.001724,36.304803,36.336896,1,50.018467,Normal
9.2,24.51,33450,23,24,cbr,1500,-------,30,1279,9108,13662000,router,server1,32.180177,32.181257,32.191377,0.00108,1016.496869,1524750,1500,0.130291,0.03252,32.158857,32.191377,1,9.960185,UDP-Flood
24.11,11,356924,23,22,ack,55,-------,12,10505,16103,885665,Router,Switch2,33.015317,33.015317,33.02532,0,328.522947,18068.8,55,0.008446,0,33.005314,33.035323,1.030019,50.046382,Normal
24.5,5,349309,21,5,ack,55,-------,6,10306,16091,885005,Switch1,clien-5,32.442133,32.442133,32.452136,0,328.204851,18051.3,55,0.00844,0,32.422126,32.452136,1.030042,50.057349,Normal
6.1,24.26,629078,23,24,cbr,1000,-------,27,5629,6250,6250000,Router,server1,70.05264,70.05312,70.0632,0.00048,124.942027,124942,1000,0.059605,0.00096,70.032,70.0632,25,75.0232,Normal
24.5,5,565692,21,5,ack,55,-------,6,15761,16091,885005,Switch1,clien-5,49.037576,49.037576,49.047579,0,328.204851,18051.3,55,0.00844,0,49.01757,49.047579,1.030042,50.057349,Normal
24.1,1,478904,21,1,ack,55,-------,2,13574,16091,885005,Switch1,clien-1,42.382176,42.382176,42.392179,0,328.297902,18056.4,55,0.00844,0,42.36217,42.392179,1.030016,50.043427,Normal
24.7,7,518117,21,7,ack,55,-------,8,14561,16090,884950,Switch1,clien-7,45.38957,45.38957,45.399573,0,328.167767,18049.2,55,0.00844,0,45.369563,45.399573,1.030051,50.059851,Normal
24.6,6,13859,23,21,ack,55,-------,7,423,16091,885005,Router,Switch1,2.524349,2.524349,2.534352,0,328.19309,18050.6,55,0.00844,0,2.514346,2.544355,1.030048,50.059112,Smurf
4,24.4,140886,4,21,tcp,1540,-------,5,4393,16091,24780100,clien-4,Switch1,14.477384,14.477384,14.487507,0,328.217832,505455,1540,0.236321,0,14.477384,14.507754,1,50.025368,Normal
2,24.16,46094,2,21,ping,65535,-------,0,-1,210,13762400,client-2,switch1,44,44,44.015243,0,23.188498,1519660,65535.2,0.131248,0.056214,44,44.056214,0,9.056214,Smurf
24.11,11,337499,22,11,ack,55,-------,12,10014,16103,885665,Switch2,clien-11,31.533911,31.533911,31.543914,0,328.522947,18068.8,55,0.008446,0,31.513905,31.543914,1.030019,50.046382,Normal
24.15,15,264811,22,15,ack,55,-------,16,8176,16092,885060,Switch2,clien-15,25.957029,25.957029,25.967032,0,328.206042,18051.3,55,0.008441,0,25.937022,25.967032,1.030045,50.060221,Normal
24.13,13,201476,22,13,ack,55,-------,14,6278,16103,885665,Switch2,clien-13,20.20168,20.20168,20.211683,0,328.460278,18065.3,55,0.008446,0,20.181674,20.211683,1.030032,50.055747,Normal
24.3,3,395290,21,3,ack,55,-------,4,11465,16091,885005,Switch1,clien-3,35.968766,35.968766,35.97877,0,328.241051,18053.3,55,0.00844,0,35.94876,35.97877,1.030029,50.051929,Normal
14,24.14,186408,22,23,tcp,1540,-------,15,5812,16103,24798600,Switch2,Router,18.77443,18.77443,18.784554,0,328.431807,505785,1540,0.236498,0,18.764307,18.794677,1,50.029965,UDP-Flood