import gradio as gr
from parrot import Parrot
from gtts import gTTS
from IPython.display import Audio
import os

# Create a Parrot instance
parrot = Parrot(model_tag="turing", use_gpu=False)

# Supported languages for paraphrasing
languages = {
    "English": "en",
    "French": "fr",
    "Spanish": "es",
    "German": "de"
}

# Supported models for paraphrasing
models = {
    "Turing": "turing",
    "Pegasus": "pegasus"
}

def paraphrase_text(text, language, model, level):
    # Call the Parrot instance to paraphrase the input text
    result = parrot.predict([text], model_tag=model, context_length=level, src_lang=language, dest_lang=language)
    return result[0]['translation_text']

def speak_paraphrased_text(text, language):
    # Use gTTS to generate an audio file of the paraphrased text and play it
    tts = gTTS(text=text, lang=language)
    tts.save("paraphrased_text.mp3")
    return Audio("paraphrased_text.mp3", autoplay=True)

# Create a Gradio interface
interface = gr.Interface(
    fn=paraphrase_text,
    inputs=[
        gr.inputs.Textbox(label="Enter text to paraphrase:"),
        gr.inputs.Dropdown(list(languages.keys()), label="Select language for paraphrasing:"),
        gr.inputs.Dropdown(list(models.keys()), label="Select model for paraphrasing:"),
        gr.inputs.Slider(minimum=1, maximum=10, default=5, label="Level of paraphrasing:"),
        gr.inputs.Checkbox(label="Generate audio output?")
    ],
    outputs=gr.outputs.Textbox(label="Paraphrased text:"),
    title="Parrot - Text Paraphrasing",
    description="Enter a text to paraphrase using Parrot.",
    theme="default"
)

# Launch the interface
interface.launch()

# If the user has selected the "Generate audio output?" checkbox, call the speak_paraphrased_text function to generate and play the audio
if interface.inputs[4].value:
    speak_paraphrased_text(interface.outputs.value, languages[interface.inputs[1].value])
    
# Delete the generated audio file
os.remove("paraphrased_text.mp3")