import streamlit as st
import torch
from transformers import AutoModel, AutoTokenizer
import pandas as pd
import io

st.set_page_config(page_title="ME2-BERT Moralities", layout="centered")

st.markdown(
    """
    <style>
    [data-testid="stFileUploadDropzone"] p,
    [data-baseweb="file-uploader-dropzone"] p {
        visibility: hidden !important;
        display: none !important;
    }
    [data-testid="stFileUploadDropzone"] p:before,
    [data-baseweb="file-uploader-dropzone"] p:before {
        content: "Drag and drop file here or browse (max 10 MB)";
        visibility: visible;
        display: block;
        color: rgb(120, 119, 119);
    }
    </style>
    <script>
    document.addEventListener('DOMContentLoaded', function() {
      let elements = document.querySelectorAll('span, p, div');
      elements.forEach(el => {
        if (el.textContent.includes('200 MB') || el.textContent.includes('200MB')) {
          el.style.display = 'none';
        }
      });
    });
    </script>
    """,
    unsafe_allow_html=True
)

MODEL_NAME = "lorenzozan/ME2-BERT"
MF_MAPPING = {
    "CH": "CARE/HARM",
    "FC": "FAIRNESS/CHEATING",
    "LB": "LOYALTY/BETRAYAL",
    "AS": "AUTHORITY/SUBVERSION",
    "PD": "PURITY/DEGRADATION"
}

@st.cache_resource
def load_model_and_tokenizer(model_name: str):
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
    model.eval()
    return model, tokenizer

model, tokenizer = load_model_and_tokenizer(MODEL_NAME)

def predict_moralities(text_list, max_seq_length=200):
    encoded_input = tokenizer(
        text_list,
        max_length=max_seq_length,
        padding="max_length",
        truncation=True,
        return_tensors="pt",
    )
    with torch.no_grad():
        outputs = model(
            input_ids=encoded_input["input_ids"],
            attention_mask=encoded_input["attention_mask"],
            return_dict=True
        )
    results = []
    for i in range(len(text_list)):
        row_scores = outputs[i]
        mapped_scores = {}
        for short_lbl, val in row_scores.items():
            mapped_scores[MF_MAPPING.get(short_lbl, short_lbl)] = float(val)
        results.append(mapped_scores)
    return results

st.title("ME2-BERT: A BERT-based model for Moral Foundations Prediction")
st.markdown(
"""
> Moralities, emotions, and events are complex aspects of human cognition, which are often treated separately since capturing their combined effects is challenging, especially due to the lack of annotated data. Leveraging their interrelations hence becomes crucial for advancing the understanding of human moral behaviors.
In this work, we propose ME²-BERT, the first holistic  framework for fine-tuning a pre-trained language model like BERT to the task of moral foundation prediction. ME²-BERT integrates events and emotions for learning domain-invariant morality-relevant text representations. 
Our extensive experiments show that ME²-BERT outperforms existing state-of-the-art methods for moral foundation prediction, with an average percentage increase up to 35% in the out-of-domain scenario.

**ME2-BERT** is a model for **Moral Foundation Prediction** based on  [Moral Foundation Theory (MFT)](https://moralfoundations.org/). 
Use this web app app to predict moral foundations (CARE/HARM, FAIRNESS/CHEATING, AUTHORITY/SUBVERSION, LOYALTY/BETRAYAL, PURITY/DEGRADATION) 
by entering a text or by uploading a csv file with at least a 'text' column. 

[Paper](https://aclanthology.org/2025.coling-main.638.pdf) | [Source code](https://github.com/lorenzozangari/ME2-BERT)
"""
)

def do_prediction():
    text_stripped = st.session_state["user_text_box"].strip()
    if text_stripped:
        single_result = predict_moralities([text_stripped])
        st.session_state["prediction_result"] = single_result[0]
    else:
        st.session_state["prediction_result"] = None

if "prediction_result" not in st.session_state:
    st.session_state["prediction_result"] = None

st.text_area(
    "Enter your text here:", 
    height=150, 
    key="user_text_box", 
    on_change=do_prediction
)

if st.button("Predict"):
    do_prediction()

if st.session_state["prediction_result"]:
    st.write("**Predicted moral foundations:**")
    st.json(st.session_state["prediction_result"])


st.markdown("## CSV File Upload")
st.write("You can upload a space-separated CSV file with a 'text' column. The file must have at most 50 rows.")

uploaded_file = st.file_uploader("Upload CSV (space-separated)", type=["csv"])
if uploaded_file is not None:
    df = pd.read_csv(uploaded_file, sep=' ')
    if len(df) > 50:
        st.error("The CSV file must have at most 50 rows.")
    else:
        # Controllo esistenza colonna 'text'
        if 'text' not in df.columns:
            st.error("The CSV file must contain a column named 'text'.")
        else:
            text_list = df['text'].tolist()
            predictions = predict_moralities(text_list)
            
       
            st.write("### Predictions for each row:")
            for i, row_text in enumerate(text_list):
                st.write(f"Row {i+1}: {row_text}")
                st.json(predictions[i])
            
            results_for_download = []
            for txt, preds in zip(text_list, predictions):
                row_data = {"text": txt}
                row_data.update(preds)  
                results_for_download.append(row_data)
            
            df_results = pd.DataFrame(results_for_download)
            
         
            csv_buffer = io.StringIO()
            df_results.to_csv(csv_buffer, index=False)
            csv_data = csv_buffer.getvalue()
            
            # Bottone per il download
            st.download_button(
                label="Download predictions as CSV",
                data=csv_data,
                file_name="predictions.csv",
                mime="text/csv"
            )
            
            # Text area per poter copiare il CSV
            st.text_area("Copy the results as CSV text", value=csv_data, height=200)
# ---------------------------------------------------------------

st.markdown(
"""
## References

```
@inproceedings{zangari-etal-2025-me2,
    title = "{ME}2-{BERT}: Are Events and Emotions what you need for Moral Foundation Prediction?",
    author = "Zangari, Lorenzo  and
      Greco, Candida M.  and
      Picca, Davide  and
      Tagarelli, Andrea",
      publisher = "Association for Computational Linguistics",
      url = "https://aclanthology.org/2025.coling-main.638/",
      pages = "9516--9532",
      abstract = "Moralities, emotions, and events are complex aspects of human cognition, which are often treated separately since capturing their combined effects is challenging, especially due to the lack of annotated data. Leveraging their interrelations hence becomes crucial for advancing the understanding of human moral behaviors. In this work, we propose ME2-BERT, the first holistic framework for fine-tuning a pre-trained language model like BERT to the task of moral foundation prediction. ME2-BERT integrates events and emotions for learning domain-invariant morality-relevant text representations. Our extensive experiments show that ME2-BERT outperforms existing state-of-the-art methods for moral foundation prediction, with an average increase up to 35{\%} in the out-of-domain scenario."
}
```
"""
)