# Copyright 2023 by Jan Philip Wahle, https://jpwahle.com/
# All rights reserved.

import os
import random
import time

import gradio as gr
from openai import OpenAI

client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))


def create_prompt(sentence, paraphrase_type):
    """
    Creates a prompt for generating a paraphrase of a given sentence with specified types.

    Args:
        sentence (str): The original sentence to be paraphrased.
        paraphrase_type (str): The type of paraphrase to be generated.

    Returns:
        dict: A dictionary containing the prompt message.
    """
    return {
        "messages": [
            {
                "role": "user",
                "content": (
                    "Given the following sentence, generate a paraphrase with"
                    f" the following types. Sentence: {sentence}. Paraphrase"
                    f" Types: {paraphrase_type}"
                ),
            }
        ]
    }


paraphrase_types = [
    "Derivational Changes",
    "Inflectional Changes",
    "Modal Verb Changes",
    "Spelling changes",
    "Change of format",
    "Same Polarity Substitution (contextual)",
    "Same Polarity Substitution (habitual)",
    "Same Polarity Substitution (named ent.)",
    "Converse substitution",
    "Opposite polarity substitution (contextual)",
    "Opposite polarity substitution (habitual)",
    "Synthetic/analytic substitution",
    "Coordination changes",
    "Diathesis alternation",
    "Ellipsis",
    "Negation switching",
    "Subordination and nesting changes",
    "Direct/indirect style alternations",
    "Punctuation changes",
    "Syntax/discourse structure changes",
    "Entailment",
    "Identity",
    "Non-paraphrase",
    "Addition/Deletion",
    "Change of order",
    "Semantic-based",
]

with gr.Blocks() as demo:
    description = gr.Markdown(
        """
        ## Paraphrase Type Generator
        This demo uses a fine-tuned gpt-4o-mini model to generate paraphrases given specific paraphrase types.
        
        **How to use:**
        1. Select one or many type of paraphrase from the dropdown menu.
        2. Enter a sentence in the text box.
        3. Click the "Submit" button or hit enter.
        4. The application will generate a paraphrase of the input sentence based on the selected type.
    """
    )
    chatbot = gr.Chatbot()
    types = gr.Dropdown(
        paraphrase_types,
        value="Syntax/discourse structure changes",
        multiselect=True,
        allow_custom_value=True,
    )
    msg = gr.Textbox()
    submit = gr.Button("Submit")
    clear = gr.Button("Clear")

    def user(user_message, history):
        """
        This function takes in a user message and a history of previous messages, and returns an empty string and an updated history list with the user message appended to it.

        Args:
        - user_message (str): The message sent by the user.
        - history (list): A list of previous messages, where each message is a list containing the message text and the bot's response.

        Returns:
        - A tuple containing an empty string and the updated history list.
        """
        return "", history + [[user_message, None]]

    def generate_paraphrase(user_message, paraphrase_type, history):
        """
        Generates a paraphrase of the user's message using OpenAI's GPT-3 model.

        Args:
            user_message (str): The message to be paraphrased.
            paraphrase_type (str): The type of paraphrase to generate.
            history (list): A list of previous messages in the conversation.

        Yields:
            list: A list of previous messages in the conversation, including the new paraphrase.
        """
        types_as_str = ",".join(paraphrase_type)
        history[-1][1] = f"[System: {types_as_str}]\n\n"
        prompt = create_prompt(history[-1][0], paraphrase_type)
        bot_message = client.chat.completions.create(
            model="ft:gpt-4o-mini-2024-07-18:personal::ADQ0IcdZ",
            messages=prompt["messages"],
        )
        for character in bot_message.choices[0].message.content:
            history[-1][1] += character
            time.sleep(0.01)
            yield history

    chatbot.value = [
        [
            (
                "These outlined a"
                " theory of the photoelectric effect, explained Brownian"
                " motion, introduced his special theory of relativity—a theory"
                " which addressed the inability of classical mechanics to"
                " account satisfactorily for the behavior of the"
                " electromagnetic field—and demonstrated that if the special"
                " theory is correct, mass and energy are equivalent to each"
                " other."
            ),
            (
                "[System: Syntax/discourse structure changes]\n\nAnother of"
                " the papers introduced Einstein's special theory of"
                " relativity, which addressed the inability of classical"
                " mechanics to account for the behavior of the electromagnetic"
                " field, and demonstrated that, if the special theory is"
                " correct, mass and energy are equivalent to each other."
            ),
        ]
    ]

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        generate_paraphrase, [msg, types, chatbot], chatbot
    )
    submit.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        generate_paraphrase, [msg, types, chatbot], chatbot
    )
    clear.click(lambda: None, None, chatbot, queue=False)

demo.queue()
demo.launch()