def read_corpus():
    """
    读取给定的语料库，并把问题列表和答案列表分别写入到 qlist, alist 里面。 在此过程中，不用对字符换做任何的处理（这部分需要在 Part 2.3里处理）
    qlist = ["问题1"， “问题2”， “问题3” ....]
    alist = ["答案1", "答案2", "答案3" ....]
    务必要让每一个问题和答案对应起来（下标位置一致）
    """
    # TODO 需要完成的代码部分 ...
    import json
    import jsonpath
    
    with open('train-v2.0.json', 'r', encoding='utf-8') as f:
        dic = json.load(f)
    
    qlist = jsonpath.jsonpath(dic,'$..question') 
    alist = jsonpath.jsonpath(dic,'$..text') 
    f.close()
    assert len(qlist) == len(alist)  # 确保长度一样
    
    return qlist, alist

# TODO: 统计一下在qlist中总共出现了多少个单词？ 总共出现了多少个不同的单词(unique word)？
#       这里需要做简单的分词，对于英文我们根据空格来分词即可，其他过滤暂不考虑（只需分词）
from nltk import word_tokenize

qlist, alist = read_corpus()

qwords = []
for sentence in qlist:
    qwords.append(word_tokenize(sentence))

mydist = {}

for sentence in qwords:
    for word in sentence:
        if word not in mydist:
            mydist[word] =1
        else:
            mydist[word] +=1

word_total = len(mydist)
print (word_total)

from nltk import word_tokenize
from nltk import FreqDist
%matplotlib inline

fdist = FreqDist(mydist)
fdist.plot(20, cumulative=True)

# TODO： 从上面的图中能观察到什么样的现象？ 这样的一个图的形状跟一个非常著名的函数形状很类似，能所出此定理吗？ 
#       hint: [XXX]'s law
# 
#  zip's law

# TODO： 需要做文本方面的处理。 从上述几个常用的方法中选择合适的方法给qlist做预处理（不一定要按照上面的顺序，不一定要全部使用）
from nltk.corpus import stopwords
from nltk.stem import WordNetLemmatizer

stop_words = set(stopwords.words('english'))
wnl = WordNetLemmatizer()
f_list = []
english_punctuations = ['(', ')', '[', ']', '&', '!!', '*', '@', '#', '$', '%', '?']

def sentencesList_preprocess(sentencesList):

    for sentence in sentencesList:
        # filter stop words
        filtered = [word for word in sentence if word not in stop_words]
        
        # to lower case
        lower_case = [word.lower() for word in filtered]
        
        # lemmatization
        lemma = [wnl.lemmatize(token) for token in lower_case]
        
        # remove unused symbols
        sym = [word for word in lemma if word not in english_punctuations]
        
        f_list.append(sym)
    
    return f_list

qlist = sentencesList_preprocess(qwords)    # 更新后的问题列表

def sentence_preprocess(sentence):
    # filter stop words
    filtered = [word for word in sentence if word not in stop_words]
    # to lower case
    lower_case = [word.lower() for word in filtered]
    # lemmatization
    lemma = [wnl.lemmatize(token) for token in lower_case]
    # remove unused symbols
    token = [word for word in lemma if word not in english_punctuations]

    return token

# TODO 
from sklearn.feature_extraction.text import TfidfVectorizer

vectorizer = TfidfVectorizer()   # 定义一个tf-idf的vectorizer

qlistnew = [' '.join(sentence) for sentence in qlist]
X_tfidf = vectorizer.fit_transform(qlistnew)  # 结果存放在X矩阵里

# load Glove
def loadGlove(path):
    vocab = {}
    embedding = []
    vocab["UNK"] = 0
    embedding.append([0]*200)
    file = open(path, 'r', encoding='utf8')
    i = 1
    for line in file:
        row = line.strip().split()
        vocab[row[0]] = i
        embedding.append(row[1:])
        i += 1
    print("Finish load Glove")
    file.close()
    return vocab, embedding

