# Hint: this cheatsheet is magic! https://cheat-sheet.streamlit.app/

import constants
import pandas as pd
import streamlit as st
from transformers import BertForSequenceClassification, AutoTokenizer


@st.cache_data
def convert_df(df):
    # IMPORTANT: Cache the conversion to prevent computation on every rerun
    return df.to_csv(index=None).encode("utf-8")


def compute_ALDi(inputs):
    return 0.5


input_type = st.sidebar.radio(
    "Select the input type:", [constants.CHOICE_FILE, constants.CHOICE_TEXT]
)

st.title(constants.TITLE)

if input_type == constants.CHOICE_TEXT:
    sent = st.text_input("Arabic Sentence:", placeholder="Enter an Arabic sentence.")

    # TODO: Check if this is needed!
    st.button("Submit")

    if sent:
        ALDi_score = compute_ALDi(sent)
        st.write(ALDi_score)

else:
    file = st.file_uploader("Upload a file", type=["txt"])
    if file is not None:
        df = pd.read_csv(file, sep="\t", header=None)
        df.columns = ["Sentence"]

        # TODO: Run the model
        df["ALDi"] = df["Sentence"].apply(lambda s: compute_ALDi(s))

        # A horizontal rule
        st.markdown("""---""")

        col1, col2 = st.columns([2, 3])

        with col1:
            # Add a download button
            csv = convert_df(df)

            st.download_button(
                label=":file_folder: Download predictions as CSV",
                data=csv,
                file_name="ALDi_scores.csv",
                mime="text/csv",
            )

            # Display the output
            st.dataframe(
                df,
                hide_index=True,
            )

        with col2:
            # TODO: Add the visualization
            st.image("https://static.streamlit.io/examples/dog.jpg")