# import section
from os import walk

import numpy as np
from sklearn.preprocessing import LabelEncoder, OneHotEncoder

from tensorflow.keras.optimizers import SGD, Adam, Adadelta
from tensorflow.keras.layers import Conv1D, Dense, MaxPooling1D, Flatten, Dropout
from tensorflow.keras.models import Sequential
from tensorflow.keras.regularizers import l2, l1

from sklearn.metrics import confusion_matrix
import tensorflow.keras as keras
import datetime
import argparse
import os

SEQUENCE_SIZE = 1000

def padding(seq_size_max, line):
    less = seq_size_max-len(line)
    padding_value = line[:less]
    line = line+padding_value
    if(len(line) < seq_size_max):
        line = padding(seq_size_max, line)
    return line

def load_sequence_dataset_from_file(file_path):
    seq_size_max = SEQUENCE_SIZE
    sequences = []
    with open(file_path) as f:
        lines = f.readlines()
        for line in lines:
            line = line.strip()
            if len(line) >= seq_size_max:
                sequences.append(line[:seq_size_max])
            else:
                line = padding(seq_size_max, line)
                sequences.append(line)
    return sequences

def check_dataset_validity(dataset, file_path):
    i = 1
    for sequence in dataset:
        if any(x not in ['A', 'T', 'C', 'G', 'N'] for x in set(list(sequence))) or len(sequence)!=SEQUENCE_SIZE:
            print('Following problem is found in dataset at ', file_path)
            print('At line number:', i)
            print('Charaters :', set(sequence))
            print('Length: ', len(sequence), 'where sequence len should be', SEQUENCE_SIZE)
            print('Sequence', sequence)
            exit()
        i += 1

