# tensorflow imports
import tensorflow
from tensorflow.keras.preprocessing.text import Tokenizer


class Vocabulary:
    def __init__(self, texts:list, threshold:int = 3):
        """ Vocabulary class we utilised in CNN-RNN models. This class creates a vocabulary for our captions with a cut-off frequency.

        Args:
            texts (list): All training captions from which we extract our vocabulary
            threshold (int, optional): The cut-off frequency in our dictionary. Defaults to 3.
        """
        self.texts = texts
        # add <pad> abd <unk> tokens in vocabulary. <unk> refer to the OOV (out-of-vocabulary) tokens.
        self.pad_token = "<pad>"
        self.unk_token = "<unk>"
        self.threshold = threshold
        # init the tokenizer
        self.tokenizer = Tokenizer(oov_token="<unk>")
        # dictionaries
        self.word2idx = {}
        self.idx2word = {}

    def build_vocab(self) ->tuple[Tokenizer, dict, dict]:
        """ This function creates the Vocabulary we want to employ for our CNN-RNN model.

        Returns:
            tuple[Tokenizer, dict, dict]: The fitted tokenizer, the word-to-idx dictionary, and idx-to-word dictionary. 
            The latters are mappers for words and index respectively
        """
        # fit in training captions
        self.tokenizer.fit_on_texts(self.texts)
        # create the vocabulary in sorted way
        sorted_vocab = dict(
            sorted(self.tokenizer.word_counts.items(), key=lambda t: t[1], reverse=True)
        )

        word_index_threshold, index_word_threshold = {}, {}

        # add pad and unk token to tokenizer dictionary
        word_index_threshold[self.pad_token] = 0
        index_word_threshold[0] = self.pad_token
        word_index_threshold[self.unk_token] = 1
        index_word_threshold[1] = self.unk_token

        # begin from index=2
        idx = 2
        for k, v in sorted_vocab.items():
            # if the occurence of current word is less than the threshold, continue.
            if v >= self.threshold:
                word_index_threshold[k] = idx
                index_word_threshold[idx] = k
                idx += 1

        # get the dictionaries
        self.tokenizer.word_index = word_index_threshold
        self.tokenizer.index_word = index_word_threshold

        dictionary = self.tokenizer.word_index

        # append to the global ones
        for k, v in dictionary.items():
            self.word2idx[k] = v
            self.idx2word[v] = k

        print(f"Made a vocabulary with {len(self.word2idx)} words!")
        
        return self.tokenizer, self.word2idx, self.idx2word