import os
import openai
import gradio as gr
from langchain import OpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
import datetime

# Console to variable
from io import StringIO
import sys

NUM_WORDS_DEFAULT = 0
FORMALITY_DEFAULT = "Casual"
TEMPERATURE_DEFAULT = 0.5
EMOTION_DEFAULT = "N/A"
TRANSLATE_TO_DEFAULT = "Don't translate"
PROMPT_TEMPLATE = PromptTemplate(
    input_variables=["original_words", "num_words", "formality", "emotions", "translate_to"],
    template="Express {num_words} in a {formality} manner, "
             "{emotions} {translate_to} the following: \n{original_words}\n",
)


def set_openai_api_key(api_key, openai_api_key, temperature, llm_chain):
    if api_key:
        print("temperature: ", temperature)
        openai_api_key = api_key

        os.environ["OPENAI_API_KEY"] = api_key
        llm = OpenAI(model_name='text-davinci-003', temperature=temperature, max_tokens=512)
        os.environ["OPENAI_API_KEY"] = ""

        llm_chain = LLMChain(llm=llm, prompt=PROMPT_TEMPLATE, verbose=True)
        return openai_api_key, llm_chain


def desc2sheet(desc, openai_api_key, temperature, llm_chain, num_words, formality,
               anticipation_level, joy_level, trust_level,
               fear_level, surprise_level, sadness_level, disgust_level, anger_level,
               translate_to):
    if not openai_api_key or openai_api_key == "":
        return "<pre>Please paste your OpenAI API key (see https://beta.openai.com)</pre>"

    num_words_prompt = ""
    if num_words and int(num_words) != 0:
        num_words_prompt = "using up to " + str(num_words) + " words, "

    # Change some arguments to the lower case
    formality = formality.lower()
    anticipation_level = anticipation_level.lower()
    joy_level = joy_level.lower()
    trust_level = trust_level.lower()
    fear_level = fear_level.lower()
    surprise_level = surprise_level.lower()
    sadness_level = sadness_level.lower()
    disgust_level = disgust_level.lower()
    anger_level = anger_level.lower()

    llm_chain.llm.temperature = temperature
    print("llm_chain.llm.temperature: ", llm_chain.llm.temperature)

    # put all emotions into a list
    emotions = []
    if anticipation_level != "n/a":
        emotions.append(anticipation_level)
    if joy_level != "n/a":
        emotions.append(joy_level)
    if trust_level != "n/a":
        emotions.append(trust_level)
    if fear_level != "n/a":
        emotions.append(fear_level)
    if surprise_level != "n/a":
        emotions.append(surprise_level)
    if sadness_level != "n/a":
        emotions.append(sadness_level)
    if disgust_level != "n/a":
        emotions.append(disgust_level)
    if anger_level != "n/a":
        emotions.append(anger_level)

    emotions_str = ""
    # If there are any emotions, join the emotions list into a string that begins with "Emotions: ", with the word "and" before the last emotion
    if len(emotions) > 0:
        if len(emotions) == 1:
            emotions_str = "with emotion of " + emotions[0] + ", "
        else:
            emotions_str = "with emotions of " + ", ".join(emotions[:-1]) + " and " + emotions[-1] + ", "

    translate_to_str = ""
    if translate_to != TRANSLATE_TO_DEFAULT:
        translate_to_str = "translated to " + translate_to + ", "

    formatted_prompt = PROMPT_TEMPLATE.format(
        original_words=desc,
        num_words=num_words_prompt,
        formality=formality,
        emotions=emotions_str,
        translate_to=translate_to_str,
    )
    generated_text = llm_chain.run({'original_words': desc, 'num_words': num_words_prompt, 'formality': formality,
                                    'emotions': emotions_str, 'translate_to': translate_to_str})

    prompt_plus_generated = "#### GPT prompt:\n" + formatted_prompt + "#### Generated text:\n" + generated_text

    return prompt_plus_generated


def update_temperature(temp_slider, temp_state):
    if temp_slider:
        temp_state = temp_slider
        return temp_state


def update_num_words(slider, state):
    if slider:
        state = slider
        return state


def update_formality(radio, state):
    if radio:
        state = radio
        return state


def update_foo(widget, state):
    if widget:
        state = widget
        return state


block = gr.Blocks(css=".gradio-container {background-color: lightgray}")

