import streamlit as st
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, GPT2Tokenizer, GPT2LMHeadModel, MT5Model, AutoModelForSeq2SeqLM, BartForConditionalGeneration, BartTokenizer

# Frasi di esempio
examples = [
    "Caesar in Gallia bellum gessit. <mask> instructae erant, sed subito",
    "Est autem et aliud genus testudinis, quod <mask>, quemadmodum quae supra scripta sunt",
    "Quemadmodum vero minores rotae duriores et <mask>, sic phalangae et iuga, in quibus partibus habent minora a centro ad capita intervalla",
    "illud additur, ne, qui certum ordinem ex <mask>, ulli vos alteri hominum generi haerere vereamini nec timeatis vos"
]
examples_correct = [
    "Omnes legiones",
    "reliqua omnia habet",
    "difficiliores habent motus",
    "nostra definitione retinetis"
]

# Imposta una chiave di sessione per il testo input, così possiamo aggiornarlo
if "input_text_value" not in st.session_state:
    st.session_state["input_text_value"] = examples[0]
if "input_text_value_correct" not in st.session_state:    
    st.session_state["input_text_value_correct"] = examples_correct[0]

st.title("Completamento di frasi in testi Latino Antico")

st.write("Esempi di testo (clicca sul bottone per copiare la frase nel campo di input):")

# Per ogni frase, creiamo una riga con la frase + bottone "Usa questa frase"
for i, example in enumerate(examples, start=1):
    cols = st.columns([4,1])  # la prima colonna più larga per il testo, la seconda più stretta per il bottone
    with cols[0]:
        st.write(f"Esempio {i}: {example}")
    with cols[1]:
        # Se il bottone viene premuto, aggiorna la session state
        if st.button(f"Usa {i}"):
            st.session_state["input_text_value"] = example
            st.session_state["input_text_value_correct"] = examples_correct[i]

# Qui la text_input è associata alla chiave "input_text_value"
input_text = st.text_area(
    label="Testo:",
    height=150,  # Altezza in pixel
    key="input_text_value"
)

# Esempio: carichiamo un modello di linguaggio
# (sostituisci con i tuoi modelli)
#model = GPT2LMHeadModel.from_pretrained("Cicciokr/GPT2-Latin-GenText")
#tokenizer = GPT2Tokenizer.from_pretrained("Cicciokr/GPT2-Latin-GenText")
#model_name = "morenolq/bart-it"
model_name_mio = "Cicciokr/BART-la-s"
model_mio = AutoModelForSeq2SeqLM.from_pretrained(model_name_mio)
tokenizer_mio = AutoTokenizer.from_pretrained(model_name_mio)
#tokenizer_mio.pad_token_id = tokenizer_mio.eos_token_id

#generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
generator_mio = pipeline("text2text-generation", model=model_mio, tokenizer=tokenizer_mio)

# Se l'utente ha inserito (o selezionato) un testo
if input_text:
    if 'input_text_value_correct' in st.session_state:
        st.write(f" Parola corretta: {st.session_state['input_text_value_correct']}\n")
        st.write(f" -----------------------------------------------------------\n")

    inputs_mio = tokenizer_mio(input_text, return_tensors="pt")
    output_mio = model_mio.generate(
        **inputs_mio,
        forced_bos_token_id=tokenizer_mio.bos_token_id,
        max_length=50, do_sample=True, top_p=0.96, num_return_sequences=5
        #num_return_sequences=1,
        #top_k=50,  # 🔹 Maggiore varietà nelle scelte
        #top_p=0.95,  # 🔹 Nucleus sampling per migliorare il realismo
        #temperature=0.8,  # 🔹 Bilancia la casualità
        #repetition_penalty=1.2,  # 🔹 Evita ripetizioni e loop
        #max_length=50  # 🔹 Previene ripetizioni infinite
    )
    generated_text_mio = tokenizer_mio.decode(output_mio[0], skip_special_tokens=True)
    #generated_text_mio = output_mio[0]["generated_text"]
    st.subheader("Risultato BART CC100:")
    st.write(f" Frase predetta: {generated_text_mio}\n")