import numpy as np
from tensorflow.keras.models import load_model

from Data.Preprocess import get_data
from GenerateAbnormalData import generate_abnormal_data
from IDS.CNN.CNNetwork import run_cnn_network, evaluate_cnn_network
from Network.BiGAN.RunBiGanNetwork import run_bigan
from Network.GAN.RunNetwork import run_network


def get_ids_train_data():
    data = get_data('train')
    _train_data = data['all'][0]
    _train_label = data['all'][1]
    labels = ['label-dos', 'label-normal', 'label-probe', 'label-r2l', 'label-u2r']
    for label in ['normal', 'u2r', 'r2l', 'probe', 'dos']:
        print('[-]' + label + '数据量为：' + str(data[label][0].shape[0]))
    while True:
        model_name = input('[-]请输入生成器模型名称(输入go开始训练IDS)：')
        if model_name == 'go':
            break
        label_name = input('[-]请输入生成器数据类型(dos,normal,probe,r2l,u2r)：')
        try:
            generator = load_model('Models/Gan/' + model_name + '.h5')
        except Exception as e:
            print('[-]没有此模型，名字输错啦')
        num = int(input('[-]请输入生成数据量：'))
        fake_abnormal_data_x = generate_abnormal_data(generator, num)
        _train_data = np.append(_train_data, fake_abnormal_data_x, axis=0)

        fake_abnormal_data_y = np.zeros((fake_abnormal_data_x.shape[0], 5))
        for i in range(5):
            if labels[i].endswith(label_name):
                fake_abnormal_data_y[::, i] = 1
        _train_label = np.append(_train_label, fake_abnormal_data_y, axis=0)

    return _train_data, _train_label


def train_gan():
    type_of_gan = input('[-]请输入神经网络模型类型( ① GAN ② BIGAN )：')
    model_name = input('[-]请输入模型名称：')
    data_class = input('[-]请输入需要生成的数据：')
    data = get_data('train')[data_class]
    epochs = int(input('[-]请输入训练次数：'))
    if int(type_of_gan) == 1:
        generator = run_network(data[0], data[1], epochs)
        generator.save('Models/Gan/' + 'gan_' + model_name + '_' + data_class + '.h5')
    elif int(type_of_gan) == 2:
        generator = run_bigan(data[0], data[1], epochs)
        generator.save('Models/Gan/' + 'bigan_' + model_name + '_' + data_class + '.h5')


def train_ids_without_gan():
    model_name = input('[-]请输入模型名称：')
    data = get_data('train')['all']
    epochs = int(input('[-]请输入训练次数：'))
    cnn = run_cnn_network(data[0], data[1], epochs)
    cnn.save('Models/IDS/' + model_name + '.h5')


def train_ids_with_gan():
    train_data, train_label = get_ids_train_data()
    model_name = input('[-]请输入IDS模型名称：')
    epochs = int(input('[-]请输入训练次数：'))
    cnn = run_cnn_network(train_data, train_label, epochs)
    cnn.save('Models/IDS/' + model_name + '.h5')


def load_and_test_ids():
    model_name = input('[-]请输入IDS模型名称：')
    cnn = load_model('Models/IDS/' + model_name + '.h5')
    data = get_data('test')
    losses = []
    accuracies = []
    flags = []
    while True:
        opt_flag = input('[-]请输入需要测试的样本集(all,normal,u2r,r2l,probe)：')
        flags.append(opt_flag)
        try:
            loss, accuracy = evaluate_cnn_network(cnn, data[opt_flag][0], data[opt_flag][1], 'KDD' + opt_flag)
            losses.append(loss)
            accuracies.append(accuracy)
            for i in range(len(losses)):
                print('[-]( 测试集: KDD-' + flags[i] + ' )  测试损失为：' + str(losses[i]) + '  准确度为：' + str(accuracies[i]))
        except KeyError as e:
            break


def main():
    flag = input('[-]请输入操作( ① GAN模型 ② IDS模型 )：')
    if int(flag) == 1:
        train_gan()
    elif int(flag) == 2:
        flag_for_ids = input('[-] ① 单独训练IDS模型 ② 使用GAN生成数据并训练IDS模型 ③ 加载并测试IDS模型 ：')
        if int(flag_for_ids) == 1:
            train_ids_without_gan()
        elif int(flag_for_ids) == 2:
            train_ids_with_gan()
        elif int(flag_for_ids) == 3:
            load_and_test_ids()


if __name__ == '__main__':
    main()