with block:
    openai_api_key_state = gr.State()
    temperature_state = gr.State(TEMPERATURE_DEFAULT)
    llm_chain_state = gr.State()
    num_words_state = gr.State(NUM_WORDS_DEFAULT)
    formality_state = gr.State(FORMALITY_DEFAULT)
    anticipation_level_state = gr.State(EMOTION_DEFAULT)
    joy_level_state = gr.State(EMOTION_DEFAULT)
    trust_level_state = gr.State(EMOTION_DEFAULT)
    fear_level_state = gr.State(EMOTION_DEFAULT)
    surprise_level_state = gr.State(EMOTION_DEFAULT)
    sadness_level_state = gr.State(EMOTION_DEFAULT)
    disgust_level_state = gr.State(EMOTION_DEFAULT)
    anger_level_state = gr.State(EMOTION_DEFAULT)
    translate_to_state = gr.State(TRANSLATE_TO_DEFAULT)

    with gr.Row():
        temperature_slider = gr.Slider(label="GPT Temperature", value=TEMPERATURE_DEFAULT, minimum=0.0, maximum=1.0,
                                       step=0.1)
        title = gr.Markdown("""<h3><center>GPT-3.5 Express-inator</center></h3>""")
        openai_api_key_textbox = gr.Textbox(placeholder="Paste your OpenAI API key",
                                            show_label=False, lines=1, type='password')

    table = gr.Markdown("")
    with gr.Row():
        request = gr.Textbox(label="Words to express: ",
                             value="The quick brown fox jumped over the lazy dog.",
                             placeholder="Ex: The quick brown fox jumped over the lazy dog.")

        submit = gr.Button(value="Express", variant="secondary").style(full_width=False)

    with gr.Row():
        with gr.Column():
            num_words_slider = gr.Slider(label="Max number of words to generate (0 for don't care)",
                                         value=NUM_WORDS_DEFAULT, minimum=0, maximum=100, step=10)
            num_words_slider.change(update_num_words,
                                    inputs=[num_words_slider, num_words_state],
                                    outputs=[num_words_state])

            formality_radio = gr.Radio(label="Formality:", choices=["Casual", "Polite", "Honorific"],
                                       value=FORMALITY_DEFAULT)
            formality_radio.change(update_formality,
                                   inputs=[formality_radio, formality_state],
                                   outputs=[formality_state])

            translate_to_radio = gr.Radio(label="Translate to:", choices=[
                "Arabic", "British English", "Chinese (Simplified)", "Chinese (Traditional)", "Czech", "Danish", "Dutch", "English",
                "Finnish", "French", "Gen Z slang", "German", "Greek", "Hebrew", "Hindi", "Hungarian", "Indonesian",
                "Italian", "Japanese", "Klingon", "Korean", "Norwegian", "Old English", "Polish", "Portuguese",
                "Romanian", "Russian", "Spanish", "Swedish", "Thai", "Turkish", "Vietnamese", "Yoda",
                TRANSLATE_TO_DEFAULT], value=TRANSLATE_TO_DEFAULT)

            translate_to_radio.change(update_foo,
                                      inputs=[translate_to_radio, translate_to_state],
                                      outputs=[translate_to_state])

        with gr.Column():
            anticipation_level_radio = gr.Radio(label="Anticipation level:",
                                                choices=[EMOTION_DEFAULT, "Interest", "Anticipation", "Vigilance"],
                                                value=EMOTION_DEFAULT)
            anticipation_level_radio.change(update_foo,
                                            inputs=[anticipation_level_radio, anticipation_level_state],
                                            outputs=[anticipation_level_state])

            joy_level_radio = gr.Radio(label="Joy level:",
                                       choices=[EMOTION_DEFAULT, "Serenity", "Joy", "Ecstasy"],
                                       value=EMOTION_DEFAULT)
            joy_level_radio.change(update_foo,
                                   inputs=[joy_level_radio, joy_level_state],
                                   outputs=[joy_level_state])

            trust_level_radio = gr.Radio(label="Trust level:",
                                         choices=[EMOTION_DEFAULT, "Acceptance", "Trust", "Admiration"],
                                         value=EMOTION_DEFAULT)
            trust_level_radio.change(update_foo,
                                     inputs=[trust_level_radio, trust_level_state],
                                     outputs=[trust_level_state])

            fear_level_radio = gr.Radio(label="Fear level:",
                                        choices=[EMOTION_DEFAULT, "Apprehension", "Fear", "Terror"],
                                        value=EMOTION_DEFAULT)
            fear_level_radio.change(update_foo,
                                    inputs=[fear_level_radio, fear_level_state],
                                    outputs=[fear_level_state])

            surprise_level_radio = gr.Radio(label="Surprise level:",
                                            choices=[EMOTION_DEFAULT, "Distraction", "Surprise", "Amazement"],
                                            value=EMOTION_DEFAULT)
            surprise_level_radio.change(update_foo,
                                        inputs=[surprise_level_radio, surprise_level_state],
                                        outputs=[surprise_level_state])

            sadness_level_radio = gr.Radio(label="Sadness level:",
                                           choices=[EMOTION_DEFAULT, "Pensiveness", "Sadness", "Grief"],
                                           value=EMOTION_DEFAULT)
            sadness_level_radio.change(update_foo,
                                       inputs=[sadness_level_radio, sadness_level_state],
                                       outputs=[sadness_level_state])

            disgust_level_radio = gr.Radio(label="Disgust level:",
                                           choices=[EMOTION_DEFAULT, "Boredom", "Disgust", "Loathing"],
                                           value=EMOTION_DEFAULT)
            disgust_level_radio.change(update_foo,
                                       inputs=[disgust_level_radio, disgust_level_state],
                                       outputs=[disgust_level_state])

            anger_level_radio = gr.Radio(label="Anger level:",
                                         choices=[EMOTION_DEFAULT, "Annoyance", "Anger", "Rage"],
                                         value=EMOTION_DEFAULT)
            anger_level_radio.change(update_foo,
                                     inputs=[anger_level_radio, anger_level_state],
                                     outputs=[anger_level_state])

    submit.click(desc2sheet,
                 inputs=[
                     request, openai_api_key_state, temperature_state, llm_chain_state, num_words_state,
                     formality_state,
                     anticipation_level_state, joy_level_state, trust_level_state, fear_level_state,
                     surprise_level_state, sadness_level_state, disgust_level_state, anger_level_state,
                     translate_to_state],
                 outputs=[table])

    request.submit(desc2sheet,
                   inputs=[
                       request, openai_api_key_state, temperature_state, llm_chain_state, num_words_state,
                       formality_state,
                       anticipation_level_state, joy_level_state, trust_level_state, fear_level_state,
                       surprise_level_state, sadness_level_state, disgust_level_state, anger_level_state,
                       translate_to_state],
                   outputs=[table])

    openai_api_key_textbox.change(set_openai_api_key,
                                  inputs=[openai_api_key_textbox, openai_api_key_state, temperature_state,
                                          llm_chain_state],
                                  outputs=[openai_api_key_state, llm_chain_state])

    temperature_slider.change(update_temperature,
                              inputs=[temperature_slider, temperature_state],
                              outputs=[temperature_state])

block.launch()