# -*- coding: utf-8 -*-

"""
@Datetime: 2019/3/28
@Author: Zhang Yafei
"""
# #################################  nltk计算tfidf ####################################333
# import nltk
# from nltk.text import TextCollection
# sents = ['this is sentence one', 'this is sentence two', 'this is sentence three']
# sents = [nltk.word_tokenize(sent) for sent in sents]
# corpus = TextCollection(sents)
#
# # 直接就能算出tfidf
# # (term: ⼀句话中的某个term, text: 这句话)
# print(corpus.idf('three'))
# print(corpus.tf('four',nltk.word_tokenize('this is a sentence four')))
# print(corpus.tf_idf('four',nltk.word_tokenize('this is a sentence four')))

# #################################  sklearn计算tfidf ####################################333
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.feature_extraction.text import CountVectorizer

corpus = ["我 来到 北京 清华大学",  # 第一类文本切词后的结果，词之间以空格隔开
          "他 来到 了 网易 杭研 大厦",  # 第二类文本的切词结果
          "小明 硕士 毕业 与 中国 科学院",  # 第三类文本的切词结果
          "我 爱 北京 天安门"]  # 第四类文本的切词结果
# vectorizer = CountVectorizer()  # 该类会将文本中的词语转换为词频矩阵，矩阵元素a[i][j] 表示j词在i类文本下的词频
# transformer = TfidfTransformer()  # 该类会统计每个词语的tf-idf权值
# tfidf = transformer.fit_transform(vectorizer.fit_transform(corpus))  # 第一个fit_transform是计算tf-idf，第二个fit_transform是将文本转为词频矩阵
# word = vectorizer.get_feature_names()  # 获取词袋模型中的所有词语
# weight = tfidf.toarray()  # 将tf-idf矩阵抽取出来，元素a[i][j]表示j词在i类文本中的tf-idf权重
# for i in range(len(weight)):  # 打印每类文本的tf-idf词语权重，第一个for遍历所有文本，第二个for便利某一类文本下的词语权重
#     print("-------这里输出第", i, "类文本的词语tf-idf权重------")
#     for j in range(len(word)):
#         print(word[j], weight[i][j])
# print(weight)


# #################################  gensim计算tfidf ####################################333
import warnings

warnings.filterwarnings(action='ignore', category=UserWarning, module='gensim')

from gensim import corpora, models, similarities

texts = [[word for word in text.split()] for text in corpus]

dictionary = corpora.Dictionary(texts)
print(dictionary)
# 9、对语料库进一步处理，得到新语料库
corpus = [dictionary.doc2bow(text) for text in texts]

# 10、将新语料库通过tf-idf model 进行处理，得到tfidf
tfidf = models.TfidfModel(corpus)

# 8、将要对比的文档通过doc2bow转化为稀疏向量
data3 = "基本思路 我 爱 北京 天安门"
new_xs = dictionary.doc2bow(data3.split())
print(new_xs)

featurenum = len(dictionary.token2id.keys())
#12、稀疏矩阵相似度，从而建立索引
index = similarities.SparseMatrixSimilarity(tfidf[corpus], num_features=featurenum)
#13、得到最终相似结果
sim = index[tfidf[new_xs]]
print(sim)