import streamlit as st
import pandas as pd
import string


DATAPATH = 'data/'


def prep_data():
    conversations = pd.read_csv(
        DATAPATH + 'movie_conversations.tsv',
        encoding='utf-8-sig', sep="\t", header=None,
        names=['character_ID_1', 'character_ID_2', 'movie_ID', 'utterances']
    )
    lines = pd.read_csv(
        DATAPATH + 'movie_lines.tsv', encoding='utf-8-sig', index_col='line_ID',
        on_bad_lines='skip', sep='\t', header=None,
        names=['line_ID', 'character_ID', 'movie_ID', 'name', 'text']
    )
    movies = pd.read_csv(
        DATAPATH + 'movie_titles_metadata.tsv',
        encoding='utf-8-sig', sep='\t', header=None,
        names=['movie_ID', 'title', 'year', 'rating', 'votes', 'genres']
    ).dropna()
    # string to list of strings
    conversations['utterances'] = conversations['utterances'].apply(
        lambda x: x.translate(str.maketrans('', '', string.punctuation)).split()
    )
    # concat dialogs
    keys = set(lines.index)
    conversations['dialog'] = conversations['utterances'].apply(
        lambda lst: [lines.loc[id]['text'] for id in lst if id in keys]
    )
    conversations = conversations.merge(movies, on='movie_ID', how='inner')
    conversations = conversations[['dialog', 'title']]
    return conversations


def app(data):
    def film_change():
        st.session_state.annotations = {}
        st.session_state.dialog_num = 0
        st.session_state.phrase_num = 1
        st.session_state.cur_ch = []
        st.session_state.data = data[data['title'] == st.session_state.film]
        st.session_state.dialog_quant = len(st.session_state.data)
        st.session_state.current_dialog = st.session_state.data.iloc[
            st.session_state.dialog_num
        ]['dialog']
        st.session_state.dialog_length = len(st.session_state.current_dialog)

    def annotate():
        st.session_state.annotations[
            st.session_state.current_dialog[st.session_state.phrase_num]
        ] = st.session_state.cur_ch
        st.session_state.phrase_num += 1
        st.session_state.cur_ch = []
        if st.session_state.phrase_num == st.session_state.dialog_length:
            st.session_state.dialog_num += 1
            st.session_state.phrase_num = 1
            st.session_state.current_dialog = st.session_state.data.iloc[
                st.session_state.dialog_num
            ]['dialog']
            st.session_state.dialog_length = len(st.session_state.current_dialog)

    def em_click(emotion):
        if emotion in st.session_state.cur_ch:
            st.session_state.cur_ch.remove(emotion)
        else:
            st.session_state.cur_ch.append(emotion)

    st.write(
        """
        ## Phrase Annotation

        Welcome to the dialog annotation tool! Label the emotion of phrase!
        """
    )

    if "annotations" not in st.session_state:
        st.session_state.annotations = {}
        st.session_state.dialog_num = 0
        st.session_state.phrase_num = 1
        st.session_state.cur_ch = []
        st.session_state.film = data['title'].unique()[0]
        st.session_state.data = data[data['title'] == st.session_state.film]
        st.session_state.dialog_quant = len(st.session_state.data)
        st.session_state.current_dialog = st.session_state.data.iloc[
            st.session_state.dialog_num
        ]['dialog']
        st.session_state.dialog_length = len(st.session_state.current_dialog)

    st.session_state.film = st.selectbox(
        'Select a movie to annotate',
        options=data['title'].unique(),
        on_change=film_change
    )

    st.write("")
    col1, col2 = st.columns(2)
    with col1:
        if st.session_state.dialog_num != st.session_state.dialog_quant:
            i = st.session_state.phrase_num
            st.write("### Previous phrase")
            st.write(st.session_state.current_dialog[i-1])
            st.write("### Phrase")
            st.write(st.session_state.current_dialog[i])
            st.write("")
            st.write(
                f'###### You have selected the following emotions:\
                {st.session_state.cur_ch}'
            )
            c1, c2 = st.columns(2)
            with c1:
                st.button(
                    "Anger 🤬", on_click=em_click,
                    args=['anger'], use_container_width=True
                )
                st.button(
                    "Disgust 🤢", on_click=em_click,
                    args=['disgust'], use_container_width=True
                )
                st.button(
                    "Fear 😨", on_click=em_click,
                    args=['fear'], use_container_width=True
                )
                st.button(
                    "Surprise 😲", on_click=em_click,
                    args=['surprise'], use_container_width=True
                )
            with c2:
                st.button(
                    "Joy 😀", on_click=em_click,
                    args=['joy'], use_container_width=True
                )
                st.button(
                    "Neutral 😐", on_click=em_click,
                    args=['neutral'], use_container_width=True
                )
                st.button(
                    "Sadness 😭", on_click=em_click,
                    args=['sadness'], use_container_width=True
                )
            st.button(
                'Annotate', on_click=annotate,
                use_container_width=True
            )
        else:
            st.success(
                "🎉 Done! All dialogues in this film annotated."
            )
    with col2:
        st.write("### Annotations")
        st.write(st.session_state.annotations)


if __name__ == "__main__":
    data = prep_data()
    app(data)