import gradio as gr
import random
import openai
import os

def tell_us_about_yo_momma(category, topic):
    openAI_key = os.getenv('openAI_key')
    if openAI_key.strip()=='':
        return '[ERROR]: Please enter you Open AI Key. Get your key here : https://platform.openai.com/account/api-keys'
    openai.api_key = openAI_key

    prompt = ""
    prompt = f"You are a comedian who is bold, cutting, and broadly appealing.  Tell us a yo momma joke about her {category} that is at least loosely related to: {topic}."
   
    completions = openai.Completion.create(
        model="text-davinci-003",
        prompt=prompt,
        max_tokens=512,
        n=1,
        stop=None,
        temperature=0.7,
    )
    message = completions.choices[0].text.strip()
    return message

def launch_demo():
    categories = [
        "Appearance",
        "Weight",
        "Intelligence",
        "Age",
        "Financial Status",
        "Cleanliness",
        "Dirtyness",
        "Hygiene",
        "Lifestyle",
        "Occupation",
        "Culture",
        "Cooking Skills"
    ]
    
    
    topics = [
        "Holidays",
        "Plants",
        "Animals",
        "Countries",
        "Pop Culture",
        "Religion",
        "Lifestyle",
        "Environment",
        "Technology",
        "Sports",
        "Music",
        "Space/Astronomy",
        "Food and Cooking",
        "History",
        "Travel",
        "Education",
        "Health and Fitness",
        "Movies/TV Shows"
    ]

    dropdown = gr.components.Dropdown(categories, label="Category")
    text_input = gr.components.Textbox(label="Topic")
    output = gr.components.Textbox(label="Lolz")

    # some random pre-seeded options
    rando_choice = random.choice(categories)
    rando_topic = random.choice(topics)
    rando_choice_b = random.choice(categories)
    rando_topic_b = random.choice(topics)
    rando_choice_c = random.choice(categories)
    rando_topic_c = random.choice(topics)

    examples = [
        [rando_choice, rando_topic],
        [rando_choice_b, rando_topic_b],
        [rando_choice_c, rando_topic_c],
    ]

    gr.Interface(
        fn=tell_us_about_yo_momma,
        inputs=[dropdown, text_input],
        outputs=output,
        examples=examples,
        title="Yo Momma so Generative...",
        description="Select a Category, enter a Topic, and click Submit to find out how generative Yo Momma is.<br><br>You can also try selecting one of the supplied examples.",
        theme="default",
    ).launch()

launch_demo()