import pandas as pd
from tqdm import tqdm
import numpy as np
import pickle
from gensim.models import KeyedVectors  # 词向量用来比较俩俩之间相似度

# 读取数据： 导入在preprocessor.ipynb中生成的data/question_answer_pares.pkl文件，并将其保存在变量QApares中
with open('data/question_answer_pares.pkl','rb') as f:
    QApares = pickle.load(f)
# 添加一列展示列表长度
QApares["num"] = [len(QApares.question_after_preprocessing[_]) for _ in range(len(QApares.question_after_preprocessing.values))]
# 去除空列表，并重新index
QApares = QApares.loc[QApares.num > 0].reset_index(drop=True)

# 查看关键词列表长度
# QApares.num.value_counts()

# 构建一个倒排表，有关倒排表的详细内容参考实验手册
# 为了能够快速检索，倒排表应用哈希表来存储。python中字典内部便是用哈希表来存储的，所以这里我们直接将倒排表保存在字典中
# 注意：在这里不需要考虑单词之间的相似度。
inverted_list = {}
for idx, sentence in enumerate(QApares.question_after_preprocessing.values):
    for each in sentence:
        if not inverted_list.get(each):
            inverted_list[each] = []
        inverted_list[each].append(idx)

#d ata/retrieve/sgns.zhihu.word是从https://github.com/Embedding/Chinese-Word-Vectors下载到的预训练好的中文词向量文件
#使 用KeyedVectors.load_word2vec_format()函数加载预训练好的词向量文件
model = KeyedVectors.load_word2vec_format('data/retrieve/sgns.zhihu.word')

def get_similar_by_word(word,topk):
    '''
        返回与一个单词word相似度最高的topk个单词所组成的单词列表
        出参：
            word_list：与word相似度最高的topk个单词所组成的单词列表。格式为[单词1，单词2，单词3，单词4，单词5]
    '''
    word_list = [word]
    try:
        similar_words = model.similar_by_word(word,topk)
        for word in similar_words:
            word_list.append(word[0])
    except Exception:
        pass
    return word_list

# TODO：
# 构造一个新的倒排表，并将结果保存在字典inverted_list_new中
# 新的倒排表键为word，值为老倒排表[word]、老倒排表[单词1]、老倒排表[单词2]、老倒排表[单词3]、老倒排表[单词4]的并集
# 即新倒排表保存了包含单词word或包含与单词word最相近的5个单词中的某一个的问题的index
inverted_list_new = {}
for word in tqdm(inverted_list):
    li_ = []
    for each in get_similar_by_word(word, 5):
        li_ += inverted_list.get(each, [])
    inverted_list_new[word] = set(li_)

# 将新的倒排表保存在文件data/retrieve/invertedList.pkl中
with open('data/retrieve/invertedList.pkl','wb') as f:
    pickle.dump(inverted_list_new,f)

#这一格的内容是从preprocessor.ipynb中粘贴而来，包含了数据预处理的几个关键函数
import emoji
import re
import jieba
def clean(content):
    content = emoji.demojize(content)
    content = re.sub('<.*>','',content)
    return content
#这一函数是用于对句子进行分词，在preprocessor.ipynb中由于数据是已经分好词的，所以我们并没有进行这一步骤，但是对于一个新的问句，这一步是必不可少的
def question_cut(content):
    return list(jieba.cut(content))
def strip(wordList):
    return [word.strip() for word in wordList if word.strip()!='']
with open("data/stopWord.json","r", encoding="utf-8") as f:
    stopWords = f.read().split("\n")
def rm_stop_word(wordList):
    return [word for word in wordList if word not in stopWords]

# 从data/retrieve/invertedList.pkl加载倒排表并将其保存在变量invertedList中
with open('data/retrieve/invertedList.pkl','rb') as f:
    invertedList = pickle.load(f)

def get_retrieve_result(sentence):
    '''
        输入一个句子sentence，根据倒排表进行快速检索，返回与该句子较相近的一些候选问题的index
        候选问题由包含该句子中任一单词或包含与该句子中任一单词意思相近的单词的问题索引组成
    '''
    sentence = clean(sentence)
    sentence = question_cut(sentence)
    sentence = strip(sentence)
    sentence = rm_stop_word(sentence)
    candidate = set()
    for word in sentence:
        if word in invertedList:
            candidate = candidate | invertedList[word]
    return candidate

len(get_retrieve_result('最快什么时候可以发货'))  # 通过倒排表返回文档IDs