import csv
import nltk
import re
from nltk.corpus import stopwords
from nltk.stem import WordNetLemmatizer
import pandas as pd

#返回类别
def getLb(data):
    if data.startswith("J"):
        return nltk.corpus.wordnet.ADJ
    elif data.startswith("V"):
        return nltk.corpus.wordnet.VERB
    elif data.startswith("N"):
        return nltk.corpus.wordnet.NOUN
    elif data.startswith("R"):
        return nltk.corpus.wordnet.ADV
    else:
        return "";

def preprocessing(data):
    newdata=[]
    punctuation = '!,;:?"\''
    data=re.sub(r'[{}]+'.format(punctuation), '', data).strip().lower();#去标点和转小写
    for i in nltk.sent_tokenize(data, "english"):  # 对文本按照句子进行分割
        for j in nltk.word_tokenize(i):  # 对句子进行分词
            newdata.append(j)
    stops = stopwords.words('english')
    newdata= [i for i in newdata if i not in stops]#去停用词
    newdata = nltk.pos_tag(newdata)#词性标注
    lem = WordNetLemmatizer()
    for i, j in enumerate(newdata):#还原词
        y = getLb(j[1])
        if y:
            newdata[i] = lem.lemmatize(j[0], y)
        else:
            newdata[i] = j[0]

    newdata=' '.join(newdata);
    return newdata

if __name__ == "__main__":
    sms = open('./_5', 'r', encoding='utf-8')
    sms_data = []
    sms_label = []
    csv_reader = csv.reader(sms, delimiter='\t')
    for line in csv_reader:
        sms_label.append(line[0])
        sms_data.append(preprocessing(line[1]))
    sms.close()
    result=pd.DataFrame({"类别":sms_label,"特征":sms_data})
    result.to_csv("./result.csv")

from sklearn.model_selection import train_test_split
import pandas as pd

data=pd.read_csv("./分类/垃圾邮件分类/result.csv")

x_train,x_test, y_train, y_test = train_test_split(data.iloc[:,2], data.iloc[:,1], test_size=0.2, random_state=0

from sklearn.feature_extraction.text import TfidfVectorizer

tfidf2 = TfidfVectorizer()
X_train=tfidf2.fit_transform(x_train.values.astype('U'))
X_text=tfidf2.transform(x_test.values.astype('U'))

#邮件与向量的关系
tfidf2.vocabulary_

X_train.toarray()
X_train.toarray().shape

#向量还原成邮件
import numpy as np
a=np.flatnonzero(X_train.toarray()[0])
a
X_train.toarray()[0][a]

b=tfidf2.vocabulary_
key_list=[]
for key,value in b.items():
    if value in a:
        key_list.append(key)
key_list
x_train.iloc[0]

from sklearn.naive_bayes import MultinomialNB
mnb=MultinomialNB()
mnb.fit(X_train,y_train)
y_mnb=mnb.predict(X_text)

from sklearn.metrics import confusion_matrix
from sklearn.metrics import classification_report

# 混淆矩阵
cm = confusion_matrix(y_test, y_mnb)
print('nb_confusion_matrix:')
print(cm)
# 主要分类指标的文本报告
cr = classification_report(y_test, y_mnb)
print('nb_classification_report:')
print(cr)
(cm[0][0]+cm[1][1])/np.sum(cm)