path = '../glove.6b/glove.6b.200d.txt'
voc, emb = loadGlove(path)

def sentence_to_vec(embedding, sentence):
    vec = np.zeros((200,), dtype=np.float64)
    for word in sentence:
        if word in voc:
            idx = voc[word]
            vec += embedding[idx].astype('float64')
        else:
            vec += embedding[0].astype('float64')
    vec = vec/len(sentence)
    return vec

def sentences_to_vec(embedding, sentences):
    vec = np.zeros((len(sentences), 200))
    for i, sentence in enumerate(sentences):
        sentence = sentence.strip().split(' ')
        vec[i] = sentence_to_vec(embedding, sentence)
                
    return vec

import numpy as np
emc = np.asarray(emb)

# TODO 基于Glove向量获取句子向量
# emb  =  # 这是 D*H的矩阵，这里的D是词典库的大小， H是词向量的大小。 这里面我们给定的每个单词的词向量，
        # 这需要从文本中读取
    
X_w2v = sentences_to_vec(emc, qlistnew)   # 初始化完emb之后就可以对每一个句子来构建句子向量了，这个过程使用average pooling来实现

# for i in voc:
#     v = emb[v]
#     res = list(cosine_similarity(emb, v)[0])
#     result = sorted(get_top_numbers(res, 10), reverse = True)

import numpy as np
from bert_embedding.bert import BertEmbedding
bert_embedding = BertEmbedding(model='bert_12_768_12', dataset_name='book_corpus_wiki_en_uncased')

def sentences_to_vec_bert(sentences):
    X_bert = np.zeros((len(sentences), 768))
    for i in range(len(sentences)):
        emb = bert_embedding(sentence, 'avg')
        vec_list = emb[0][1]
#         v_bert = np.zeros((768,), dtype=np.float32)
        
#         for vec in vec_list:
#             v_bert += vec.astype('float32')
#         v_bert = v_bert/len(vec_list)
        v_bert = np.average(vec_list, axis=0)
        X_bert[i] = v_bert
    return X_bert

def sentence_to_vec_bert(sentence):
    emb = bert_embedding([sentence], 'avg')
    vec_list = emb[0][1]
    v_bert = np.sum(vec_list, axis=0)
    
    return v_bert

# TODO 基于BERT的句子向量计算

X_bert = sentences_to_vec_bert(qlistnew) # 每一个句子的向量结果存放在X_bert矩阵里。行数为句子的总个数，列数为一个句子embedding大小。 

import heapq
from sklearn.metrics.pairwise import cosine_similarity

# return the top k numbers from the list using prority queue
def get_top_numbers(tlist, k):
    max_heap = []
    l = len(tlist)
    if l<=0 or k<=0 or k>l: return None
    
    for i in tlist:
        if k > len(max_heap): heapq.heappush(max_heap, i)
        else: heapq.heappushpop(max_heap, i)
    
    return max_heap

def get_top_results_tfidf_noindex(query):
    # TODO 需要编写
    """
    给定用户输入的问题 query, 返回最有可能的TOP 5问题。这里面需要做到以下几点：
    1. 对于用户的输入 query 首先做一系列的预处理(上面提到的方法)，然后再转换成tf-idf向量（利用上面的vectorizer)
    2. 计算跟每个库里的问题之间的相似度
    3. 找出相似度最高的top5问题的答案
    """
    query = word_tokenize(query)
    query= sentence_preprocess(query)
    query = ' '.join(query)
    
    q_tfidf = vectorizer.transform([query])
    res = list(cosine_similarity(q_tfidf, X_tfidf)[0])
    result = sorted(get_top_numbers(res, 5), reverse = True)
    
    top_idxs = []  # top_idxs存放相似度最高的（存在qlist里的）问题的下标 
                   # hint: 请使用 priority queue来找出top results. 思考为什么可以这么做？ 
    dict_visited = {}

    for r in result:
        for i, n in enumerate(res):
            if n == r and i not in dict_visited: 
                top_idxs.append(i)
                dict_visited[i] = True

    ans = [alist[i] for i in top_idxs]
    
    return ans  # 返回相似度最高的问题对应的答案，作为TOP5答案    

