import streamlit as st
from transformers import pipeline, AutoModelForMaskedLM, AutoTokenizer
from cltk.data import download as cltk_download
cltk_download("lat")

# Esempio: prova a importare CLTK (se presente)
try:
    from cltk import NLP
    nlp_lat = NLP(language="lat")
except ImportError:
    nlp_lat = None

# 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"] = "Lorem ipsum dolor sit amet, [MASK] adipiscing elit."

# Frasi di esempio
examples = [
    "Asdrubal, frater Annibalis, qui secundo Punico bello [MASK] ingentibus copiis ab Hispania veniens ...",
    "hanno et mago qui [MASK] punico bello cornelium consulem aput liparas ceperunt ...",
    "Lorem ipsum dolor sit amet, [MASK] adipiscing elit.",
    "Populus Romanus cum Macedonibus [MASK] ter gessit"
]

st.title("Completamento di parole in testi Latino Antico con Analisi Morfologica")

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

# Qui la text_input è associata alla chiave "input_text_value"
input_text = st.text_input("Testo:", key="input_text_value")

# Esempio: carichiamo un modello di linguaggio
# (sostituisci con i tuoi modelli)
tokenizer_roberta = AutoTokenizer.from_pretrained("Cicciokr/Roberta-Base-Latin-Uncased")
model_roberta = AutoModelForMaskedLM.from_pretrained("Cicciokr/Roberta-Base-Latin-Uncased")
fill_mask_roberta = pipeline("fill-mask", model=model_roberta, tokenizer=tokenizer_roberta)

# Se l'utente ha inserito (o selezionato) un testo
if input_text:
    # Sostituiamo [MASK] con <mask> (lo tokenizer Roberta se lo aspetta così)
    input_text_roberta = input_text.replace("[MASK]", "<mask>")
    predictions_roberta = fill_mask_roberta(input_text_roberta)
    
    st.subheader("Risultati delle previsioni (RoBERTa):")
    for pred in predictions_roberta:
        st.write(
            f"Token: {pred['token_str']} - "
            f"Probabilità: {pred['score']:.4f}\n"
            f"Sequence: {pred['sequence']}\n"
        )
        st.write("---")
    
    # Esempio di analisi con CLTK (se installato e importato)
    if nlp_lat is not None:
        st.subheader("Analisi Morfologica con CLTK (opzionale)")

        for i, pred in enumerate(predictions_roberta, start=1):
            # Ricostruisci la frase sostituendo <mask> con la parola predetta
            predicted_text = input_text_roberta.replace("<mask>", pred['token_str'])
            doc = nlp_lat(predicted_text)
            
            st.write(f"Frase {i}: {predicted_text}")
            for token in doc.tokens:
                st.write(
                    f"- Token: {token.string}\n"
                    f"  - Lemma: {token.lemma}\n"
                    f"  - UPOS: {token.upos}\n"
                    f"  - Morfologia: {token.features}\n"
                )
            st.write("---")
    else:
        st.warning("CLTK non installato (o non importato). Esegui 'pip install cltk' per abilitare l'analisi.")