#Import the libraries we know we'll need for the Generator.
import pandas as pd, spacy, nltk, numpy as np, re, os
from spacy.matcher import Matcher
from nltk.corpus import wordnet

#Attempting to fix the issue with spacy model in a more intuitive way.
try:
    nlp = spacy.load("en_core_web_lg")
except:
    script = "python -m spacy download en_core_web_lg"
    os.system("bash -c '%s'" % script)
nlp = spacy.load("en_core_web_lg")

#Import the libraries to support the model and predictions.
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TextClassificationPipeline
import lime
import torch
import torch.nn.functional as F
from lime.lime_text import LimeTextExplainer

#Import the libraries for human interaction and visualization.
import altair as alt
import streamlit as st
from annotated_text import annotated_text as ant

#Import functions needed to build dataframes of keywords from WordNet
from WNgen import *
from NLselector import *

@st.experimental_singleton
def set_up_explainer():
    class_names = ['negative', 'positive']
    explainer = LimeTextExplainer(class_names=class_names)
    return explainer

@st.experimental_singleton
def prepare_model():
    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
    model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
    pipe = TextClassificationPipeline(model=model, tokenizer=tokenizer, return_all_scores=True) 
    return tokenizer, model, pipe

@st.experimental_singleton
def prepare_lists():
    try:
        wordnet.synsets("bias")
    except:
        nltk.download('omw-1.4')
        nltk.download('wordnet')
    countries = pd.read_csv("Assets/Countries/combined-countries.csv")
    professions = pd.read_csv("Assets/Professions/soc-professions-2018.csv")
    word_lists = [list(countries.Words.apply(lambda x: x.lower())),list(professions.Words)]
    return countries, professions, word_lists

#Provide all the functions necessary to run the app
#get definitions for control flow in Streamlit
def get_def(word, POS=False):
    pos_options = ['NOUN','VERB','ADJ','ADV']
    m_word = re.sub("(\W\s|\s)","_",word)
    if POS in pos_options:
        seed_definitions = [syn.definition() for syn in wordnet.synsets(m_word, pos=getattr(wordnet, POS))]
    else:
        seed_definitions = [syn.definition() for syn in wordnet.synsets(m_word)]
    if len(seed_definitions) > 0:
        seed_definition = col1.selectbox("Which definition is most relevant?", seed_definitions, key= "WN_definition")
        if col1.button("Choose Definition"):
            col1.write("You've chosen a definition.")
            st.session_state.definition = seed_definition
            return seed_definition
        else:
            col1.write("Please choose a definition.")
    else:
        col1.error("The word you've chosen does not have a definition within WordNet.")
        
###Start coding the actual app###
st.set_page_config(layout="wide", page_title="Natural Language Counterfactuals (NLC)")
layouts = ['Natural Language Explanation', 'Lime Explanation', 'MultiNLC', 'MultiNLC + Lime', 'VizNLC']
alternatives = ['Similarity', 'Sampling (Random)', 'Sampling (Fixed)', 'Probability']
alt_choice = "Similarity"

#Content in the Sidebar.
st.sidebar.info('This is an interface for exploring how different interfaces for exploring natural language explanations (NLE) may appear to people. It is intended to allow individuals to provide feedback on specific versions, as well as to compare what one offers over others for the same inputs.')
layout = st.sidebar.selectbox("Select a layout to explore.", layouts)
alt_choice = st.sidebar.selectbox("Choose the way you want to display alternatives.", alternatives) #Commented out until we decide this is useful functionality.

#Set up the Main Area Layout
st.title('Natural Language Counterfactuals (NLC) Prototype')
st.subheader(f'Current Layout: {layout}')
text = st.text_input('Provide a sentence you want to evaluate.', placeholder = "I like you. I love you.", key="input")

#Prepare the model, data, and Lime. Set starting variables.
tokenizer, model, pipe = prepare_model()
countries, professions, word_lists = prepare_lists()
explainer = set_up_explainer()
text2 = ""
text3 = ""
cf_df = pd.DataFrame()
if 'definition' not in st.session_state:
    st.session_state.definition = "<(^_')>"
    
#Outline the various user interfaces we have built.