# TODO: 编写几个测试用例，并输出结果
print (get_top_results_tfidf_noindex("when did Beyonce start becoming popular"))
print (get_top_results_tfidf_noindex("what languge does the word of 'symbiosis' come from"))

# TODO 请创建倒排表
inverted_idx = {}  # 定一个一个简单的倒排表，是一个map结构。 循环所有qlist一遍就可以
for index, sentence in enumerate(qlist):
    for word in sentence:
        if word not in inverted_idx: inverted_idx[word] = [index]
        else: inverted_idx[word].append(index)

# TODO 读取语义相关的单词
def get_related_words(file):
    related_words = {}
    for line in open(file, mode='r', encoding='utf-8'):
        item = line.split(",")
        word, s_list = item[0], [value for value in item[1].strip().split()]
        related_words[word] = s_list
    return related_words

related_words = get_related_words('related_words.txt') # 直接放在文件夹的根目录下，不要修改此路径。

def get_inverted_index_sentence(query):
    query = word_tokenize(query)
    query= sentence_preprocess(query)
    
    r_list = []
    for q in query:
        if q in related_words: 
            for word in related_words[q]:
                r_list.append(word)
    
    total_list = query
    for word in r_list:
        total_list.append(word)
    
    idx_list = []    
    for word in total_list:
        if word in inverted_idx:
            indx = inverted_idx[word]
            idx_list.extend(indx)
    return query, idx_list

def get_top_index(result):
    top_idxs = []
    dict_visited = {}
    top_result = sorted(get_top_numbers(result, 5), reverse = True)
    
    for r in top_result:
        for i, n in enumerate(result):
            if n == r and i not in dict_visited: 
                top_idxs.append(i)
                dict_visited[i] = True
    return top_idxs

def get_top_results_tfidf(query):
    """
    给定用户输入的问题 query, 返回最有可能的TOP 5问题。这里面需要做到以下几点：
    1. 利用倒排表来筛选 candidate （需要使用related_words). 
    2. 对于候选文档，计算跟输入问题之间的相似度
    3. 找出相似度最高的top5问题的答案
    """
    query, sentence_idxs = get_inverted_index_sentence(query)
    query = ' '.join(query)
    q_tfidf = vectorizer.transform([query])
    
    
    X_tfidf_idx = []
    for indx in sentence_idxs:
        X_tfidf_idx.append(X_tfidf[indx].toarray()[0])
    
    res = list(cosine_similarity(q_tfidf, X_tfidf_idx)[0])
    
    top_idxs = []   
    top_idxs = get_top_index(res)
    
    ans = [alist[i] for i in top_idxs]
    
    return ans   # 返回相似度最高的问题对应的答案，作为TOP5答案

test_query1 = "when did Beyonce start becoming popular"
test_query2 = "what languge does the word of symbiosis come from"

print (get_top_results_tfidf(test_query1))
print (get_top_results_tfidf(test_query2))

def get_top_results_w2v(query):
    """
    给定用户输入的问题 query, 返回最有可能的TOP 5问题。这里面需要做到以下几点：
    1. 利用倒排表来筛选 candidate （需要使用related_words). 
    2. 对于候选文档，计算跟输入问题之间的相似度
    3. 找出相似度最高的top5问题的答案
    """
        
    query, sentence_idxs = get_inverted_index_sentence(query)
    query = ' '.join(query)
    q_w2v = sentence_to_vec(emc,query)
    
    
    q_w2v_idx = []
    for indx in sentence_idxs:
        q_w2v_idx.append(X_w2v[indx])
    
    res = list(cosine_similarity(np.array([q_w2v]), np.array([q_w2v_idx])[0])[0])
    
    top_idxs = []   
    top_idxs = get_top_index(res)
    
    ans = [alist[i] for i in top_idxs]
    
    return ans  # 返回相似度最高的问题对应的答案，作为TOP5答案

