import requests
from sentence_transformers import SentenceTransformer, CrossEncoder, util
import os, re
import torch
from rank_bm25 import BM25Okapi
from sklearn.feature_extraction import _stop_words
import string
from tqdm.autonotebook import tqdm
import numpy as np
from bs4 import BeautifulSoup
from nltk import sent_tokenize
import time
from newspaper import Article
import base64 
import docx2txt
from io import StringIO
from PyPDF2 import PdfFileReader
import validators

nltk.download('punkt')

from nltk import sent_tokenize

warnings.filterwarnings("ignore")

def extract_text_from_url(url: str):
    
    '''Extract text from url'''
    
    article = Article(url)
    article.download()
    article.parse()
    
    # get text
    text = article.text
    
    # get article title
    title = article.title
    
    return title, text

def extract_text_from_file(file):
    
    '''Extract text from uploaded file'''

    # read text file
    if file.type == "text/plain":
        # To convert to a string based IO:
        stringio = StringIO(file.getvalue().decode("utf-8"))

        # To read file as string:
        file_text = stringio.read()

    # read pdf file
    elif file.type == "application/pdf":
        pdfReader = PdfFileReader(file)
        count = pdfReader.numPages
        all_text = ""
        
        for i in range(count):
            page = pdfReader.getPage(i)
            all_text += page.extractText()
            file_text = all_text

    # read docx file
    elif (
        file.type
        == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
    ):
        file_text = docx2txt.process(file)

    return file_text
    
def preprocess_plain_text(text,window_size=window_size):

    text = text.encode("ascii", "ignore").decode()  # unicode
    text = re.sub(r"https*\S+", " ", text)  # url
    text = re.sub(r"@\S+", " ", text)  # mentions
    text = re.sub(r"#\S+", " ", text)  # hastags
    #text = re.sub(r"\s{2,}", " ", text)  # over spaces
    text = re.sub("[^.,!?%$A-Za-z0-9]+", " ", text)  # special characters except .,!?
    
    #break into lines and remove leading and trailing space on each
    lines = [line.strip() for line in text.splitlines()]
    
    # #break multi-headlines into a line each
    chunks = [phrase.strip() for line in lines for phrase in line.split("  ")]
    
    # # drop blank lines
    text = '\n'.join(chunk for chunk in chunks if chunk)
    
    ## We split this article into paragraphs and then every paragraph into sentences
    paragraphs = []
    for paragraph in text.replace('\n',' ').split("\n\n"):
        if len(paragraph.strip()) > 0:
            paragraphs.append(sent_tokenize(paragraph.strip()))

    #We combine up to 3 sentences into a passage. You can choose smaller or larger values for window_size
    #Smaller value: Context from other sentences might get lost
    #Lager values: More context from the paragraph remains, but results are longer
    window_size = window_size
    passages = []
    for paragraph in paragraphs:
        for start_idx in range(0, len(paragraph), window_size):
            end_idx = min(start_idx+window_size, len(paragraph))
            passages.append(" ".join(paragraph[start_idx:end_idx]))
            
    print("Paragraphs: ", len(paragraphs))
    print("Sentences: ", sum([len(p) for p in paragraphs]))
    print("Passages: ", len(passages))

    return passages

@st.cache(allow_output_mutation=True)
def bi_encoder(bi_enc,passages):
    
    #We use the Bi-Encoder to encode all passages, so that we can use it with sematic search
    bi_encoder = SentenceTransformer(bi_enc)
    #Start the multi-process pool on all available CUDA devices
    pool = bi_encoder.start_multi_process_pool()

    #Compute the embeddings using the multi-process pool
    print('encoding passages into a vector space...')
    corpus_embeddings = bi_encoder.encode_multi_process(passages, pool)
    print("Embeddings computed. Shape:", corpus_embeddings.shape)

    #Optional: Stop the proccesses in the pool
    bi_encoder.stop_multi_process_pool(pool)
    
    return corpus_embeddings
    