def main(file_path_clean, file_path_trojan, result_file, log_dir):
    dataset_a = load_sequence_dataset_from_file(file_path_clean)
    check_dataset_validity(dataset_a, file_path_clean)
    dataset_b = load_sequence_dataset_from_file(file_path_trojan)
    check_dataset_validity(dataset_b, file_path_trojan)

    print('dataset_a len:', len(dataset_a), 'dataset_b len:', len(dataset_b))
    label_a = list(np.zeros(len(dataset_a)))
    label_b = list(np.ones(len(dataset_b)))

    sequences = dataset_a + dataset_b
    labels = label_a + label_b

    # data preprocessing
    # encoding sequences and labels
    integer_encoder = LabelEncoder()
    one_hot_encoder = OneHotEncoder()
    input_features = []

    one_hot_encoder.fit(np.array(list('ACTGN')).reshape(-1, 1))
    for sequence in sequences:
        one_hot_encoded = one_hot_encoder.transform(np.array(list(sequence)).reshape(-1, 1)).toarray()
        input_features.append(one_hot_encoded)
    input_features = np.stack(input_features)

    one_hot_encoder = OneHotEncoder()
    labels = np.array(labels).reshape(-1, 1)
    input_labels = one_hot_encoder.fit_transform(labels).toarray()

    # splitting test and training data
    from sklearn.model_selection import train_test_split
    train_features, test_features, train_labels, test_labels = train_test_split(
        input_features, input_labels, test_size=0.25, random_state=42)

    # constructing cnn
    model = Sequential()

    model.add(Conv1D(filters=32, kernel_size=5, padding='same',
                     kernel_initializer='he_uniform',
                     input_shape=(train_features.shape[1], 5), activation='relu',  # 5 because A C T G N
                     kernel_regularizer=l2(0.001)))
    model.add(Dropout(0.2))
    model.add(Conv1D(filters=32, kernel_size=5, padding='same',
                     kernel_initializer='he_uniform',
                     activation='relu',
                     kernel_regularizer=l2(0.001)))
    model.add(Dropout(0.2))
    model.add(MaxPooling1D(pool_size=2))

    model.add(Flatten())
    model.add(Dense(16, kernel_initializer='he_uniform', activation='relu',
                    kernel_regularizer=l2(0.001)))
    model.add(Dropout(0.1))
    model.add(Dense(2, activation='softmax'))
    # model.add(Dense(2, activation='softmax'))

    # setting up training run
    epochs = 3000
    lrate = 0.001
    decay = lrate / epochs
    sgd = SGD(lr=lrate, momentum=0.9, decay=decay, nesterov=False)

    # run training and get training summary
    model.compile(loss='binary_crossentropy', optimizer=sgd, metrics=['binary_accuracy'])
    model.summary()
    # The fit() method - trains the model
    #log_dir = "logs/fit/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    tensorboard_callback = keras.callbacks.TensorBoard(log_dir=log_dir)
    model.fit(train_features, train_labels, nb_epoch=epochs, batch_size=100,
              callbacks=[tensorboard_callback], validation_split=0.10)

    # test results
    predicted_labels = model.predict(np.stack(test_features))
    cm = confusion_matrix(np.argmax(test_labels, axis=1),
                          np.argmax(predicted_labels, axis=1))
    cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]  # normalized
    print('Confusion matrix:\n', cm)

    scores = model.evaluate(test_features, test_labels, verbose=0)
    print("Accuracy: %.2f%%" % (scores[1] * 100))
    TN = cm[0][0]
    FN = cm[0][1]
    FP = cm[1][0]
    TP = cm[1][1]
    accuracy = scores[1] * 100
    with open(result_file, 'w') as f:
        f.write('TN, FN, FP, TP, accuracy\n')
        f.write('{}, {}, {}, {}, {}\n'.format(TN, FN, FP, TP, accuracy))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--fragment_len", help="slice size, between 1-5")
    parser.add_argument("--retention_pos", help="retention position, between 0-5")
    parser.add_argument("--encryption_key", help="encryption key, examle 10, 20, 30 etc. 0 means no encryption.")
    parser.add_argument("--dataset_num", help="dataset number, between 0-9")
    parser.add_argument("--scenario", help="trojan scenarios. random, nw_best, "
                                           "nw_worst, nw_greedy_best, nw_greedy_worst")
    args = parser.parse_args()
    fragment_len = args.fragment_len if int(args.fragment_len) in range(1,6) else None
    dataset_num = args.dataset_num if int(args.dataset_num) in range(0, 10) else None
    retention_pos = args.dataset_num if int(args.retention_pos) in range(0, 5) else None
    encryption_key = args.dataset_num if int(args.encryption_key) in range(0, 127) else None
    file_names = {
        'random': 'random_trojan_insertion_dataset',
        'nw_best': 'nw_best_straight_forward_trojan_insertion_dataset',
        'nw_worst': 'nw_worst_straight_forward_trojan_insertion_dataset',
        'nw_greedy_best': 'nw_best_greedy_trojan_insertion_dataset',
        'nw_greedy_worst': 'nw_worst_greedy_trojan_insertion_dataset'
    }
    scenario = file_names[args.scenario] if args.scenario in file_names.keys() else None
    if any(x is None for x in [fragment_len, dataset_num, scenario, retention_pos, encryption_key]):
        print('The parameters are not correct')
        exit()
    data_path = 'experiment_data/datasets/fragment_len_{}/retention_pos_{}/encryption_key_{}/dataset_{}/'. \
        format(fragment_len, retention_pos, encryption_key, dataset_num)
    file_path_clean = data_path+'non_trojan_dataset.txt'
    file_path_trojan = data_path+scenario+'.txt'
    result_dir = 'experiment_results/results/'
    os.makedirs(result_dir, exist_ok=True)
    logs_dir = 'experiment_results/logs/'
    os.makedirs(logs_dir, exist_ok=True)
    result_file = result_dir+'fragment_len_{}.scenario_{}.dataset_num_{}.txt'.format(fragment_len, scenario, dataset_num)
    logs_location = logs_dir + 'fragment_len_{}/scenario_{}/dataset_num_{}/'.format(fragment_len, scenario, dataset_num)
    # print(file_path_clean, file_path_trojan, result_file, logs_location)
    main(file_path_clean, file_path_trojan, result_file, logs_location)