from transformers import T5ForConditionalGeneration, T5Tokenizer
import streamlit as st
from PIL import Image
import os


@st.cache(allow_output_mutation=True)
def load_model_cache():
    auth_token = os.environ.get("TOKEN_FROM_SECRET") or True

    tokenizer_pl = T5Tokenizer.from_pretrained(
        "Voicelab/vlt5-base-rfc-v1_2", use_auth_token=auth_token
    )
    model_pl = T5ForConditionalGeneration.from_pretrained(
        "Voicelab/vlt5-base-rfc-v1_2", use_auth_token=auth_token
    )
    model_det_pl = T5ForConditionalGeneration.from_pretrained(
        "Voicelab/vlt5-base-rfc-detector-1.0", use_auth_token=auth_token
    )

    return tokenizer_pl, model_pl, model_det_pl


img_full = Image.open("images/vl-logo-nlp-blue.png")
img_short = Image.open("images/sVL-NLP-short.png")
img_favicon = Image.open("images/favicon_vl.png")
max_length: int = 5000
cache_size: int = 100

st.set_page_config(
    page_title="DEMO - Reason for Contact generation",
    page_icon=img_favicon,
    initial_sidebar_state="expanded",
)

tokenizer_pl, model_pl, model_det_pl = load_model_cache()


def get_predictions(text, mode):
    input_ids = tokenizer_pl(text, return_tensors="pt", truncation=True).input_ids
    if mode == "Polish - RfC Generation":
        output = model_pl.generate(
            input_ids,
            no_repeat_ngram_size=1,
            num_beams=3,
            num_beam_groups=3,
            min_length=10,
            max_length=100,
        )
    elif mode == "Polish - RfC Detection":
        output = model_det_pl.generate(
            input_ids,
            no_repeat_ngram_size=2,
            num_beams=3,
            num_beam_groups=3,
            repetition_penalty=1.5,
            diversity_penalty=2.0,
            length_penalty=2.0,
        )
    predicted_rfc = tokenizer_pl.decode(output[0], skip_special_tokens=True)
    return predicted_rfc


def trim_length():
    if len(st.session_state["input"]) > max_length:
        st.session_state["input"] = st.session_state["input"][:max_length]


if __name__ == "__main__":
    st.sidebar.image(img_short)
    st.image(img_full)
    st.title("VLT5 - Reason for Contact generator")

    generated_keywords = ""
    user_input = st.text_area(
        label=f"Input text (max {max_length} characters)",
        value="",
        height=300,
        on_change=trim_length,
        key="input",
    )

    mode = st.sidebar.title("Model settings")
    mode = st.sidebar.radio(
        "Select model to test",
        [
            "Polish - RfC Generation",
            "Polish - RfC Detection",
        ],
    )

    result = st.button("Find reason for contact")
    if mode == "Polish - RfC Generation (accepts whole conversation)":
        st.markdown("### You selected RfC Generation model.")
        st.markdown("-- *Input*: Whole conversation. Should specify roles (e.g. **AGENT: Hello, how can I help you? CLIENT: Hi, I would like to report a stolen card.**")
        st.markdown("-- *Output*: Reason for calling for the whole conversation.")
        text_area = "Put a whole conversation or full e-mail here."
        
    elif mode == "Polish - RfC Detection (accepts one turn)":
        st.markdown("### You selected RfC Detection model.")
        st.markdown("-- *Input*: A single turn from the conversation e.g. **'Hello, how can I help you?'** or **'Hi, I would like to report a stolen card.'**")
        st.markdown("-- *Output*: Model will return an empty string if a turn possibly does not includes Reason for Calling, or a sentence if the RfC is detected.")
        text_area = "Put a single turn or a few sentences here."
        
    if result:
        generated_rfc = get_predictions(text=user_input, mode=mode)
        st.text_area(text_area, generated_rfc)
        print(f"Input: {user_input} ---> Reason for contact: {generated_rfc}")