test_query1 = "when did Beyonce start becoming popular"
test_query2 = "what languge does the word of symbiosis come from"

print (get_top_results_w2v(test_query1))
print (get_top_results_w2v(test_query2))

def get_top_results_bert(query):
    """
    给定用户输入的问题 query, 返回最有可能的TOP 5问题。这里面需要做到以下几点：
    1. 利用倒排表来筛选 candidate （需要使用related_words). 
    2. 对于候选文档，计算跟输入问题之间的相似度
    3. 找出相似度最高的top5问题的答案
    """
    
    query, sentence_idxs = get_inverted_index_sentence(query)
    query = ' '.join(query)
    q_bert = sentence_to_vec_bert(query)
     
    
    
    q_bert_idx = []
    for indx in sentence_idxs:
        q_bert_idx.append(X_bert[indx])
    
    res = list(cosine_similarity(np.array([q_bert]), np.array([q_bert_idx])[0])[0])
    
    top_idxs = []   
    top_idxs = get_top_index(res)
    
    ans = [alist[i] for i in top_idxs]
    
    return ans  # 返回相似度最高的问题对应的答案，作为TOP5答案

test_query1 = "when did Beyonce start becoming popular"
test_query2 = "what language does the word of symbiosis come from"

print (get_top_results_bert(test_query1))
print (get_top_results_bert(test_query2))

# TODO: 编写几个测试用例，并输出结果

test_query1 = "when did Beyonce start becoming popular"
test_query2 = "what languge does the word of symbiosis come from"

print (get_top_results_tfidf(test_query1))
print (get_top_results_w2v(test_query1))
# print (get_top_results_bert(test_query1))

print (get_top_results_tfidf(test_query2))
print (get_top_results_w2v(test_query2))
# print (get_top_results_bert(test_query2))

from nltk.corpus import reuters
import numpy as np
import codecs
# 读取语料库的数据
categories = reuters.categories()
corpus = reuters.sents(categories=categories)

word2index = {}
index2word = {}
corpus2 = []

for sentence in corpus:
    corpus2.append(['<s> '] + sentence + [' </s>'])

for sentence in corpus2:
    for word in sentence:
        word = word.lower()
        if word in word2index: continue
        index2word[len(word2index)] = word
        word2index[word] = len(word2index)

word_count = len(word2index)
uni_count = np.zeros(word_count)
bi_count = np.zeros((word_count, word_count))

for sentence in corpus2:
    for i, word in enumerate(sentence):
        word = word.lower()
        uni_count[word2index[word]] += 1
        if i <len(sentence) -1:
            pre = word2index[word]
            curr = word2index[sentence[i+1].lower()]
            bi_count[pre, curr] +=1

bigram = np.zeros((word_count, word_count))

for i in range(word_count):
    for j in range(word_count):
        if bi_count[i,j]==0:
            bigram[i,j] = 1.0 / (uni_count[i] + word_count)
        else:
            bigram[i,j] = (1.0 + bi_count[i,j]) / (word_count + uni_count[i])

def checkCount(pre,word):
    if pre.lower() in word2index and word.lower() in word2index:
        return bigram[word2index[pre.lower()],word2index[word.lower()]]
    else:
        return 0.0

# TODO 构建channel probability  
channel = {}

spell_error_dict = {}
for line in open('spell-errors.txt'):
    word = line.split(":")
    c_word = word[0] # correct word is the key
    spell_error_dict[c_word] = [e_word.strip( )for e_word in word[1].strip().split(",")]

# TODO
for c_word in spell_error_dict:
    if c_word not in channel:
        channel[c_word] = {}
    for e_word in spell_error_dict[c_word]:
        channel[c_word][e_word] = 1/len(spell_error_dict[c_word])
    
