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
    )

    return tokenizer_pl, model_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 detection",
    page_icon=img_favicon,
    initial_sidebar_state="expanded",
)

tokenizer_en, model_en, tokenizer_pl, model_pl = load_model_cache()


def get_predictions(text):
    input_ids = tokenizer_pl(text, return_tensors="pt", truncation=True).input_ids
    output = model_pl.generate(
        input_ids,
        no_repeat_ngram_size=1,
        num_beams=3,
        num_beam_groups=3,
        min_length=10,
        max_length=100,
    )
    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 - RfC generation")

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

    language = st.sidebar.title("Model settings")
    language = st.sidebar.radio(
        "Select model to test",
        [
            "Polish",
        ],
    )

    result = st.button("Find reason for contact")
    if result:
        generated_rfc = get_predictions(text=user_input)
        st.text_area("Reason", generated_rfc)
        print(f"Input: {user_input} ---> Reason for contact: {generated_rfc}")