import time
import streamlit as st
import string
from io import StringIO 
import json
from transformers import BertTokenizer, BertForMaskedLM

MAX_INPUT = 1000

model_names = [
            {   "name":"SGPT-125M", 
                "model":"Muennighoff/SGPT-125M-weightedmean-nli-bitfit",
                "mark":False,
                "class":"SGPTModel"},


            {   "name":"SGPT-5.8B",
                "model": "Muennighoff/SGPT-5.8B-weightedmean-msmarco-specb-bitfit" ,
                "fork_url":"https://github.com/taskswithcode/sgpt",
                "orig_author_url":"https://github.com/Muennighoff",
                "orig_author":"Niklas Muennighoff",
                "sota_info": {   
                                 "task":"#1 in multiple information retrieval & search tasks",
                                 "sota_link":"https://paperswithcode.com/paper/sgpt-gpt-sentence-embeddings-for-semantic",
                            },
                "paper_url":"https://arxiv.org/abs/2202.08904v5",
                "mark":True,
                "class":"SGPTModel"},

            {   "name":"SGPT-1.3B",
                "model": "Muennighoff/SGPT-1.3B-weightedmean-msmarco-specb-bitfit",
                "mark":False,
                "class":"SGPTModel"},

            {   "name":"sentence-transformers/all-MiniLM-L6-v2", 
                "model":"sentence-transformers/all-MiniLM-L6-v2",
                "fork_url":"https://github.com/taskswithcode/sentence_similarity_hf_model",
                "orig_author_url":"https://github.com/UKPLab",
                "orig_author":"Ubiquitous Knowledge Processing Lab",
                "sota_info": {   
                                 "task":"Nearly 4 million downloads from huggingface",
                                 "sota_link":"https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2"
                            },
                "paper_url":"https://arxiv.org/abs/1908.10084",
                "mark":True,
                "class":"HFModel"},

            ]



example_file_names = {
"Machine learning terms (30+ phrases)": "tests/small_test.txt",
"Customer feedback mixed with noise (50+ sentences)":"tests/larger_test.txt"
}


def construct_model_info_for_display():
    options_arr  = []
    markdown_str = "<div style=\"font-size:16px; color: #2f2f2f; text-align: left\"><br/><b>Models evaluated</b></div>"
    for node in model_names:
        options_arr .append(node["name"])
        if (node["mark"] == True):
            markdown_str += f"<div style=\"font-size:16px; color: #5f5f5f; text-align: left\">&nbsp;•&nbsp;Model:&nbsp;<a href=\'{node['paper_url']}\' target='_blank'>{node['name']}</a><br/>&nbsp;&nbsp;&nbsp;&nbsp;Code released by:&nbsp;<a href=\'{node['orig_author_url']}\' target='_blank'>{node['orig_author']}</a><br/>&nbsp;&nbsp;&nbsp;&nbsp;Model info:&nbsp;<a href=\'{node['sota_info']['sota_link']}\' target='_blank'>{node['sota_info']['task']}</a><br/>&nbsp;&nbsp;&nbsp;&nbsp;Forked <a href=\'{node['fork_url']}\' target='_blank'>code</a><br/><br/></div>"
    markdown_str += "<div style=\"font-size:12px; color: #9f9f9f; text-align: left\"><b>Note:</b><br/>•&nbsp;Uploaded files are loaded into non-persistent memory for the duration of the computation. They are not saved</div>"
    limit = "{:,}".format(MAX_INPUT)
    markdown_str += f"<div style=\"font-size:12px; color: #9f9f9f; text-align: left\">•&nbsp;User uploaded file has a maximum limit of {limit} sentences.</div>"
    return options_arr,markdown_str


st.set_page_config(page_title='TWC - Compare state-of-the-art models for Sentence Similarity task', page_icon="logo.jpg", layout='centered', initial_sidebar_state='auto',
            menu_items={
             'About': 'This app was created by taskswithcode. http://taskswithcode.com'
              })
col,pad = st.columns([85,15])

with col:
    st.image("long_form_logo_with_icon.png")


@st.experimental_memo
def load_model(model_name):
    try:
        ret_model = None
        for node in model_names:
            if (model_name.startswith(node["name"])):
                obj_class = globals()[node["class"]]
                ret_model = obj_class()
                ret_model.init_model(node["model"])
        assert(ret_model is not None)
    except Exception as e:
        st.error("Unable to load model:" + model_name + " " +  str(e))
        pass
    return ret_model

  
@st.experimental_memo
def cached_compute_similarity(sentences,_model,model_name,main_index):
    texts,embeddings = _model.compute_embeddings(sentences,is_file=False)
    results = _model.output_results(None,texts,embeddings,main_index)
    return results


def uncached_compute_similarity(sentences,_model,model_name,main_index):
    with st.spinner('Computing vectors for sentences'):
        texts,embeddings = _model.compute_embeddings(sentences,is_file=False)
        results = _model.output_results(None,texts,embeddings,main_index)
    #st.success("Similarity computation complete")
    return results

