import openai
import gradio as gr
from langchain import OpenAI
from langchain.chains import PALChain
import datetime

gpt_only_prompt = "Calculate the following, giving only the final answer:\n"
prompt = ""

def openai_create(prompt):
    print("prompt: " + prompt)

    # We use temperature of 0.0 because it gives the most predictable, factual answer (i.e. avoids hallucination).
    response = openai.Completion.create(
        model="text-davinci-003",
        prompt=prompt,
        temperature=0.0,
        max_tokens=300,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0
    )

    return response.choices[0].text


def calc_gpt_only(math_problem):
    answer = openai_create(gpt_only_prompt + math_problem + "\n")

    print("math problem: " + math_problem)
    print("calc_gpt_only answer: " + answer)

    html = "<pre>" + answer + "</pre>"
    return html


def calc_gpt_pal(math_problem):
    llm = OpenAI(model_name='code-davinci-002', temperature=0, max_tokens=512)
    pal_chain = PALChain.from_math_prompt(llm, verbose=True)
    answer = pal_chain.run(math_problem)

    print("math problem: " + math_problem)
    print("calc_gpt_pal answer: " + answer)

    html = "<pre>" + answer + "</pre>"
    return html


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

with block:
    title = gr.Markdown("""<h3><center>Comparing GPT math results</center></h3>""")
    answer_html = gr.Markdown()

    request = gr.Textbox(label="Math question:",
                         placeholder="Ex: What is the sum of the first 10 prime numbers?")
    with gr.Row():
        gpt_only = gr.Button(value="GPT Only", variant="secondary").style(full_width=False)
        gpt_pal = gr.Button(value="GPT w/PAL", variant="secondary").style(full_width=False)

    gr.Examples(
        examples=["42 times 81",
                  "Olivia has $23. She bought five bagels for $3 each. How much money does she have left?",
                  "What is the sum of the first 21 Fibonacci numbers?",
                  "Jane quit her job on Mar 20, 2020. 176 days have passed since then. What is the date tomorrow in MM/DD/YYYY?"],
        inputs=request
    )

    gr.HTML("See <a href='https://reasonwithpal.com/'>PAL: Program-aided Language Models</a>")

    gr.HTML("<center>Powered by <a href='https://github.com/hwchase17/langchain'>LangChain 🦜️🔗</a></center>")

    gpt_only.click(calc_gpt_only, inputs=[request], outputs=[answer_html])
    gpt_pal.click(calc_gpt_pal, inputs=[request], outputs=[answer_html])

block.launch()