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

# Create a Parrot instance
parrot = Parrot(model_tag="t5-small", use_gpu=False)

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

# Supported models for paraphrasing
models = {
    "T5 Small": "t5-small",
    "T5 Base": "t5-base",
    "Pegasus": "google/pegasus-multi_news"
}

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)

def highlight_diffs(original_text, paraphrased_text, level):
    # Use diff-match-patch library to highlight the differences between original and paraphrased text
    dmp = difflib.Differ()
    diffs = list(dmp.compare(original_text.split(), paraphrased_text.split()))
    if level == "only_changes":
        return " ".join([diff[2:] for diff in diffs if diff.startswith("+ ")])
    else:
        return " ".join([diff[2:] if diff.startswith("+ ") else diff[2:].replace(diff[0], "<span style='background-color:yellow'>"+diff[0]+"</span>") for diff in diffs])

# 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?"),
        gr.inputs.File(label="Upload file containing text to paraphrase (Optional):")
    ],
    outputs=[
        gr.outputs.Textbox(label="Paraphrased text:"),
        gr.outputs.File(label="Download paraphrased text as text file:"),
        gr.outputs.HTML(label="Highlighted differences:")
    ],
    title="Parrot Paraphrasing Tool",
    description="Enter text to paraphrase, select language and model, and choose level of paraphrasing."
)

def app(text, language, model, level, audio_output, file_upload):
    if file_upload:
        with open(file_upload.name, "r", encoding="utf-8") as f:
            text = f.read()
    paraphrased_text = paraphrase_text(text, languages[language], models[model], level)
    if audio_output:
        speak_paraphrased_text(paraphrased_text, languages[language])
        download_output = gr.outputs.File(paraphrased_text, label="Download paraphrased text as text file:", type="file", file_name="paraphrased_text.txt")
        highlighted_diffs = highlight_diffs(text, paraphrased_text, level)
    return paraphrased_text, download_output, highlighted_diffs

interface.launch()