# numpy import

import numpy as np

# tensorflow imports

import tensorflow as tf

import tensorflow

from tensorflow.keras.preprocessing.text import Tokenizer

from tensorflow.keras.models import Model

class BeamSearch:

    def __init__(self, start_token:str, end_token:str, max_length:int, tokenizer:Tokenizer, idx_to_word:dict, word_to_idx:dict, beam_index:int):

        """ The Beam Search sampling method for generating captions. An illustration of the algorithm is provided in my Thesis paper.

        Args:

            start_token (str): The start-token used during pre-processing of the training captions

            end_token (str): The end-token used during pre-processing of the training captions

            max_length (int): The maximum length (limit) for the generated captions

            tokenizer (Tokenizer): The fitted tokenizer from the Vocabulary object

            idx_to_word (dict): Dictionary with keys to be the index number and values the words in the created vocabulary

            word_to_idx (dict): Dictionary with keys to be the words and values the index number in the created vocabulary  

            beam_index (int): The beam size for the Beam Seach algorithm.

        """

        self.start_token = start_token

        self.end_token = end_token

        self.max_length = max_length

        self.tokenizer = tokenizer

        self.idx_to_word = idx_to_word

        self.word_to_idx = word_to_idx

        self.beam_index = beam_index

    def get_word(self, idx:int) -> str:

        """ Fetches the word from the index-to-word vocab, which was created after the pre-processing of the Training captions

        Args:

            idx (int): The index for the index-to-word vocab.

        Returns:

            str: The word for the given index if exist in the created index-to-word vocab, else None

        """

        return self.idx_to_word.get(idx, None)

    def get_idx(self, word:str)->int:

        """ Fetches the index number from the word-to-index vocab, which was created after the pre-processing of the Training captions

        Args:

            word (str): The word for which we want its index in the word-to-index dictionary.

        Returns:

            int: The index for the given word if exist in the created word-to-index vocab, else -1. The latter number refer to None

        """

        return self.word_to_idx.get(word, -1)

    def beam_search_predict(self, model:Model, image:np.array, tag:np.array, dataset:str='iuxray', multi_modal:bool=False)->str:

        """ Executes the beam search algorithm employing the pre-trained model along with the test instance's data.

        Args:

            model (Model): The model we want to evaluate on our employed dataset

            image (np.array): Current test image embedding

            tag (np.array): The tag embedding for the current test instance. This is used only for IU X-Ray dataset.

            dataset (str, optional): The dataset we employed for the model. Defaults to 'iuxray'.

            multi_modal (bool, optional): If we want to use the multi-modal version of model. This is used only for IU X-Ray dataset. Defaults to False.

        Returns:

            str: The generated description for the given image using the beam search.

        """

        start = [self.get_idx(self.start_token)]

        start_word = [[start, 0.0]]

        while len(start_word[0][0]) < self.max_length:

            # store current word,probs pairs

            temp = []

            # for current sequence

            for s in start_word:

                # pad the sequence in order to fetch the next token

                par_caps = tf.keras.preprocessing.sequence.pad_sequences([s[0]], maxlen=self.max_length, padding="post")

                if multi_modal:

                    if dataset == 'iuxray':

                        preds = model.predict(

                            [image[0], image[1], tag, par_caps], verbose=0)

                    else:

                        preds = model.predict(

                            [image, tag, par_caps], verbose=0)

                else:

                    if dataset == 'iuxray':

                        preds = model.predict(

                            [image[0], image[1], par_caps], verbose=0)

                    else:

                        preds = model.predict([image, par_caps], verbose=0)

                # get the best paths

                word_preds = np.argsort(preds[0])[-self.beam_index:]

                # Getting the top <self.self.beam_index>(n) predictions and creating a

                # new list so as to put them via the model again

                for w in word_preds:

                    next_cap, prob = s[0][:], s[1]

                    next_cap.append(w)

                    prob += preds[0][w]

                    temp.append([next_cap, prob])

            start_word = temp

            # Sorting according to the probabilities

            start_word = sorted(start_word, reverse=False, key=lambda l: l[1])

            # Getting the top words

            start_word = start_word[-self.beam_index:]

        # get the best path

        start_word = start_word[-1][0]

        intermediate_caption = [self.get_word(i) for i in start_word]

        final_caption = []

        for i in intermediate_caption:

            if i != self.end_token:

                final_caption.append(i)

            else:

                break

        final_caption = " ".join(final_caption[1:])

        return final_caption

    def ensemble_beam_search(self, models:list, images_list:list)->str:

        """ Executes the beam search algorithm employing the pre-trained models along with the test instances data.

        This utilises the beam search algorithm for each model in models list.

        Args:

            models (list): The models we want to evaluate on our employed dataset

            images_list (list): Current test images embeddings for each encoder we used.

        Returns:

            str: The generated description for the given image using the beam search.

        """

        start = [self.get_idx(self.start_token)]

        start_word = [[start, 0.0]]

        while len(start_word[0][0]) < self.max_length:

            # for current seq

            for s in start_word:

                # pad current caption

                current_caption = tf.keras.preprocessing.sequence.pad_sequences([s[0]], maxlen=self.max_length, padding="post")

                # get all predictions from the pre-trained models

                ensemble_predictions = [ ensemble_member.predict([image, current_caption], verbose=0) for ensemble_member, image in zip(models, images_list) ]

                # get the best pairs

                ensemble_word_predictions = [ np.argsort(prediction[0])[-self.beam_index:] for prediction in ensemble_predictions ]

                # and store them with word,pairs for each ensemble member

                ensemble_current_probs = list()

                for member_prediction, member_word_predictions in zip(ensemble_predictions, ensemble_word_predictions):

                    temp_current_seq = list()

                    for word in member_word_predictions:

                        next_cap, prob = s[0][:], s[1]

                        next_cap.append(word)

                        prob += member_prediction[0][word]

                        temp_current_seq.append([next_cap, prob])

                    ensemble_current_probs.append(temp_current_seq)

            # get all the best candidates

            ensemble_starting_words = [

                sorted(current_probs, reverse=False, key=lambda l: l[1])[-self.beam_index:] for current_probs in ensemble_current_probs

            ]

            start_words = [member_starting_words[-1] for member_starting_words in ensemble_starting_words]

            start_word = sorted(start_words, reverse=False, key=lambda l: l[1])

        # create the caption

        start_word = start_word[-1][0]

        intermediate_caption = [self.get_word(i) for i in start_word]

        final_caption = []

        for i in intermediate_caption:

            if i != self.end_token:

                final_caption.append(i)

            else:

                break

        final_caption = " ".join(final_caption[1:])

        return final_caption