import streamlit as st
import torch
from pandas import options
from transformers import BartForConditionalGeneration, BartTokenizer

# initialize model + tok variables
model = None
tok = None

# Examples for each models
examples = [
    ["interview-question-remake", "I have a cat named dolche and he's not very friendly with strangers. I've had him for 9 years now and it has been a pleasure to see him grow closer to us every year."],
    ["interview-length-tagged","Today's weather was really nice."],
    ["reverse-interview-question", "There are so many incredible musicians out there and so many really compelling big hits this year that it makes for a really interesting way to recap some of those big events."]
]

# Descriptions for each models
# descriptions = "Interview question remake is a model that..."

# pass in Strings of model choice and input text for context
@st.cache
def genQuestion(model_choice, context):
    # global descriptions
    if model_choice=="interview-question-remake":
        model = BartForConditionalGeneration.from_pretrained("hyechanjun/interview-question-remake")
        tok = BartTokenizer.from_pretrained("hyechanjun/interview-question-remake")
        # descriptions = "Interview question remake is a model that..."
    elif model_choice=="interview-length-tagged":
        model = BartForConditionalGeneration.from_pretrained("hyechanjun/interview-length-tagged")
        tok = BartTokenizer.from_pretrained("hyechanjun/interview-length-tagged")
        # descriptions = "Interview question tagged is a model that..."
    elif model_choice=="reverse-interview-question":
        model = BartForConditionalGeneration.from_pretrained("hyechanjun/reverse-interview-question")
        tok = BartTokenizer.from_pretrained("hyechanjun/reverse-interview-question")
        # descriptions = "Reverse interview question  is a model that..."

    inputs = tok(context, return_tensors="pt")
    output = model.generate(inputs["input_ids"], num_beams=16, max_length=64, min_length=9, num_return_sequences=16, diversity_penalty = 1.0, num_beam_groups=4)
    final_output = ''
    for i in range(16):
        final_output +=  [tok.decode(beam, skip_special_tokens=True, clean_up_tokenization_spaces=False) for beam in output][i] + "\n\n"

    return final_output


# Wide page layout (instead of having a narrower, one-column page layout)
st.set_page_config(layout="wide")

# Title
st.title("Interview AI Test Website")

# Adding a Session State to store stateful variables and for saving user's labels/tags for generated questions
if 'button_sent' not in st.session_state:
    st.session_state.button_sent = False

# Input fields
input = st.text_input('Context')                                    # user inputs context to construct a response (str)

maxl, minl = st.columns(2)

option = st.selectbox(
                'Please select a model.',
                ('interview-question-remake', 'interview-length-tagged', 'reverse-interview-question'))

if option == 'interview-question-remake':
    st.write("This is the re-fine-tuned base model for our interview AI. It returns strings terminating in a question mark (?).")
elif option == 'interview-length-tagged':
    st.write("This is a length-tagged version of our interview AI. You can specify how long its responses should be (ranges of multiples of 10)")
elif option == 'reverse-interview-question':
    st.write("This model asks a question that would have resulted in the context you provide (a.k.a. it traverses backward through the interview)")


# Column layout to display generated responses alongside tags
col1, col2 = st.columns((3, 1))

if st.button('Submit') or st.session_state.button_sent:
    with st.spinner('Generating a response...'):
        output = genQuestion(option, input)
        print(output)
    # st.write(output)
    st.session_state.button_sent = True
    col1.text_area(label="Generated Responses:", value=output, height=200)



# TODO:
#   - disable multiselect widget when responses are being generated AND when a question is not selected to be tagged
#   - connect tags with an individual question
#   - save session state so tags associated with their respective questions can also be saved
#   - write/store the saved state data to some database for future use?
#   - brainstorm good names for tags/labels OR allow users to enter their own tag names if possible