'''
Created on 2017/8/23

@author: ful6ru04
'''
import os
import numpy as np
import itertools
from matplotlib import pyplot as plt
import pickle

from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import LinearSVC
from sklearn import metrics
from sklearn.naive_bayes import GaussianNB
from sklearn.neural_network import MLPClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.neighbors import KNeighborsClassifier


def plot_confusion_matrix(confusion_matrix, classes, title='Confusion matrix', path=None):
    plt.figure()
    plt.imshow(confusion_matrix, interpolation='nearest', cmap=plt.cm.Blues)
    plt.title(title)
    plt.colorbar()
    tick_marks = np.arange(len(classes))
    plt.xticks(tick_marks, classes, rotation=45)
    plt.yticks(tick_marks, classes)

    fmt = '.2f' if False else 'd'
    thresh = confusion_matrix.max() / 2.
    for i, j in itertools.product(range(confusion_matrix.shape[0]), range(confusion_matrix.shape[1])):
        plt.text(j, i, format(confusion_matrix[i, j], fmt),
                 horizontalalignment="center",
                 color="white" if confusion_matrix[i, j] > thresh else "black")
    plt.tight_layout()
    plt.ylabel('True label')
    plt.xlabel('Predicted label')
    if not path is None:
        plt.savefig(path, bbox_inches='tight')

