import gradio as gr
import tensorflow as tf
import numpy as np
from tensorflow.keras.preprocessing.sequence import pad_sequences
import json
import os

# Initialize global variables
model = None
tokenizer = None
max_len_seq = None

def load_model_artifacts():
    global model, tokenizer, max_len_seq
    
    # Load model directly from root directory
    model = tf.keras.models.load_model('shakespeare_model.h5')
    
    # Load tokenizer from root
    with open('tokenizer.json', 'r') as f:
        tokenizer = tf.keras.preprocessing.text.tokenizer_from_json(f.read())
    
    # Load config from root
    with open('config.json', 'r') as f:
        config = json.load(f)
        max_len_seq = config['max_len_seq']

def generate_shakespeare_quote(seed_text, num_words):
    """Generate Shakespeare-style text from a seed text"""
    if model is None:
        load_model_artifacts()
    
    try:
        for _ in range(int(num_words)):
            # Convert the seed text to sequences
            token_list = tokenizer.texts_to_sequences([seed_text])[0]
            # Pad the sequences
            token_list = pad_sequences([token_list], maxlen=max_len_seq-1, padding='pre')
            # Predict the next word
            predicted = model.predict(token_list, verbose=0)
            # Get the word with highest probability
            predicted_word = tokenizer.index_word[np.argmax(predicted, axis=-1).item()]
            # Add the predicted word to the seed text
            seed_text += " " + predicted_word
        
        return seed_text
    except Exception as e:
        return f"Error generating text: {str(e)}"

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_shakespeare_quote,
    inputs=[
        gr.Textbox(
            label="Enter your seed text",
            placeholder="Start your quote here...",
            value="to be or"
        ),
        gr.Slider(
            minimum=1,
            maximum=50,
            value=10,
            step=1,
            label="Number of words to generate"
        )
    ],
    outputs=gr.Textbox(label="Generated Quote"),
    title="Shakespeare Quote Generator",
    description="""Generate Shakespeare-style quotes using AI! 
                Enter a seed text and choose how many words you want to generate.
                The model will continue your text in Shakespeare's style.""",
    examples=[
        ["to be or", 10],
        ["love is", 15],
        ["life is", 12],
        ["death be not", 10],
        ["shall i compare thee", 8]
    ],
    theme=gr.themes.Base()
)

# Launch the app
if __name__ == "__main__":
    iface.launch()