"""
Script containing the class to construct the index. It uses the SPIMI (Single PAss In-Memory Index) Algorithm
"""

import json
import sys
import shutil
import os
from collections import OrderedDict

try:
    import cPickle as pickle
except:
    import pickle


class IndexConstructor:
    """
    Index Contructor that gets a token stream and block size in bytes and then creates the index
    based on the SPIMI alogrithm
    """

    def __init__(self, token_stream, block_size=10240):
        """
        Constructor receiving a token stream and the block size in bytes
        :param token_stream: String with pairs of <term and doc_id>
        :param block_size: The block size for the SPIMI algorithm to use in bytes
        """
        self.token_stream = token_stream
        self.block_size = block_size

        # To keep track of the no of blocks and for the file name of the block
        self.block_no = 1

        # Path to save the block file
        self.tmp_file_dir_path = os.path.join(os.path.abspath(os.path.dirname(os.path.relpath(__file__))),
                                              "tmp_block_dir")

        # Path to save the final inverted index
        self.inverted_index_path = os.path.join(os.path.relpath(__file__), "..", "..", "parsed_data", "inverted_index.bin")

        # Inverted index will be an ordered dictionary
        self.inverted_index = OrderedDict()

    def construct_index(self, clear_tmp=True, get_stats=False, save_json=False):
        """
        Main method to construct the inverted index by a sequence of steps
        :param clear_tmp: True to clear the temp blocks, False not to clear
        :param get_stats: True to get stats on the inverted index, False otherwise
        :param save_json: True to save the inverted index in JSON
        :return: the path to the complete inverted index
        """

        print("Making new directory for blocks' file")
        self.__clear_tmp_block_dir()
        self.__make_tmp_block_dir()

        print("Inverting the block files with SPIMI")
        self.spimi_invert()

        print("Merging the blocks to make index")
        self.merge_blocks()

        if clear_tmp:
            print("Clearing temporary folder for the blocks file")
            self.__clear_tmp_block_dir()
            self.__dump_inverted_index()

        if save_json:
            self.save_index_json()

        if get_stats:
            self.get_stats()

        return self.inverted_index_path

    def spimi_invert(self):
        """
        SPIMI Algorithm to create the inverted index using blocks.
        - The algorithm will pass through the token stream
        - For each term, if it does not exist for a new dict for each block, add it to the dict
        - If term exists in dict, add doc id to the term's postings list
        - The algorihtm checks that the dictionary does not grow more thatn the block size in bytes
        - If this occurs, the block is saved in memory
        :return: None
        """

        list_idx = 0

        # Make sure to go through all tokens
        while self.__check_idx(list_idx):

            block_dict = dict()

            # Make sure to respect block size and that list indexing does not overflow
            while self.__get_dump_size(block_dict) < self.block_size and self.__check_idx(list_idx):

                token_pair = self.token_stream[list_idx]
                term, doc_id = token_pair[0], int(token_pair[1])
                term_posting_list = []
                if term not in block_dict:
                    block_dict[term] = term_posting_list
                else:
                    term_posting_list = block_dict[term]

                if doc_id not in term_posting_list:
                    term_posting_list.insert(len(term_posting_list), doc_id)

                list_idx += 1

            sorted_terms = sorted(block_dict)

            self.save_block_data(sorted_terms, block_dict)
            self.block_no += 1

    def merge_blocks(self):
        """
        Method to merge the different blocks generated by the SPIMI Algorithm into the complete inverted index
        - The algorithm goes from the smallest to the largest block (in bytes)
        :return: None
        """

        # Get the block files in ascending order of size
        file_list = self.get_sorted_block_files()
        tmp_inverted_idx = dict()

        for block_file in file_list:

            # Open file and merge with in memory inverted index
            try:
                with open(os.path.join(self.tmp_file_dir_path, block_file[1]), "rb") as file_obj:
                    sorted_terms = pickle.load(file_obj)
                    block_dict = pickle.load(file_obj)
            except (IOError, OSError):
                print("Unable to load block file")
                exit(1)
            else:
                for term in sorted_terms:
                    # If new term in final index, add to it, if not merge the postings list
                    if term not in tmp_inverted_idx:
                        tmp_inverted_idx[term] = block_dict[term]
                    else:
                        tmp_inverted_idx[term] = self.merge_postings_list(tmp_inverted_idx[term], block_dict[term])

        # Get sorted tmp index and store
        for term in sorted(tmp_inverted_idx):
            self.inverted_index[term] = tmp_inverted_idx[term]

    def get_sorted_block_files(self):
        """
        Method to get block files in ascending order of size
        :return: dictionary with file_no: tmp_file_path (block)
        """

        files = os.listdir(self.tmp_file_dir_path)
        file_pairs = list()
        for file in files:
            file_path = os.path.join(self.tmp_file_dir_path, file)
            # Store as a tuple pair
            file_pairs.append((os.path.getsize(file_path), file))

        # Sort by the size
        file_pairs.sort(key=lambda pair: pair[0])

        return file_pairs

    def merge_postings_list(self, first_list, second_list):
        """
        Method to merge two postings lists
        :param first_list: First list to merge
        :param second_list: Second list to merge
        :return:
        """

        iter_first_list = iter(first_list)
        iter_scnd_list = iter(second_list)

        first_list_post = next(iter_first_list, None)
        scnd_list_post = next(iter_scnd_list, None)

        merged_list = list()
        # Merge the lists linearly by keeping an iterator on each list
        while first_list_post is not None and scnd_list_post is not None:

            if first_list_post == scnd_list_post:
                merged_list.append(first_list_post)
                first_list_post = next(iter_first_list, None)
                scnd_list_post = next(iter_scnd_list, None)

            elif first_list_post < scnd_list_post:
                merged_list.append(first_list_post)
                first_list_post = next(iter_first_list, None)

            else:
                merged_list.append(scnd_list_post)
                scnd_list_post = next(iter_scnd_list, None)

        # Merge remaining items in list
        if first_list_post is not None:
            merged_list += first_list[first_list.index(first_list_post):]
        elif scnd_list_post is not None:
            merged_list += second_list[second_list.index(scnd_list_post):]

        return merged_list

    def save_block_data(self, sorted_keys, block_dict):
        """
        Method to save datas in a block file
        :param sorted_keys: Sorted terms
        :param block_dict: The entire block dict with the postings list
        :return: None
        """

        block_path = os.path.join(self.tmp_file_dir_path, "block_%05d.bin" % self.block_no)

        try:
            with open(block_path, "wb") as tmp_file:
                self.__dump_pickle(sorted_keys, tmp_file)
                self.__dump_pickle(block_dict, tmp_file)
        except (IOError, OSError):
            print("Error saving block file")
            exit(1)

    def get_stats(self):
        """
        Method to print stats on the inverted index
        :return: None
        """

        print("No. of distinct terms: %d" % len(self.inverted_index))

        postings_count = sum(len(post_list) for post_list in self.inverted_index.values())
        print("No. of nonpositional postings: %d" % postings_count)

    def save_index_json(self):
        """
        Method to save the inverted index as JSON
        :return: None
        """
        try:
            with open(os.path.join(os.path.relpath(__file__), "..", "..", "parsed_data", "inverted_index.json"), "w") as dump_file:
                json.dump(self.inverted_index, dump_file)
        except (IOError, OSError):
            print("Unable to dump inverted index (JSON)")
            exit(1)

    def __clear_tmp_block_dir(self):
        """
        Helper method to clear the temporary directory to save blocks
        :return: None
        """

        if os.path.exists(self.tmp_file_dir_path):
            shutil.rmtree(self.tmp_file_dir_path)

    def __make_tmp_block_dir(self):
        """
        Helper method to make a temporary directory for the blocks
        :return: None
        """
        os.makedirs(self.tmp_file_dir_path)

    def __dump_pickle(self, obj, file_obj):
        """
        Method to save two objs toa binary file
        :param obj: first object to save (will be the sorted terms for the invered index)
        :param file_obj: second obj (will be the whole inverted index)
        :return: None
        """
        pickle.dump(obj=obj, file=file_obj, protocol=pickle.HIGHEST_PROTOCOL)

    def __get_dump_size(self, dict_obj):
        """
        Method to get calculate size of dictionary if dumped to disk (n bytes)
        :param dict_obj:  Dictionary to calculate size
        :return: Size in bytes of dictionary if saved to disk
        """
        terms_size = sys.getsizeof(pickle.dumps(list(dict_obj.keys())))
        whole_dict_size = sys.getsizeof(pickle.dumps(dict_obj, pickle.HIGHEST_PROTOCOL))
        return terms_size + whole_dict_size

    def __check_idx(self, list_idx):
        """
        Method to check the index of list compared to len of the token stream
        :param list_idx: List index to check
        :return: True if index is less than len of token stream, False otherwise
        """
        return list_idx < len(self.token_stream)

    def __dump_inverted_index(self):
        """
        Helper method to dump the complete inverted index
        :return: None
        """
        try:
            with open(self.inverted_index_path, "wb") as dump_file:
                self.__dump_pickle(self.inverted_index, dump_file)
        except (IOError, OSError):
            print("Unable to dump inverted index")
            exit(1)