col1, col2, col3 = st.columns(3)    
if layout == 'Natural Language Explanation':
    with col1:
        if st.session_state.input != "":
            st.caption("This is the sentence you provided.")
            st.write(text)
            probability, sentiment = eval_pred(text, return_all=True)
            nat_lang_explanation = construct_nlexp(text,sentiment,probability)

if layout == 'Lime Explanation':
    with col1:
        #Use spaCy to make the sentence into a doc so we can do NLP.
        doc = nlp(st.session_state.input)
        #Evaluate the provided sentence for sentiment and probability.
        if st.session_state.input != "":
            st.caption("This is the sentence you provided.")
            st.write(text)
            probability, sentiment = eval_pred(text, return_all=True)
            options, lime = critical_words(st.session_state.input,options=True)
            nat_lang_explanation = construct_nlexp(text,sentiment,probability)
            st.write(" ")
            st.altair_chart(lime_viz(lime))
            
if layout == 'MultiNLC':
    with col1:
        #Use spaCy to make the sentence into a doc so we can do NLP.
        doc = nlp(st.session_state.input)
        #Evaluate the provided sentence for sentiment and probability.
        if st.session_state.input != "":
            st.caption("This is the sentence you provided.")
            st.write(text)
            probability, sentiment = eval_pred(text, return_all=True)
            options, lime = critical_words(st.session_state.input,options=True)
            nat_lang_explanation = construct_nlexp(text,sentiment,probability)

        #Allow the user to pick an option to generate counterfactuals from.
            option = st.radio('Which word would you like to use to generate alternatives?', options, key = "option")
            lc_option = option.lower()
            if (any(lc_option in sublist for sublist in word_lists)):
                st.write(f'You selected {option}. It matches a list.')
            elif option:
                st.write(f'You selected {option}. It does not match a list.')
                definition = get_def(option)
            else:
                st.write('Awaiting your selection.')

            if st.button('Generate Alternatives'):
                if lc_option in word_lists[0]:
                    cf_df = gen_cf_country(countries, doc, option)
                    st.success('Alternatives created.')
                elif lc_option in word_lists[1]:
                    cf_df = gen_cf_profession(professions, doc, option)
                    st.success('Alternatives created.')
                else:
                    with st.sidebar:
                        ant("Generating alternatives for",(option,"opt","#E0FBFB"), "with a definition of: ",(st.session_state.definition,"def","#E0FBFB"),".")
                    cf_df = cf_from_wordnet_df(option,text,seed_definition=st.session_state.definition)
                    st.success('Alternatives created.')

                if len(cf_df) != 0:
                    if alt_choice == "Similarity":
                        text2, text3 = get_min_max(cf_df, option)
                        col2.caption(f"This sentence is 'similar' to {option}.")
                        col3.caption(f"This sentence is 'not similar' to {option}.")
                    elif alt_choice == "Sampling (Random)":
                        text2, text3 = sampled_alts(cf_df, option)
                        col2.caption(f"This sentence is a random sample from the alternatives.")
                        col3.caption(f"This sentence is a random sample from the alternatives.")
                    elif alt_choice == "Sampling (Fixed)":
                        text2, text3 = sampled_alts(cf_df, option, fixed=True)
                        col2.caption(f"This sentence is a fixed sample of the alternatives.")
                        col3.caption(f"This sentence is a fixed sample of the alternatives.")
                    elif alt_choice == "Probability":
                        text2, text3 = abs_dif(cf_df, option)
                        col2.caption(f"This sentence is the closest prediction in the model.")
                        col3.caption(f"This sentence is the farthest prediction in the model.")
                    with st.sidebar:
                        st.info(f"Alternatives generated: {len(cf_df)}")
        
    with col2:
        if text2 != "":
            sim2 = cf_df.loc[cf_df['text'] == text2, 'similarity'].iloc[0]
            st.write(text2)
            probability2, sentiment2 = eval_pred(text2, return_all=True)
            nat_lang_explanation = construct_nlexp(text2,sentiment2,probability2)
            #st.info(f" Similarity Score: {np.round(sim2, 2)}, Num Checked: {len(cf_df)}") #for QA purposes

    with col3:
        if text3 != "":
            sim3 = cf_df.loc[cf_df['text'] == text3, 'similarity'].iloc[0]
            st.write(text3)
            probability3, sentiment3 = eval_pred(text3, return_all=True)
            nat_lang_explanation = construct_nlexp(text3,sentiment3,probability3)
            #st.info(f"Similarity Score: {np.round(sim3, 2)}, Num Checked: {len(cf_df)}") #for QA purposes

