import streamlit as st
import openai

#streamlit run main.py
# import os
# os.environ.getattribute("openai.api_key")

header = st.container()
about = st.container()
model_output = st.container()

with header:
  st.title('RHYME THYME')
  #st.text('This this bot-bop-shop uses Artificail Ingredients \
  #(AI) to produce natural flavors.')

with model_output:
  st.header('Half-baked Takes:')
  sel_col, disp_col = st.columns(2)

  sel_col.subheader('Input ideas')

  #input_poem = sel_col.selectbox('What type of poem do you want to write?', options=['limerick', 'haiku', 'sonnet', 'acrostic', 'villanelle', 'ode', 'elegy', 'ballad', 'couplet', 'tercet', 'quatrain', 'cinquan', 'sestet'])

  input_word1 = sel_col.text_input('Input a word to use in the poem:')
 
  input_word2 = sel_col.text_input('Input another word to use in the poem:')

 
  # This is the NLP model


 
  def generate_poem(input_word1, input_word2):
    # Use GPT-3 to generate a limerick based on the user's input word
    prompt = f"Write a limerick in AABBA rhyme scheme about {input_word1} and '{input_word2}':"
    completions = openai.Completion.create(
      engine="text-davinci-002",
      prompt=prompt,
      max_tokens=1024,
      n=1,
      stop=None,
      temperature=0.5,
    )
 
    # Return the generated limerick
    return completions.choices[0].text

  limerick = generate_poem(input_word1.lower(), input_word2.lower())

  # This is where the results are displayed on the app
  disp_col.subheader('Limerick output')
  disp_col.write(limerick)