# !/usr/bin/env python
# encoding: utf-8


"""
  @author: gaogao
  @file: search_engine_base.py
  @time: 2021/6/28 11:14
  @desc:
"""
import re
import pylru


class LRUCache:
    def __init__(self, size=32):
        self.cache = pylru.lrucache(size)

    def has(self, key):
        return key in self.cache

    def get(self, key):
        return self.cache[key]

    def set(self, key, value):
        self.cache[key] = value


class SearchEngineBase:
    def __init__(self):
        pass

    def add_corpus(self, file_path):
        with open(file_path, 'r') as fin:
            text = fin.read()
        self.process_corpus(file_path, text)

    def process_corpus(self, id, text):
        raise Exception("process_corpus not implemented.")

    def search(self, query):
        raise Exception("search not implemented.")


class SimpleEngine(SearchEngineBase):
    def __init__(self):
        super().__init__()
        self.__id_to_texts = {}

    def process_corpus(self, id, text):
        self.__id_to_texts[id] = text

    def search(self, query):
        results = []
        for id, text in self.__id_to_texts.items():
            if query in text:
                results.append(id)
        return results


class BOWEngine(SearchEngineBase):
    def __init__(self):
        super().__init__()
        self.__id_to_words = {}

    def process_corpus(self, id, text):
        self.__id_to_words[id] = self.parse_text_to_words(text)

    def search(self, query):
        query_words = self.parse_text_to_words(query)
        results = []
        for id, words in self.__id_to_words.items():
            if self.query_match(query_words, words):
                results.append(id)
        return results

    @staticmethod
    def query_match(query_words, words):
        for query_word in query_words:
            if query_word not in words:
                return False
        return True

    @staticmethod
    def parse_text_to_words(text):
        # 使用正则表达式去掉标点符号和换行符
        text = re.sub(r'[^\w]', ' ', text)
        # 转为小写
        text = text.lower()
        # 生成所有单词的列表
        word_list = text.split(' ')
        # 去除空白单词
        word_list = filter(None, word_list)
        # 返回单词的set
        return set(word_list)


class BOWInvertedIndexEngine(SearchEngineBase):
    def __init__(self):
        super(BOWInvertedIndexEngine, self).__init__()
        self.inverted_index = {}

    def process_corpus(self, id, text):
        words = self.parse_text_to_words(text)
        for word in words:
            if word not in self.inverted_index:
                self.inverted_index[word] = []
            self.inverted_index[word].append(id)
        print("self.inverted_index", self.inverted_index)

    def mergeKLists(self, lists):
        """
        :type lists: List[ListNode]
        :rtype: ListNode
        """
        print("list", lists)
        heap = []
        for node in lists:
            if node:
                heap.append((node.val, node))  # 堆中放入tuple：值，地址
        heapq.heapify(heap)  # 做成堆
        dummy = ListNode(0)
        curr = dummy

        while heap:
            pop = heapq.heappop(heap)  # 从堆中取最小的值
            curr.next = ListNode(pop[0])  # 将值放入
            curr = curr.next
            if pop[1].next:  # 如果该链表没结束
                heapq.heappush(heap, (pop[1].next.val, pop[1].next))  # 将该链表下个节点压入栈中
        return dummy.next

    def search(self, query):
        query_words = list(self.parse_text_to_words(query))
        print("query_words", query_words)
        query_words_index = []
        for query_word in query_words:
            query_words_index.append(0)
        # 如果某一个查询单词的倒叙索引为空，我们就立即返回
        for query_word in query_words:
            if query_word not in self.inverted_index:
                return []
        result = []

        while True:
            # 首先，获得当前状态下所有倒叙索引的index
            current_ids = []
            for idx, query_word in enumerate(query_words):
                current_index = query_words_index[idx]
                current_inverted_list = self.inverted_index[query_word]

                # 已经遍历到了某一个倒序索引的末尾，结束search
                if current_index >= len(current_inverted_list):
                    print("result", result)
                    return result
                current_ids.append(current_inverted_list[current_index])
                print("current_ids", current_ids)
                # 然后，如果current_ids的所有元素都一样,那么表明这个单词在这个元素对应的文档中都出现了。
                if all(x == current_ids[0] for x in current_ids):
                    result.append(current_ids[0])
                    query_words_index = [x + 1 for x in query_words_index]
                    continue
                # 如果不是，我们就把最小的元素加一
                min_val = min(current_ids)
                min_val_pos = current_ids.index(min_val)
                query_words_index[min_val_pos] += 1

    @staticmethod
    def parse_text_to_words(text):
        # 使用正则表达式去掉标点符号和换行符
        text = re.sub(r'[^\w]', ' ', text)
        # 转为小写
        text = text.lower()
        # 生成所有单词的列表
        word_list = text.split(' ')
        # 去除空白单词
        word_list = filter(None, word_list)
        # 返回单词的set
        return set(word_list)


class BOWInvertedIndexEngineWithCache(BOWInvertedIndexEngine, LRUCache):
    def __init__(self):
        super(BOWInvertedIndexEngineWithCache, self).__init__()
        LRUCache.__init__(self)

    def search(self, query):
        if self.has(query):
            print("cache hit!")
            return self.get(query)

        result = super(BOWInvertedIndexEngineWithCache, self).search(query)
        self.set(query, result)
        return result


def main(search_engine):
    print("se", search_engine)
    for file_path in [r'1.txt', r'2.txt', r'3.txt', r'4.txt', r'5.txt']:
        print("file_path", file_path)
        search_engine.add_corpus(file_path)
    while True:
        query = input()
        results = search_engine.search(query)
        print("found {} result(s):".format(len(results)))
        for result in results:
            print(result)


search_engine = BOWInvertedIndexEngineWithCache()
main(search_engine)
