# re imports
import re

# numpy imports 
import numpy as np

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

# inflect imports
import inflect
# create the NUMBER2TEXT object which helps us to convert each numerical text to its textual represantation
NUMBER_TO_TEXT = inflect.engine()

# utils imports
from utils.vocabulary import Vocabulary
# nltk imports
import nltk
nltk.download("punkt")

class TextHandler:
    def __init__(self, clean:bool=False, use_sep:bool=True):
        """ Text Hanlder class we used to pre-process our captions.
        The steps are provided in my Thesis.

        Args:
            clean (bool, optional): If we want to clean our text from special words like x-XXXX. Defaults to False.
            use_sep (bool, optional): If we want to separate our sentences with a SEQ_SEP token. Defaults to True.
        """
        self.__clean = clean
        self.__start_token = "startsequence"
        self.__end_token = "endsequence"
        self.__seq_sep = None
        if use_sep:
            self.__seq_sep = " endofsequence "

    def get_basic_token(self) -> tuple[str, str, str]:
        """ Returns the start, end, and seq_sep special tokens

        Returns:
            tuple[str, str, str]: start, end, and seq_sep tokens
        """
        return self.__start_token, self.__end_token, self.__seq_sep
    
    def remove_punctuations(self, text:str) -> str:
        """ Removes punctuations from training captions as well as cpecial characters

        Args:
            text (str): Text to pre-process

        Returns:
            str: Pre-processed text, without punctuation
        """
        return re.sub(r"[-()\"#/@;:<>{}`+=~|!.?$%^&*'/+\[\]_]+", "", text)
    
    def num2words(self, text:str) -> str:
        """This function converts each numerical text to its textual represantation. Like 10 to ten, and not onezero.

        Args:
            text (str): Text to pre-process

        Returns:
            str: Pre-processed text, with textual numbers
        """
        sentences = text.split('.')
        new_seqs = list()
        # get all sequences
        for s in sentences:
            tokens = s.split()
            new_tokens = list()
            # for each seq, get all words
            for token in tokens:
                # find the number
                try:
                    number = int(token)
                    # convert to each textual represantion. This also converts 10 to ten, and not onezero
                    word = NUMBER_TO_TEXT.number_to_words(number)
                except:
                    word = token
                new_tokens.append(word)
            new_seqs.append(' '.join(new_tokens))

        # connect again whole sentence
        modified_text = '. '.join(new_seqs)
        return modified_text
            

    def __preprocess_text(self, text:str) -> str:
        """ Exetures the pre-processed steps. More details are provided in my Thesis

        Args:
            text (str): Text to pre-process

        Returns:
            str: Pre-processed text. 
        """
        text = re.sub(r"won\'t", "will not", text)
        text = re.sub(r"can\'t", "can not", text)
        text = re.sub(r"n\'t", " not", text)
        text = re.sub(r"\'re", " are", text)
        text = re.sub(r"\'s", " is", text)
        text = re.sub(r"\'d", " would", text)
        text = re.sub(r"\'ll", " will", text)
        text = re.sub(r"\'t", " not", text)
        text = re.sub(r"\'ve", " have", text)
        text = re.sub(r"\'m", " am", text)
        text = re.sub("&", "and", text)
        text = re.sub("@", "at", text)
        text = re.sub("year old", "", text)
        text = re.sub("yearold", "", text)
        
        text = self.num2words(text)
        
        if self.__clean:
            text = self.__clean_text(text)

        text = text.strip().lower()
        text = " ".join(text.split())  # removes unwanted spaces
        if text == "":
            text = np.nan

        return text

    def __clean_text(self, text:str) -> str:
        """ This function cleans the text from special words.

        Args:
            text (str): Text to pre-process

        Returns:
            str: Pre-processed text, without special wortds.
        """
        regex = r"\d."
        text = re.sub(regex, "", text)

        regex = r"X+"
        text = re.sub(regex, "", text)

        regex = r"[^.a-zA-Z]"
        text = re.sub(regex, " ", text)

        regex = r"http\S+"
        text = re.sub(regex, "", text)

        return text

    def separate_sequences(self, text:str) -> str:
        """ This function reads a sequence of texts and appends a SEQ_SEP token between sentences, for better training.
        More details are provided in my Thesis

        Args:
            text (str): Text to pre-process

        Returns:
            str: Pre-processed text, with SEQ_SEP special token.
        """
        start, end, seq_sep = self.get_basic_token()
        if seq_sep is not None:
            sequences = nltk.tokenize.sent_tokenize(text)
            sequences = [s for s in sequences if len(s) > 5]
            text = seq_sep.join(sequences)
            text = self.remove_punctuations(text)
        return start + " " + text + " " + end

    def preprocess_all(self, texts:list) -> list:
        """ Begins the pre-processing for a list of texts.

        Args:
            texts (list): All texts in which we want to apply the pre-process.

        Returns:
            list: Pre-processed texts
        """
        preprocessed_texts = [self.__preprocess_text(text) for text in texts]
        separated_texts = [self.separate_sequences(text) for text in preprocessed_texts]
        return separated_texts