import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.preprocessing.text import Tokenizer


# Function to convert sequences back to text
def sequences_to_text(sequences: list, tokenizer) -> list:
    """
    Converts sequences of token indices back into their corresponding text representation.

    Parameters:
    sequences (list): A list of sequences, where each sequence is a list of token indices.
    tokenizer (Tokenizer): The tokenizer that was used for converting text to sequences.

    Returns:
    list: A list of strings, where each string represents the decoded text of the corresponding input sequence.
    """
    # Create a reverse mapping of the tokenizer's word index to convert indices back to words
    reverse_word_index = dict(
        [(value, key) for (key, value) in tokenizer.word_index.items()]
    )

    # Add a special token for padding indices
    reverse_word_index[0] = "<PAD>"  # For padding indices

    # Initialize a list to store the converted text
    texts = []

    # Iterate over each sequence in the input
    for seq in sequences:
        # Convert each index in the sequence back to the corresponding word
        words = [reverse_word_index.get(i, "?") for i in seq]

        # Join words into a single string, remove padding tokens, and strip any extra spaces
        texts.append(" ".join(words).replace("<PAD>", "").strip())

    # Return the list of decoded texts
    return texts


def predict_text(
    input_text: str,
    transformer: tf.keras.Model,
    tokenizer: Tokenizer,
    max_length: int = 10,
) -> str:
    """
    Predict the response from the Transformer model for a given input text.

    Args:
        input_text (str): The input string to predict.
        transformer (tf.keras.Model): The trained Transformer model.
        tokenizer (Tokenizer): The tokenizer used for training.
        max_length (int, optional): The maximum length for padding sequences. Default is 10.

    Returns:
        str: The predicted text output generated by the Transformer model.
    """
    # Step 1: Tokenize the input text to convert it into a sequence of integers
    input_sequence = tokenizer.texts_to_sequences([input_text])

    # Step 2: Pad the sequence to ensure it matches the model's expected input size
    # This helps in handling varying input lengths
    input_padded = pad_sequences(input_sequence, maxlen=max_length, padding="post")

    # Step 3: Make a prediction using the Transformer model
    # The model outputs a probability distribution over the vocabulary for each position in the sequence
    predicted_output = transformer.predict(input_padded)

    # Step 4: Convert the predicted output back to text
    # Use argmax to select the index with the highest probability as the predicted word
    predicted_sequence = np.argmax(predicted_output, axis=-1)

    # Convert the sequence of indices back to readable text using the tokenizer
    predicted_text = sequences_to_text(predicted_sequence, tokenizer)

    # Return the first element of the predicted text as a single string
    return predicted_text[0]


# Note: `sequences_to_text` function should be defined elsewhere in your code to convert sequences back to text
