#coding=utf-8
import tensorflow as tf
from tensorflow import keras
import numpy as np
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.naive_bayes import GaussianNB
from sklearn import metrics
from sklearn.svm import SVC
import re
from sklearn.utils import shuffle
from sklearn.neural_network import MLPClassifier
from xgboost import XGBClassifier
from sklearn import preprocessing
import glob
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics import accuracy_score, confusion_matrix, classification_report

#数据载入
def load_file_from_dir(dir):
    """
    :param dir:文件地址
    :return:
    texts:当前目录下的全部文本，list
    """
    files = glob.glob(dir)
    texts = []
    for file in files:
        print('loading files from %s'% file)
        with open(file) as f:
            lines = f.readlines()
            liness = " ".join(lines)
            # 文本中出现标签，变换为空,re.I:忽略大小写
            liness = re.sub(r"[APT1|Crypto|Locker|Zeus]", ' ', liness, flags=re.I)
            texts.append(liness)
    return texts

def load_files():
    """
    :return:
    x：全部文本，list,4762
    y:标签，list, 4762
    """
    malware_class = ['APT1', 'Crypto', 'Locker', 'Zeus']
    x = []
    y = []
    for i, name in enumerate(malware_class):
        dir = '/home/hezhouyu/projects/dataset/malware/%s/*' % name
        print('loading files from %s index %d' % (dir, i))
        v = load_file_from_dir(dir)
        x += v
        y += [i]*len(v)

    print(type(x), len(x), type(y), len(y))
    return x, y

#特征提取
def get_features_ngram(x, y):
    """
    3-gram词袋模型
    :param x:文本内容，list of string
    :param y: 文本标签，list of int
    :return:
    x_train :训练样本内容，array，(n_samples,output_dim)=(3333,500)
    x_test = 测试样本内容，array，(n_samples,output_dim)=(1429,500)
    y_train = 训练样本标签，array，(n_samples, )=(3333, )
    y_test = 测试样本标签，array，(n_samples, )=(1429, )
    """
    vectorizer = CountVectorizer(
        decode_error='ignore',
        ngram_range=(2, 4),
        token_pattern=r'\b\w+\b',#按单词切分
        strip_accents='ascii',
        max_features=max_words,
        stop_words='english',
        max_df=1.0,#作为一个阈值，词是否当作关键词。表示词出现的次数与语料库文档数的百分比
        min_df=1
    )
    print(vectorizer)
    x = vectorizer.fit_transform(x)
    transformer = TfidfTransformer(smooth_idf=False)
    x = transformer.fit_transform(x)

    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.3, shuffle=True)  # 划分数据集
    x_train = x_train.toarray()  # 转矩阵
    x_test = x_test.toarray()
    y_train = np.array(y_train)
    y_test = np.array(y_test)

    print("Transform texts into 2-gram successfully")
    print(x_train.shape, y_train.shape, x_test.shape, y_test.shape)

    return x_train, y_train, x_test, y_test

def get_features_tfidf(x, y):
    tfidf_vectorizer = TfidfVectorizer(ngram_range=(2, 2), decode_error="ignore", lowercase=True,
                                       stop_words="english", max_features=max_words, binary=False, token_pattern=r'\b\w+\b')
    x = tfidf_vectorizer.fit_transform(x)

    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.3, shuffle=True)  # 划分数据集
    x_train = x_train.toarray()  # 转矩阵
    x_test = x_test.toarray()
    y_train = np.array(y_train)
    y_test = np.array(y_test)

    print("Transform texts into tfidf-2gram successfully")
    print(x_train.shape, y_train.shape, x_test.shape, y_test.shape)

    return x_train, y_train, x_test, y_test
#模型的训练与验证
def svm():
    """
    2,4-gram-tfidf-acc:0.606
    :return:返回svm模型
    """
    s = SVC(C=1.0, kernel='linear', decision_function_shape='ovr')
    return s

def xgb():
    """
    2,4-gram-tfidf-acc:0.6815
    :return: 返回xgboosting模型
    """
    model = XGBClassifier(n_estimate=100, max_depth=8, n_jobs=1)#n_estimate:决策树个数
    return model

def mlp():
    """
    2,4-gram-tfidf-acc:0.6102
    :return:返回mlp模型
    """
    model = MLPClassifier(hidden_layer_sizes=(10, 4), activation='relu', solver='adam', max_iter=500,
                          shuffle=True, verbose=1)
    return model

if __name__ == "__main__":
    max_words = 1000
    x, y = load_files()
    x_train, y_train, x_test, y_test = get_features_tfidf(x, y)
    #model = svm()
    #model = xgb()
    model = mlp()
    trained_model = model.fit(x_train, y_train)
    y_pred = trained_model.predict(x_test)
    acc = accuracy_score(y_test, y_pred)
    clf = classification_report(y_test, y_pred)
    print(acc, '\n', clf)
