import streamlit as st
import streamlit.components.v1 as components
import pandas as pd

@st.experimental_memo
def read_xsum_samples():
    df = pd.read_csv('./data/small_x_sum.csv')
    df = df[['ID', 'Content', 'Summary']]
    all_index = list(df.index)
    all_text = list(df.Content)
    text_dict = dict(zip(all_index, all_text))
    return text_dict

@st.experimental_memo
def read_squad_samples():
    df = pd.read_csv('./data/squad_sample.tsv', sep='\t')
    df = df[['id', 'context', 'question']]
    all_index = list(df.index)
    all_text = list(df.context)
    text_dict = dict(zip(all_index, all_text))
    return text_dict

@st.experimental_memo
def read_conll_samples():
    df = pd.read_pickle('./data/conll_df.pkl')
    df = df[['doc-text', 'ner-tag']]
    all_index = list(df.index)
    all_text = list(df['doc-text'])
    text_dict = dict(zip(all_index, all_text))
    return text_dict


def get_default_texts(chosen_datasets):
    if "bbc-xsum-summarization" in chosen_datasets:
        the_dict = read_xsum_samples()
    elif "conll-ner" in chosen_datasets:
        the_dict = read_conll_samples()
    else:
        the_dict = read_squad_samples()
    return the_dict


def display_output(all_outputs):
    all_tabs = [tab[0] for tab in all_outputs]
    tabs = st.tabs(all_tabs)
    for tab_index in range(len(all_tabs)):
        if all_outputs[tab_index][0] == 'Predicted Answer':
            with tabs[tab_index]:
                st.text_area('Best Answer found:', value=all_outputs[tab_index][1], disabled=True)
        elif all_outputs[tab_index][0] == 'Predicted Entities':
            with tabs[tab_index]:
                # st.markdown(all_outputs[tab_index][1], unsafe_allow_html=True)
                components.html(all_outputs[tab_index][1], scrolling=True)
        else:
            with tabs[tab_index]:
                st.text_area('Best Prediction found', value=all_outputs[tab_index][1], disabled=True)

if __name__ == '__main__':
    read_conll_samples()