import numpy as np
from os import listdir
from os.path import isfile, join
import csv
from sklearn import model_selection, preprocessing, linear_model, naive_bayes, metrics, svm
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn import ensemble
import xgboost
import numpy
from sklearn.metrics import confusion_matrix
import matplotlib.pyplot as plt
import pandas as pd
import time
import sys

# check correct parameter
if len(sys.argv) != 2 or (sys.argv[1] != 'opcode' and sys.argv[1] != 'API'):
    print("Insufficient arguments")
    sys.exit()

# getting parameters
mode = sys.argv[1]

print("\tstart malware classifier\n")
start = time.time()

trainLabelFile = "1st_answer.csv"
with open(trainLabelFile) as f:
    labelDict = dict(filter(None, csv.reader(f)))
print("\tCreate LabelDict done!\n")


# Loading the created Dataset
labels, sequences = [], []
if mode == 'opcode':
    malware_path = './opcode/1/'
    normal_path = './opcode/0/'
else:
    malware_path = './API/1/'
    normal_path = './API/0/'

files1 = [f for f in listdir(malware_path)
          if isfile(join(malware_path, f))]
files2 = [f for f in listdir(normal_path) if isfile(join(normal_path, f))]

for i, file in enumerate(files1):
    if not file:
        continue
    with open(malware_path + file, 'r') as f:
        try:
            labels.append(int(labelDict[file[:len(file)-4] + '.vir']))
            lines = f.readlines()
            sequences.append(" ".join(lines))
        except:
            print(i)
            pass

print("\tLoad malware dataset done!")

for i, file in enumerate(files2):
    if not file:
        continue
    with open(normal_path + file, 'r') as f:
        try:
            labels.append(int(labelDict[file[:len(file)-4] + '.vir']))
            lines = f.readlines()
            sequences.append(" ".join(lines))
        except:
            print(i)
            pass

print("\tLoad normal dataset done!")

# create a dataframe using texts and lables
trainDF = pd.DataFrame()
trainDF['sequences'] = sequences
trainDF['label'] = labels
print("len(trainDF)   :", len(trainDF))
print(trainDF.head())

# split the dataset into training and validation datasets
train_x, valid_x, train_y, valid_y = model_selection.train_test_split(
    trainDF['sequences'], trainDF['label'])
print(len(train_x))

# label encode the target variable
encoder = preprocessing.LabelEncoder()
train_y = encoder.fit_transform(train_y)
valid_y = encoder.fit_transform(valid_y)

# To get the distribution of the malware family dataset
(unique, counts) = numpy.unique(valid_y, return_counts=True)
frequencies = numpy.asarray((unique, counts)).T
print(frequencies)


# ngram(Tri-gram) level tf-idf
tfidf_vect_ngram = TfidfVectorizer(
    analyzer='word', token_pattern=r'\w{1,}', ngram_range=(3, 3), max_features=1000)
# 학습
tfidf_vect_ngram.fit(trainDF['sequences'])
# 학습한 변환 적용
xtrain_tfidf_ngram = tfidf_vect_ngram.transform(train_x)
xvalid_tfidf_ngram = tfidf_vect_ngram.transform(valid_x)
print("Ngram 기반 TFIDF vector 추출 완료\n")

# To Get CSV File Features of the N grams
feature_names = tfidf_vect_ngram.get_feature_names()
dense = xtrain_tfidf_ngram.todense()
denselist = dense.tolist()
df = pd.DataFrame(denselist, columns=feature_names)
df.insert(0, column="label", value=train_y)
df.to_csv("done.csv")
print("save csv done")


def plot_confusion_matrix(cm, classes,
                          normalize=False,
                          title='Confusion matrix',
                          cmap=plt.cm.Blues):
    """
    This function prints and plots the confusion matrix.
    Normalization can be applied by setting `normalize=True`.
    """
    import itertools
#     if normalize:
#         cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
#         #print("Normalized confusion matrix")
#     else:
#         #print('Confusion matrix, without normalization')

#     #print(cm)

    plt.imshow(cm, interpolation='nearest', cmap=cmap)
    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 normalize else 'd'
    thresh = cm.max() / 2.
    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
        plt.text(j, i, format(cm[i, j], fmt),
                 horizontalalignment="center",
                 color="white" if cm[i, j] > thresh else "black")

    plt.ylabel('True label')
    plt.xlabel('Predicted label')
    plt.tight_layout()


def train_model(classifier, feature_vector_train, label, feature_vector_valid, pic_title, is_neural_net=False):
    # fit the training dataset on the classifier
    classifier.fit(feature_vector_train, label)

    # predict the labels on validation dataset
    predictions = classifier.predict(feature_vector_valid)
    if is_neural_net:
        predictions = predictions.argmax(axis=-1)
    cnf_matrix = confusion_matrix(valid_y, predictions, labels=[0, 1])
    # numpy float 출력옵션 변경 / 소숫점 2
    np.set_printoptions(precision=2)

    # Plot non-normalized confusion matrix
    plt.figure()
    plot_confusion_matrix(cnf_matrix, classes=[0, 1],
                          title='Confusion Matrix')
    plt.savefig(pic_title)

    return metrics.accuracy_score(predictions, valid_y)


# Naive Bayes on Ngram Level TF IDF Vectors
accuracy1 = train_model(naive_bayes.MultinomialNB(),
                        xtrain_tfidf_ngram, train_y, xvalid_tfidf_ngram, "NB.png")
print("NB, N-Gram Vectors: ", accuracy1)

# Linear Classifier on Ngram Level TF IDF Vectors
accuracy2 = train_model(linear_model.LogisticRegression(
), xtrain_tfidf_ngram, train_y, xvalid_tfidf_ngram, "LR.png")
print("LR, N-Gram Vectors: ", accuracy2)

# SVM on Ngram Level TF IDF Vectors
accuracy3 = train_model(svm.SVC(kernel='rbf', gamma=0.7),
                        xtrain_tfidf_ngram, train_y, xvalid_tfidf_ngram, "SVM.png")
print("SVM, N-Gram Vectors: ", accuracy3)

# RF on Ngram Level TF IDF Vectors
accuracy4 = train_model(ensemble.RandomForestClassifier(
), xtrain_tfidf_ngram, train_y, xvalid_tfidf_ngram, "RF.png")
print("RF, Ngrams: ", accuracy4)

# Extereme Gradient Boosting on Word Level TF IDF Vectors
accuracy5 = train_model(xgboost.XGBClassifier(
), xtrain_tfidf_ngram.tocsc(), train_y, xvalid_tfidf_ngram.tocsc(), "Xgb.png")
print("Xgb, Ngram TF-IDF: ", accuracy5)

print("time: ", time.time()-start)
