import gradio as gr
import tensorflow as tf
from tensorflow.keras.preprocessing.text import Tokenizer
from tensorflow.keras.preprocessing.sequence import pad_sequences

# Example jokes data (for training purposes)
jokes = [
    "Why don't scientists trust atoms? Because they make up everything!",
    "What do you call a fish with no eyes? Fsh.",
    "Why did the bicycle fall over? Because it was two tired!",
    "What did the left eye say to the right eye? Between you and me, something smells!",
    "Parallel lines have so much in common. It's a shame they'll never meet.",
    "I used to hate facial hair, but then it grew on me.",
    "Why did the scarecrow win an award? Because he was outstanding in his field!",
    "What do you call a lazy kangaroo? Pouch potato!",
    "I'm reading a book about anti-gravity. It's impossible to put down!",
    "What do you call a bear with no teeth? A gummy bear!",
    "Why did the coffee go to the police? It got mugged.",
    "I'm friends with 25 letters of the alphabet. I don't know why.",
    "Why did the golfer wear two pairs of pants? In case he got a hole in one.",
    "I told my wife she was drawing her eyebrows too high. She seemed surprised."
]

# 1. Tokenize the jokes
tokenizer = Tokenizer()
tokenizer.fit_on_texts(jokes)
total_words = len(tokenizer.word_index) + 1

# 2. Create input sequences for training the model
input_sequences = []
for joke in jokes:
    token_list = tokenizer.texts_to_sequences([joke])
    for i in range(1, len(token_list[0])):  # Corrected loop range (token_list[0] instead of token_list)
        n_gram_sequence = token_list[0][:i+1]  # Create n-gram sequence
        input_sequences.append(n_gram_sequence)

# 3. Pad sequences and create predictors and labels
max_sequence_len = max([len(x) for x in input_sequences])
input_sequences = pad_sequences(input_sequences, maxlen=max_sequence_len, padding='pre')
xs, labels = input_sequences[:,:-1], input_sequences[:,-1]

# 4. Convert labels to categorical
ys = tf.keras.utils.to_categorical(labels, num_classes=total_words)

# 5. Build the model
model = tf.keras.Sequential([
    tf.keras.layers.Embedding(total_words, 100, input_length=max_sequence_len-1),
    tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(150)),
    tf.keras.layers.Dense(total_words, activation='softmax')
])
model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

# 6. Train the model (this might take some time depending on your data and machine)
model.fit(xs, ys, epochs=100, verbose=1)

# 7. Function to generate jokes
def generate_joke(seed_text, next_words=10):
    for _ in range(next_words):
        token_list = tokenizer.texts_to_sequences([seed_text])
        token_list = pad_sequences(token_list, maxlen=max_sequence_len-1, padding='pre')  # Pad the sequence
        predicted = model.predict(token_list, verbose=0)  # Predict the next word
        predicted_index = tf.math.argmax(predicted, axis=1).numpy()[0]  # Get the predicted word index
        
        output_word = ""
        for word, index in tokenizer.word_index.items():
            if index == predicted_index:
                output_word = word
                break
        
        seed_text += " " + output_word
    return seed_text

# 8. Gradio Interface (Frontend)
def gradio_interface(selected_category):
    if selected_category == "Animal Jokes":
        prompt = "Tell me a funny joke about animals."
    elif selected_category == "Technology Jokes":
        prompt = "Give me a funny technology joke."
    elif selected_category == "Dad Jokes":
        prompt = "Tell me a dad joke."
    else:
        prompt = "Give me a random joke."

    # Generate the joke using the model
    return generate_joke(prompt, next_words=10)

# 9. Gradio Interface
iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Dropdown(choices=["Animal Jokes", "Technology Jokes", "Dad Jokes", "Random Jokes"], label="Select Joke Category"),
    outputs="text",
    title="AI-Powered Joke Generator",
    description="Click the button to generate a funny joke from the selected category using AI!",
)

# Launch the Gradio app
iface.launch(share=True)