if layout == 'MultiNLC + Lime':
    with col1:

        #Use spaCy to make the sentence into a doc so we can do NLP.
        doc = nlp(st.session_state.input)
        #Evaluate the provided sentence for sentiment and probability.
        if st.session_state.input != "":
            st.caption("This is the sentence you provided.")
            st.write(text)
            probability, sentiment = eval_pred(text, return_all=True)
            options, lime = critical_words(st.session_state.input,options=True)
            nat_lang_explanation = construct_nlexp(text,sentiment,probability)
            st.write(" ")
            st.altair_chart(lime_viz(lime))

        #Allow the user to pick an option to generate counterfactuals from.
            option = st.radio('Which word would you like to use to generate alternatives?', options, key = "option")
            lc_option = option.lower()
            if (any(lc_option in sublist for sublist in word_lists)):
                st.write(f'You selected {option}. It matches a list.')
            elif option:
                st.write(f'You selected {option}. It does not match a list.')
                definition = get_def(option)
            else:
                st.write('Awaiting your selection.')

            if st.button('Generate Alternatives'):
                if lc_option in word_lists[0]:
                    cf_df = gen_cf_country(countries, doc, option)
                    st.success('Alternatives created.')
                elif lc_option in word_lists[1]:
                    cf_df = gen_cf_profession(professions, doc, option)
                    st.success('Alternatives created.')
                else:
                    with st.sidebar:
                        ant("Generating alternatives for",(option,"opt","#E0FBFB"), "with a definition of: ",(st.session_state.definition,"def","#E0FBFB"),".")
                    cf_df = cf_from_wordnet_df(option,text,seed_definition=st.session_state.definition)
                    st.success('Alternatives created.')

                if len(cf_df) != 0:
                    if alt_choice == "Similarity":
                        text2, text3 = get_min_max(cf_df, option)
                        col2.caption(f"This sentence is 'similar' to {option}.")
                        col3.caption(f"This sentence is 'not similar' to {option}.")
                    elif alt_choice == "Sampling (Random)":
                        text2, text3 = sampled_alts(cf_df, option)
                        col2.caption(f"This sentence is a random sample from the alternatives.")
                        col3.caption(f"This sentence is a random sample from the alternatives.")
                    elif alt_choice == "Sampling (Fixed)":
                        text2, text3 = sampled_alts(cf_df, option, fixed=True)
                        col2.caption(f"This sentence is a fixed sample of the alternatives.")
                        col3.caption(f"This sentence is a fixed sample of the alternatives.")
                    elif alt_choice == "Probability":
                        text2, text3 = abs_dif(cf_df, option)
                        col2.caption(f"This sentence is the closest prediction in the model.")
                        col3.caption(f"This sentence is the farthest prediction in the model.")
                    with st.sidebar:
                        st.info(f"Alternatives generated: {len(cf_df)}")

    with col2:
        if text2 != "":
            sim2 = cf_df.loc[cf_df['text'] == text2, 'similarity'].iloc[0]
            st.write(text2)
            probability2, sentiment2 = eval_pred(text2, return_all=True)
            nat_lang_explanation = construct_nlexp(text2,sentiment2,probability2)
            exp2 = explainer.explain_instance(text2, predictor, num_features=15, num_samples=2000)
            lime_results2 = exp2.as_list()
            st.write(" ")
            st.altair_chart(lime_viz(lime_results2))

    with col3:
        if text3 != "":
            sim3 = cf_df.loc[cf_df['text'] == text3, 'similarity'].iloc[0]
            st.write(text3)
            probability3, sentiment3 = eval_pred(text3, return_all=True)
            nat_lang_explanation = construct_nlexp(text3,sentiment3,probability3)
            exp3 = explainer.explain_instance(text3, predictor, num_features=15, num_samples=2000)
            lime_results3 = exp3.as_list()
            st.write(" ")
            st.altair_chart(lime_viz(lime_results3))
            