@st.cache(allow_output_mutation=True)
def cross_encoder():
    
    #The bi-encoder will retrieve 100 documents. We use a cross-encoder, to re-rank the results list to improve the quality
    cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-12-v2')
    return cross_encoder
    
@st.cache(allow_output_mutation=True)
def bm25_tokenizer(text):
    
# We also compare the results to lexical search (keyword search). Here, we use 
# the BM25 algorithm which is implemented in the rank_bm25 package.
# We lower case our text and remove stop-words from indexing
    tokenized_doc = []
    for token in text.lower().split():
        token = token.strip(string.punctuation)

        if len(token) > 0 and token not in _stop_words.ENGLISH_STOP_WORDS:
            tokenized_doc.append(token)
    return tokenized_doc

@st.cache(allow_output_mutation=True)
def bm25_api(passages):

    tokenized_corpus = []
    print('implementing BM25 algo for lexical search..')
    for passage in tqdm(passages):
        tokenized_corpus.append(bm25_tokenizer(passage))

    bm25 = BM25Okapi(tokenized_corpus)
    
    return bm25

bi_enc_options = ["multi-qa-mpnet-base-dot-v1","all-mpnet-base-v2","multi-qa-MiniLM-L6-cos-v1"]

# This function will search all wikipedia articles for passages that
# answer the query
def search(query, top_k=2):
    st.write(f"Search Query: {query}")
    
    st.write("Document Header: ")

    ##### BM25 search (lexical search) #####
    bm25_scores = bm25.get_scores(bm25_tokenizer(query))
    top_n = np.argpartition(bm25_scores, -5)[-5:]
    bm25_hits = [{'corpus_id': idx, 'score': bm25_scores[idx]} for idx in top_n]
    bm25_hits = sorted(bm25_hits, key=lambda x: x['score'], reverse=True)
    
    st.write(f"Top-{top_k} lexical search (BM25) hits")
    for hit in bm25_hits[0:top_k]:
        st.write("\t{:.3f}\t{}".format(hit['score'], passages[hit['corpus_id']].replace("\n", " ")))

    ##### Sematic Search #####
    # Encode the query using the bi-encoder and find potentially relevant passages
    question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
    question_embedding = question_embedding.gpu()
    hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k)
    hits = hits[0]  # Get the hits for the first query

    ##### Re-Ranking #####
    # Now, score all retrieved passages with the cross_encoder
    cross_inp = [[query, 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 bi-encoder
    st.markdown("\n-------------------------\n")
    st.write(f"Top-{top_k} Bi-Encoder Retrieval hits")
    hits = sorted(hits, key=lambda x: x['score'], reverse=True)
    for hit in hits[0:top_k]:
        st.write("\t{:.3f}\t{}".format(hit['score'], passages[hit['corpus_id']].replace("\n", " ")))

    # Output of top-3 hits from re-ranker
    st.markdown("\n-------------------------\n")
    st.write(f"Top-{top_k} Cross-Encoder Re-ranker hits")
    hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
    for hit in hits[0:top_k]:
        st.write("\t{:.3f}\t{}".format(hit['cross-score'], passages[hit['corpus_id']].replace("\n", " ")))
        
#Streamlit App
    
st.title("Semantic Search with Retrieve & Rerank 📝")

window_size = st.sidebar.slider("Paragraph Window Size",min_value=1,max_value=10,value=3)

bi_encoder_type = st.sidebar.selectbox(
    "Bi-Encoder", options=bi_enc_options
)

top_k = st.sidebar.slider("Number of Top Hits Generated",min_value=1,max_value=5,value=2)

st.markdown(
    """The app supports asymmetric Semantic search which seeks to improve search accuracy of documents/URL by understanding the content of the search query in contrast to traditional search engines which only find documents based on lexical matches, semantic search can also find synonyms. 
    The idea behind semantic search is to embed all entries in your corpus, whether they be sentences, paragraphs, or documents, into a vector space. At search time, the query is embedded into the same vector space and the closest embeddings from your corpus are found. These entries should have a high semantic overlap with the query.
    The all-* models where trained on all available training data (more than 1 billion training pairs) and are designed as general purpose models. The all-mpnet-base-v2 model provides the best quality, while all-MiniLM-L6-v2 is 5 times faster and still offers good quality. The models used have been trained on broad datasets, however, if your document/corpus is specialised, such as for science or economics, the results returned might be unsatisfactory.
     
    There models available to choose from:""")
    
st.markdown(
    """Model Source: 
    Bi-Encoders - [multi-qa-mpnet-base-dot-v1](https://huggingface.co/sentence-transformers/multi-qa-mpnet-base-dot-v1), [all-mpnet-base-v2](https://huggingface.co/sentence-transformers/all-mpnet-base-v2), [multi-qa-MiniLM-L6-cos-v1](https://huggingface.co/sentence-transformers/multi-qa-MiniLM-L6-cos-v1) and [all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2)
    Cross-Encoder - [cross-encoder/ms-marco-MiniLM-L-12-v2](https://huggingface.co/cross-encoder/ms-marco-MiniLM-L-12-v2)
    
    Code and App Inspiration Source:
    [Sentence Transformers](https://www.sbert.net/examples/applications/retrieve_rerank/README.html)
    
    Quick summary of the purposes of a Bi and Cross-encoder below, the image and info were adapted from [www.sbert.net](https://www.sbert.net/examples/applications/semantic-search/README.html):
    
    Bi-Encoder (Retrieval): The Bi-encoder is responsible for independently embedding the sentences and search queries into a vector space. The result is then passed to the cross-encoder for checking the relevance/similarity between the query and sentences.
    
    Cross-Encoder (Re-Ranker): A re-ranker based on a Cross-Encoder can substantially improve the final results for the user. The query and a possible document is passed simultaneously to transformer network, which then outputs a single score between 0 and 1 indicating how relevant the document is for the given query. The cross-encoder further boost the performance, especially when you search over a corpus for which the bi-encoder was not trained for.   """
)

st.image('encoder.png', caption='Retrieval and Re-Rank')

st.markdown("""   
    In order to use the app:
    - Select the preferred Sentence Transformer model (Bi-Encoder).
    - Select the number of sentences per paragraph to partition your corpus (Window-Size), if you choose a small value the context from the other sentences might get lost and for larger values the results might take longer to generate.
    - Paste the URL with your corpus or upload your preferred document in txt, pdf or Word format
    - Semantic Search away!! """
)

st.markdown("---")

url_text = st.text_input("Please Enter a url here")

st.markdown(
    "<h3 style='text-align: center; color: red;'>OR</h3>",
    unsafe_allow_html=True,
)

st.markdown(
    "<h3 style='text-align: center; color: red;'>OR</h3>",
    unsafe_allow_html=True,
)

upload_doc = st.file_uploader(
    "Upload a .txt, .pdf, .docx file"
)

search_query = st.text_input("Please Enter your search query here")

if validators.url(url_text):
    #if input is URL
    title, text = extract_text_from_url(url_text)
    passages = preprocess_plain_text(text,window_size=window_size)
    
elif upload_doc:
    
    passages = preprocess_plain_text(extract_text_from_file(upload_doc),window_size=window_size)

search = st.button("Search")

if search:
    if bi_encoder_type:

        with st.spinner(
            text=f"Loading {bi_encoder_type} Bi-Encoder and embedding document into vector space. This might take a few seconds depending on the length of your document..."
        ):
            corpus_embeddings = bi_encoder(bi_encoder_type,passages)
            cross_encoder = cross_encoder()
            bm25 = bm25_api(passages)
            
        with st.spinner(
            text="Embedding completed, searching for relevant text for given query and hits..."):
            
            search(search_query,top_k)