if __name__ == '__main__':
    ###################################################################################################
    # !Notice
    # Set to True if need to re-train a new model. Use old saved model, if set to False.
    RETRAIN = False
    if RETRAIN:
        print('Use re-train phase! Will train new  models!')
    else:
        print('Disable re-train phase! Will use saved models!')

    # Input file
    input_dir = './output-1'
    input_np_name = 'data.npz'

    # Output file
    output_dir = './output-2'
    output_figure = 'ROCcurve.png'
    ###################################################################################################

    ### Check if output directory exists
    if not os.path.exists(output_dir):
        os.mkdir(output_dir)

    ### Read data
    path = os.path.join(input_dir, input_np_name)
    data = np.load(path)
    try:
        train_y = np.ravel(data['train_y'])
        train_x = np.reshape(np.ravel(data['train_x']), (train_y.shape[0], -1))
        test_y = np.ravel(data['test_y'])
        test_x = np.reshape(np.ravel(data['test_x']), (test_y.shape[0], -1))
        print('Read data from "' + path + '"')
    except:
        print("Data doesn't exist! Please execute Main3 first!")
        exit(1)

    ### Random Forest
    model_name_rf = 'classifier_rf.pkl'
    if RETRAIN:
        classifier_rf = RandomForestClassifier(n_estimators=10)
        classifier_rf.fit(train_x, train_y)
        with open(os.path.join(output_dir, model_name_rf), 'wb') as fid:
            pickle.dump(classifier_rf, fid)
    else:
        with open(os.path.join(output_dir, model_name_rf), 'rb') as fid:
            classifier_rf = pickle.load(fid)

    pred_y = classifier_rf.predict(test_x)
    acc_rf = metrics.accuracy_score(test_y, pred_y)
    recall_rf = metrics.recall_score(test_y, pred_y)
    precision_rf = metrics.precision_score(test_y, pred_y)
    confusion_rf = metrics.confusion_matrix(test_y, pred_y)
    fpr_rf, tpr_rf, th = metrics.roc_curve(test_y, pred_y)
    roc_auc_rf = metrics.auc(fpr_rf, tpr_rf)
    print('Random forest Accuracy/Recall/Precision: ('
          + "{0:.2f}".format(acc_rf*100) + '%, '
          + "{0:.2f}".format(recall_rf*100) + '%, '
          + "{0:.2f}".format(precision_rf*100) + '%)')

    ### Linear SVM
    model_name_svm = 'classifier_svm.pkl'
    if RETRAIN:
        classifier_svm = LinearSVC()
        classifier_svm.fit(train_x, train_y)
        with open(os.path.join(output_dir, model_name_svm), 'wb') as fid:
            pickle.dump(classifier_svm, fid)
    else:
        with open(os.path.join(output_dir, model_name_svm), 'rb') as fid:
            classifier_svm = pickle.load(fid)

    pred_y = classifier_svm.predict(test_x)
    acc_svm = metrics.accuracy_score(test_y, pred_y)
    recall_svm = metrics.recall_score(test_y, pred_y)
    precision_svm = metrics.precision_score(test_y, pred_y)
    confusion_svm = metrics.confusion_matrix(test_y, pred_y)
    fpr_svm, tpr_svm, _ = metrics.roc_curve(test_y, pred_y)
    roc_auc_svm = metrics.auc(fpr_svm, tpr_svm)
    print('Linear SVM Accuracy/Recall/Precision: ('
          + "{0:.2f}".format(acc_svm*100) + '%, '
          + "{0:.2f}".format(recall_svm*100) + '%, '
          + "{0:.2f}".format(precision_svm*100) + '%)')
    
    ### Naive Bayes
    model_name_gnb = 'classifier_gnb.pkl'
    if RETRAIN:
        classifier_gnb = GaussianNB()
        classifier_gnb.fit(train_x, train_y)
        with open(os.path.join(output_dir, model_name_gnb), 'wb') as fid:
            pickle.dump(classifier_gnb, fid)
    else:
        with open(os.path.join(output_dir, model_name_gnb), 'rb') as fid:
            classifier_gnb = pickle.load(fid)

    pred_y = classifier_gnb.predict(test_x)
    acc_gnb = metrics.accuracy_score(test_y, pred_y)
    recall_gnb = metrics.recall_score(test_y, pred_y)
    precision_gnb = metrics.precision_score(test_y, pred_y)
    confusion_gnb = metrics.confusion_matrix(test_y, pred_y)
    fpr_gnb, tpr_gnb, _ = metrics.roc_curve(test_y, pred_y)
    roc_auc_gnb = metrics.auc(fpr_gnb, tpr_gnb)
    print('Naive Bayes Accuracy/Recall/Precision: ('
          + "{0:.2f}".format(acc_svm*100) + '%, '
          + "{0:.2f}".format(recall_gnb*100) + '%, '
          + "{0:.2f}".format(precision_gnb*100) + '%)')
    
    ### Neural Network
    model_name_nn = 'classifier_nn.pkl'
    if RETRAIN:
        classifier_nn = MLPClassifier()
        classifier_nn.fit(train_x, train_y)
        with open(os.path.join(output_dir, model_name_nn), 'wb') as fid:
            pickle.dump(classifier_nn, fid)
    else:
        with open(os.path.join(output_dir, model_name_nn), 'rb') as fid:
            classifier_nn = pickle.load(fid)

    pred_y = classifier_nn.predict(test_x)
    acc_nn = metrics.accuracy_score(test_y, pred_y)
    recall_nn = metrics.recall_score(test_y, pred_y)
    precision_nn = metrics.precision_score(test_y, pred_y)
    confusion_nn = metrics.confusion_matrix(test_y, pred_y)
    fpr_nn, tpr_nn, _ = metrics.roc_curve(test_y, pred_y)
    roc_auc_nn = metrics.auc(fpr_nn, tpr_nn)
    print('Neural Network Accuracy/Recall/Precision: ('
          + "{0:.2f}".format(acc_nn * 100) + '%, '
          + "{0:.2f}".format(recall_nn * 100) + '%, '
          + "{0:.2f}".format(precision_nn * 100) + '%)')
    
    ### Decision Tree
    model_name_dt = 'classifier_dt.pkl'
    if RETRAIN:
        classifier_dt = DecisionTreeClassifier(max_depth=5)
        classifier_dt.fit(train_x, train_y)
        with open(os.path.join(output_dir, model_name_dt), 'wb') as fid:
            pickle.dump(classifier_dt, fid)
    else:
        with open(os.path.join(output_dir, model_name_dt), 'rb') as fid:
            classifier_dt = pickle.load(fid)

    pred_y = classifier_dt.predict(test_x)
    acc_dt = metrics.accuracy_score(test_y, pred_y)
    recall_dt = metrics.recall_score(test_y, pred_y)
    precision_dt = metrics.precision_score(test_y, pred_y)
    confusion_dt = metrics.confusion_matrix(test_y, pred_y)
    fpr_dt, tpr_dt, _ = metrics.roc_curve(test_y, pred_y)
    roc_auc_dt = metrics.auc(fpr_dt, tpr_dt)
    print('Decision Tree Accuracy/Recall/Precision: ('
          + "{0:.2f}".format(acc_dt * 100) + '%, '
          + "{0:.2f}".format(recall_dt * 100) + '%, '
          + "{0:.2f}".format(precision_dt * 100) + '%)')

    ### K-nn
    model_name_knn = 'classifier_knn.pkl'
    if RETRAIN:
        classifier_knn = KNeighborsClassifier(3)
        classifier_knn.fit(train_x, train_y)
        with open(os.path.join(output_dir, model_name_knn), 'wb') as fid:
            pickle.dump(classifier_knn, fid)
    else:
        with open(os.path.join(output_dir, model_name_knn), 'rb') as fid:
            classifier_knn = pickle.load(fid)

    pred_y = classifier_knn.predict(test_x)
    acc_knn = metrics.accuracy_score(test_y, pred_y)
    recall_knn = metrics.recall_score(test_y, pred_y)
    precision_knn = metrics.precision_score(test_y, pred_y)
    confusion_knn = metrics.confusion_matrix(test_y, pred_y)
    fpr_knn, tpr_knn, _ = metrics.roc_curve(test_y, pred_y)
    roc_auc_knn = metrics.auc(fpr_knn, tpr_knn)
    print('K-nn Accuracy/Recall/Precision: ('
          + "{0:.2f}".format(acc_knn * 100) + '%, '
          + "{0:.2f}".format(recall_knn * 100) + '%, '
          + "{0:.2f}".format(precision_knn * 100) + '%)')

    # Draw ROC Curve
    plt.plot(fpr_dt, tpr_dt, label='DecisionTree (AUC = %0.3f)' % roc_auc_dt)
    plt.plot(fpr_gnb, tpr_gnb, label='NaiveBayes (AUC = %0.3f)' % roc_auc_gnb)
    plt.plot(fpr_knn, tpr_knn, label='K-nn (AUC = %0.3f)' % roc_auc_knn)
    plt.plot(fpr_nn, tpr_nn, label='NeuralNetwork (AUC = %0.3f)' % roc_auc_nn)
    plt.plot(fpr_rf, tpr_rf, label='RandomForest (AUC = %0.3f)' % roc_auc_rf)
    plt.plot(fpr_svm, tpr_svm, label='LinearSVM (AUC = %0.3f)' % roc_auc_svm)
    plt.plot([0, 1], [0, 1], 'k--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.0])
    plt.xlabel('False Positive')
    plt.ylabel('True Positive')
    plt.title('ROC-AUC Curve')
    plt.legend(loc='lower right')
    path = os.path.join(output_dir, output_figure)
    plt.savefig(path, bbox_inches='tight')

    # Draw Confusion Matrix
    classes = ['Normal', 'Malware']
    plot_confusion_matrix(confusion_dt, classes, 'Decision Tree Confusion Matrix', os.path.join(output_dir, 'dt.png'))
    plot_confusion_matrix(confusion_gnb, classes, 'Naive Bayes Confusion Matrix', os.path.join(output_dir, 'gnb.png'))
    plot_confusion_matrix(confusion_knn, classes, 'K-nn Confusion Matrix', os.path.join(output_dir, 'knn.png'))
    plot_confusion_matrix(confusion_nn, classes, 'Neural Network Confusion Matrix', os.path.join(output_dir, 'nn.png'))
    plot_confusion_matrix(confusion_rf, classes, 'Random Forest Confusion Matrix', os.path.join(output_dir, 'rf.png'))
    plot_confusion_matrix(confusion_svm, classes, 'Linear SVM Confusion Matrix', os.path.join(output_dir, 'svm.png'))
    plt.show()