if layout == 'VizNLC':
    with col1:

        #Use spaCy to make the sentence into a doc so we can do NLP.
        doc = nlp(st.session_state.input)
        #Evaluate the provided sentence for sentiment and probability.
        if st.session_state.input != "":
            st.caption("This is the sentence you provided.")
            st.write(text)
            probability, sentiment = eval_pred(text, return_all=True)
            options, lime = critical_words(st.session_state.input,options=True)
            nat_lang_explanation = construct_nlexp(text,sentiment,probability)
            st.write(" ")
            st.altair_chart(lime_viz(lime))

        #Allow the user to pick an option to generate counterfactuals from.
            option = st.radio('Which word would you like to use to generate alternatives?', options, key = "option")
            lc_option = option.lower()
            if (any(lc_option in sublist for sublist in word_lists)):
                st.write(f'You selected {option}. It matches a list.')
            elif option:
                st.write(f'You selected {option}. It does not match a list.')
                definition = get_def(option)
            else:
                st.write('Awaiting your selection.')

            if st.button('Generate Alternatives'):
                if lc_option in word_lists[0]:
                    cf_df = gen_cf_country(countries, doc, option)
                    st.success('Alternatives created.')
                elif lc_option in word_lists[1]:
                    cf_df = gen_cf_profession(professions, doc, option)
                    st.success('Alternatives created.')
                else:
                    with st.sidebar:
                        ant("Generating alternatives for",(option,"opt","#E0FBFB"), "with a definition of: ",(st.session_state.definition,"def","#E0FBFB"),".")
                    cf_df = cf_from_wordnet_df(option,text,seed_definition=st.session_state.definition)
                    st.success('Alternatives created.')

                if len(cf_df) != 0:
                    if alt_choice == "Similarity":
                        text2, text3 = get_min_max(cf_df, option)
                        col2.caption(f"This sentence is 'similar' to {option}.")
                        col3.caption(f"This sentence is 'not similar' to {option}.")
                    elif alt_choice == "Sampling (Random)":
                        text2, text3 = sampled_alts(cf_df, option)
                        col2.caption(f"This sentence is a random sample from the alternatives.")
                        col3.caption(f"This sentence is a random sample from the alternatives.")
                    elif alt_choice == "Sampling (Fixed)":
                        text2, text3 = sampled_alts(cf_df, option, fixed=True)
                        col2.caption(f"This sentence is a fixed sample of the alternatives.")
                        col3.caption(f"This sentence is a fixed sample of the alternatives.")
                    elif alt_choice == "Probability":
                        text2, text3 = abs_dif(cf_df, option)
                        col2.caption(f"This sentence is the closest prediction in the model.")
                        col3.caption(f"This graph represents the {len(cf_df)} alternatives to {option}.")
                    with st.sidebar:
                        st.info(f"Alternatives generated: {len(cf_df)}")
    
    with col2:
        if text2 != "":
            sim2 = cf_df.loc[cf_df['text'] == text2, 'similarity'].iloc[0]
            st.write(text2)
            probability2, sentiment2 = eval_pred(text2, return_all=True)
            nat_lang_explanation = construct_nlexp(text2,sentiment2,probability2)
            exp2 = explainer.explain_instance(text2, predictor, num_features=15, num_samples=2000)
            lime_results2 = exp2.as_list()
            st.write(" ")
            st.altair_chart(lime_viz(lime_results2))

    with col3:
        if not cf_df.empty:
            single_nearest = alt.selection_single(on='mouseover', nearest=True)
            full = alt.Chart(cf_df).encode(
                alt.X('similarity:Q', scale=alt.Scale(zero=False)),
                alt.Y('pred:Q'),
                color=alt.Color('Categories:N', legend=alt.Legend(title="Color of Categories")),
                size=alt.Size('seed:O'),
                tooltip=('Categories','text','pred')
            ).mark_circle(opacity=.5).properties(width=450, height=450).add_selection(single_nearest)
            st.altair_chart(full)