def run_test(model_name,sentences,display_area,main_index,user_uploaded):
    display_area.text("Loading model:" + model_name)
    model = load_model(model_name)
    display_area.text("Model " + model_name  + " load complete")
    try:
            if (user_uploaded):
                results = uncached_compute_similarity(sentences,model,model_name,main_index)
            else:
                display_area.text("Computing vectors for sentences")
                results = cached_compute_similarity(sentences,model,model_name,main_index)
                display_area.text("Similarity computation complete")
            return results
            
    except Exception as e:
        st.error("Some error occurred during prediction" + str(e))
        st.stop()
    return {}



    

def display_results(orig_sentences,main_index,results,response_info):
    main_sent = f"<div style=\"font-size:14px; color: #2f2f2f; text-align: left\">{response_info}<br/><br/></div>"
    main_sent += "<div style=\"font-size:14px; color: #6f6f6f; text-align: left\">Results sorted by cosine distance. Closest(1) to furthest(-1) away from main sentence</div>"
    main_sent += f"<div style=\"font-size:16px; color: #2f2f2f; text-align: left\"><b>Main sentence:</b>&nbsp;&nbsp;{orig_sentences[main_index]}</div>"
    body_sent = []
    download_data = {}
    for key in results:
        index = orig_sentences.index(key) + 1
        body_sent.append(f"<div style=\"font-size:16px; color: #2f2f2f; text-align: left\">{index}]&nbsp;{key}&nbsp;&nbsp;&nbsp;<b>{results[key]:.2f}</b></div>")
        download_data[key] =  f"{results[key]:.2f}" 
    main_sent = main_sent + "\n" + '\n'.join(body_sent)
    st.markdown(main_sent,unsafe_allow_html=True)
    st.session_state["download_ready"] = json.dumps(download_data,indent=4)


def init_session():
    st.session_state["download_ready"] = None    
    st.session_state["model_name"] = "ss_test"
    st.session_state["main_index"] = 1
    st.session_state["file_name"] = "default"
 
def main():
  init_session()
  st.markdown("<h4 style='text-align: center;'>Compare state-of-the-art models for Sentence Similarity task</h4>", unsafe_allow_html=True)


  try:
      
      
      with st.form('twc_form'):

        uploaded_file = st.file_uploader("Step 1. Upload text file(one sentence in a line) or choose an example text file below.", type=".txt")

        selected_file_index = st.selectbox(label='Example files ',  
                    options = list(dict.keys(example_file_names)), index=0,  key = "twc_file")
        st.write("")
        options_arr,markdown_str = construct_model_info_for_display()
        selected_model = st.selectbox(label='Step 2. Select Model',  
                    options = options_arr, index=0,  key = "twc_model")
        st.write("")
        main_index = st.number_input('Step 3. Enter index of sentence in file to make it the main sentence:',value=1,min_value = 1)
        st.write("")
        submit_button = st.form_submit_button('Run')

        
        input_status_area = st.empty()
        display_area = st.empty()
        if submit_button:
            start = time.time()
            if uploaded_file is not None:
                st.session_state["file_name"]  = uploaded_file.name
                sentences = StringIO(uploaded_file.getvalue().decode("utf-8")).read()
            else:
                st.session_state["file_name"]  = example_file_names[selected_file_index]
                sentences = open(example_file_names[selected_file_index]).read()
            sentences = sentences.split("\n")[:-1]
            if (len(sentences) < main_index):
                main_index = len(sentences)
                st.info("Selected sentence index is larger than number of sentences in file. Truncating to " + str(main_index)) 
            if (len(sentences) > MAX_INPUT):
                st.info(f"Input sentence count exceeds maximum sentence limit. First {MAX_INPUT} out of {len(sentences)} sentences chosen")
                sentences = sentences[:MAX_INPUT]
            st.session_state["model_name"] = selected_model
            st.session_state["main_index"] = main_index
            results = run_test(selected_model,sentences,display_area,main_index - 1,(uploaded_file is not None))
            display_area.empty()
            with display_area.container():
                response_info = f"Response time - {time.time() - start:.2f} secs for {len(sentences)} sentences"
                display_results(sentences,main_index - 1,results,response_info)
                #st.json(results)
      st.download_button(
         label="Download results as json",
         data= st.session_state["download_ready"] if st.session_state["download_ready"] != None else "",
         disabled = False if st.session_state["download_ready"] != None else True,
         file_name= (st.session_state["model_name"] + "_" +  str(st.session_state["main_index"]) + "_" + '_'.join(st.session_state["file_name"].split(".")[:-1]) + ".json").replace("/","_"),
         mime='text/json',
         key ="download" 
        )
      
      

  except Exception as e:
    st.error("Some error occurred during loading" + str(e))
    st.stop()  
	
  st.markdown(markdown_str, unsafe_allow_html=True)
  
 

if __name__ == "__main__":
   main()