# print(channel)   

alphabet = "abcdefghijklmnopqrstuvwxyz"

def known(words):
    return set(w for w in words if w in word2index)

def edits1(word):
    n = len(word)
    #删除
    s1 = [word[0:i] + word[i+1:] for i in range(n)]
    #调换相连的两个字母
    s2 = [word[0:i] + word[i+1] + word[i] + word[i+2:] for i in range(n-1)]
    #replace
    s3 = [word[0:i] + c + word[i + 1:] for i in range(n) for c in alphabet]
    #插入
    s4 = [word[0:i] + c + word[i:] for i in range(n + 1) for c in alphabet]
    edit1_words = set(s1 + s2 + s3 + s4)

    if word in edit1_words:
        edit1_words.remove(word)

    edit1_words = known(edit1_words)
    return edit1_words

def edits2(word, edit1_words):
    edit2_words = set(e2 for e1 in edit1_words for e2 in edits1(e1))
    
    if word in edit2_words:
        edit2_words.remove(word)
    
    edit2_words = known(edit2_words)
    return edit2_words

def generate_candidates(word):
    # 基于拼写错误的单词，生成跟它的编辑距离为1或者2的单词，并通过词典库的过滤。
    # 只留写法上正确的单词。 
    
    word_edit1 = edits1(word)
    word_edit2 = edits2(word, word_edit1)
    
    words = word_edit1 | word_edit2
    
    return words

import numpy as np
from queue import PriorityQueue
def spell_corrector(line):
    # 1. 首先做分词，然后把``line``表示成``tokens``
    # 2. 循环每一token, 然后判断是否存在词库里。如果不存在就意味着是拼写错误的，需要修正。 
    #    修正的过程就使用上述提到的``noisy channel model``, 然后从而找出最好的修正之后的结果。  
    
    corrected_words = []
    tokens = []
    tokens = ['<s>']+word_tokenize(line)+['</s>']
    for i, token in enumerate(tokens):
        if i == len(tokens)-1: break
        if token.lower() not in word2index:
            pre, nxt = tokens[i-1].lower(), tokens[i+1].lower()
            token = word_corrector(token, pre, nxt)
            corrected_words.append(token)
        else: corrected_words.append(token)
    newline = ' '.join(corrected_words)
    
    return newline   # 修正之后的结果，假如用户输入没有问题，那这时候``newline = line``

def word_corrector(word, pre_word, next_word):
    candidates = generate_candidates(word)
    correctors = PriorityQueue()
    
    if len(candidates) == 0: return word
    
    for candidate in candidates:
        if candidate in channel and word in channel[candidate] and candidate in word2index:
            bi_pre = checkCount(pre_word, candidate)
            bi_nxt = checkCount(candidate, next_word)
            p = np.log(channel[candidate][word] + 0.001) + bi_pre + bi_nxt
            correctors.put((-1*p, candidate))
    
    if correctors.empty(): return word
    
    return correctors.get()[1]

sentence = spell_corrector("when did Beyonce start beeome popular?")
print(sentence)

test_query1 = "What counted for more of the poplation change"  # 拼写错误的
test_query2 = "What counted for more of the population chenge"  # 拼写错误的

test_query1 = spell_corrector(test_query1)
test_query2 = spell_corrector(test_query2)

print(test_query1)
print(test_query2)

test_query1 = "when did Beyonce starte becoming popular?"  # 拼写错误的
test_query2 = "What counted for more of the population chenge"  # 拼写错误的

test_query1 = spell_corector(test_query1)
test_query2 = spell_corector(test_query2)

print (get_top_results_tfidf(test_query1))
print (get_top_results_w2v(test_query1))
print (get_top_results_bert(test_query1))

print (get_top_results_tfidf(test_query2))
print (get_top_results_w2v(test_query2))
print (get_top_results_bert(test_query2))