import streamlit as st
from functions import *

st.set_page_config(page_title="Earnings Semantic Search", page_icon="🔎")
st.sidebar.header("Semantic Search")
st.markdown("## Earnings Semantic Search with SBert")

search_input = st.text_input(
        label='Enter Your Search Query, e.g "What challenges did the business face?"', key='search')
        
top_k = st.sidebar.slider("Number of Top Hits Generated",min_value=1,max_value=5,value=2)

window_size = st.sidebar.slider("Number of Sentences Generated in Search Response",min_value=1,max_value=5,value=3)
    
if "sen_df" not in st.session_state:
    st.session_state['sen_df'] = ''
    
if "earnings_passages" not in st.session_state:
    st.session_state["earnings_passages"] = ''

if search_input is not None:

    if any(st.session_state["sen_df"]) is not None and st.session_state["earnings_passages"] is not None:
    
        ## Save to a dataframe for ease of visualization
        sen_df = st.session_state['sen_df']
            
        passages = preprocess_plain_text(st.session_state['earnings_passages'],window_size=window_size)
        
        ##### Sematic Search #####
        # Encode the query using the bi-encoder and find potentially relevant passages
        corpus_embeddings = sbert.encode(passages, convert_to_tensor=True, show_progress_bar=True)
        question_embedding = sbert.encode(search_input, convert_to_tensor=True)
        question_embedding = question_embedding.cpu()
        hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k,score_function=util.dot_score)
        hits = hits[0]  # Get the hits for the first query
        
        ##### Re-Ranking #####
        # Now, score all retrieved passages with the cross_encoder
        cross_inp = [[search_input, passages[hit['corpus_id']]] for hit in hits]
        cross_scores = cross_encoder.predict(cross_inp)
        
        # Sort results by the cross-encoder scores
        for idx in range(len(cross_scores)):
            hits[idx]['cross-score'] = cross_scores[idx]
        
        # Output of top-3 hits from re-ranker
        hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
        
        score='cross-score'
        df = pd.DataFrame([(hit[score],passages[hit['corpus_id']]) for hit in hits[0:int(top_k)]],columns=['Score','Text'])
        df['Score'] = round(df['Score'],2)
        
        print(f'Test: {df}')
        
        def gen_annotated_text(para):
            tag_list = []
            for i in sent_tokenize(para):
                label = sen_df.loc[sen_df['text']==i, 'label'].values[0]
                if label == 'Negative':
                    tag_list.append((i,label,'#faa'))
                elif label == 'Positive':
                    tag_list.append((i,label,'#afa'))
                else:
                    tag_list.append((i,label,'#fea'))
            return tag_list  
        
        text_to_annotate = [gen_annotated_text(para) for para in df.Text.tolist()]
        
        first,second = text_to_annotate[0],text_to_annotate[-1]
        
        with st.container():
            annotated_text(*first)
            
        with st.container():
            annotated_text(*second)
            
    else:
        
        st.write('Please ensure you have entered the YouTube URL or uploaded the Earnings Call file')
else:

    st.write('Please ensure you have entered the YouTube URL or uploaded the Earnings Call file')