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"
PROMPT_TEMPLATE = PromptTemplate(
    input_variables=["original_words", "num_words", "formality", "emotions"],
    template="Restate {num_words} in a {formality} manner, "
             "{emotions} 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):
    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]

    formatted_prompt = PROMPT_TEMPLATE.format(
        original_words=desc,
        num_words=num_words_prompt,
        formality=formality,
        emotions=emotions_str,
    )
    generated_text = llm_chain.run({'original_words': desc, 'num_words': num_words_prompt, 'formality': formality,
                                    'emotions': emotions_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)

    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="Create", variant="secondary").style(full_width=False)

    num_words_slider = gr.Slider(label="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])

    with gr.Accordion("Emotions", open=True):
        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],
                 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],
                   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()