import gradio as gr

description = """Do you have a long document and a bunch of questions that can be answered given the data in this file? 
Fear not for this demo is for you. 
Upload your pdf, ask your questions and wait for the magic to happen.
DISCLAIMER: I do no have idea what happens to the pdfs that you upload and who has access to them so make sure there is nothing confidential there. 
On top of that, to speed up examples calculation, your query and the name of the document will be logged in.
"""
title = "QA answering from a pdf."

import numpy as np
import time
import hashlib
import torch
from transformers import AutoTokenizer, AutoModel, AutoModelForQuestionAnswering, pipeline
from tqdm import tqdm
import os
device = "cuda:0" if torch.cuda.is_available() else "cpu"
import textract
from scipy.special import softmax
import pandas as pd
from datetime import datetime
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/multi-qa-mpnet-base-dot-v1")
model = AutoModel.from_pretrained("sentence-transformers/multi-qa-mpnet-base-dot-v1").to(device).eval()
tokenizer_ans = AutoTokenizer.from_pretrained("deepset/roberta-large-squad2")
model_ans = AutoModelForQuestionAnswering.from_pretrained("deepset/roberta-large-squad2").to(device).eval()
if device == 'cuda:0':
    pipe = pipeline("question-answering",model_ans,tokenizer =tokenizer_ans,device = 0)
else:
    pipe = pipeline("question-answering",model_ans,tokenizer =tokenizer_ans)
    
def cls_pooling(model_output):
    return model_output.last_hidden_state[:,0]

def encode_query(query):
    encoded_input = tokenizer(query, truncation=True, return_tensors='pt').to(device)

    with torch.no_grad():
        model_output = model(**encoded_input, return_dict=True)

    embeddings = cls_pooling(model_output)

    return embeddings.cpu()


def encode_docs(docs,maxlen = 64, stride = 32):
    encoded_input = []
    embeddings = []
    spans = []
    file_names = []
    name, text = docs
    
    text = text.split(" ")
    if len(text) < maxlen:
        text = " ".join(text)
        
        encoded_input.append(tokenizer(temp_text, return_tensors='pt', truncation = True).to(device))
        spans.append(temp_text)
        file_names.append(name)

    else:
        num_iters = int(len(text)/maxlen)+1
        for i in range(num_iters):
            if i == 0:
                temp_text = " ".join(text[i*maxlen:(i+1)*maxlen+stride])
            else:
                temp_text = " ".join(text[(i-1)*maxlen:(i)*maxlen][-stride:] + text[i*maxlen:(i+1)*maxlen])

            encoded_input.append(tokenizer(temp_text, return_tensors='pt', truncation = True).to(device))
            spans.append(temp_text)
            file_names.append(name)

    with torch.no_grad():
        for encoded in tqdm(encoded_input): 
            model_output = model(**encoded, return_dict=True)
            embeddings.append(cls_pooling(model_output))
    
    embeddings = np.float32(torch.stack(embeddings).transpose(0, 1).cpu())
    
    np.save("emb_{}.npy".format(name),dict(zip(list(range(len(embeddings))),embeddings))) 
    np.save("spans_{}.npy".format(name),dict(zip(list(range(len(spans))),spans)))
    np.save("file_{}.npy".format(name),dict(zip(list(range(len(file_names))),file_names)))
    
    return embeddings, spans, file_names
   
def predict(query,data):
    name_to_save = data.name.split("/")[-1].split(".")[0][:-8]

    st = str([query,name_to_save])
    st_hashed = str(hashlib.sha256(st.encode()).hexdigest()) #just to speed up examples load
    hist = st + " " + st_hashed 
    now = datetime.now()
    current_time = now.strftime("%H:%M:%S")
    try:
        df = pd.read_csv("{}.csv".format(st_hashed))
        return df
    except Exception as e:
        print(e)
        print(st)

    if name_to_save+".txt" in os.listdir():
        doc_emb = np.load('emb_{}.npy'.format(name_to_save),allow_pickle='TRUE').item()
        doc_text = np.load('spans_{}.npy'.format(name_to_save),allow_pickle='TRUE').item()
        file_names_dicto = np.load('file_{}.npy'.format(name_to_save),allow_pickle='TRUE').item()
        
        doc_emb = np.array(list(doc_emb.values())).reshape(-1,768)
        doc_text = list(doc_text.values())
        file_names = list(file_names_dicto.values())
    
    else:
        text = textract.process("{}".format(data.name)).decode('utf8')
        text = text.replace("\r", " ")
        text = text.replace("\n", " ")
        text = text.replace(" . "," ")
        
        doc_emb, doc_text, file_names = encode_docs((name_to_save,text),maxlen = 64, stride = 32)
        
        doc_emb = doc_emb.reshape(-1, 768)
        with open("{}.txt".format(name_to_save),"w",encoding="utf-8") as f:
            f.write(text)
    start = time.time()
    query_emb = encode_query(query)
    
    scores = np.matmul(query_emb, doc_emb.transpose(1,0))[0].tolist()
    doc_score_pairs = list(zip(doc_text, scores, file_names))
    doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)
    k = 5
    probs_sum = 0
    probs = softmax(sorted(scores,reverse = True)[:k])
    table = {"Passage":[],"Answer":[],"Probabilities":[],"Source":[]}
    
    for i, (passage, _, names) in enumerate(doc_score_pairs[:k]):
        passage = passage.replace("\n","")
        passage = passage.replace(" . "," ")
        
        if probs[i] > 0.1 or (i < 3 and probs[i] > 0.05): #generate answers for more likely passages but no less than 2
            QA = {'question':query,'context':passage}
            ans = pipe(QA)
            probabilities = "P(a|p): {}, P(a|p,q): {}, P(p|q): {}".format(round(ans["score"],5), 
                                                                          round(ans["score"]*probs[i],5), 
                                                                          round(probs[i],5))
            passage = passage.replace(str(ans["answer"]),str(ans["answer"]).upper()) 
            table["Passage"].append(passage)
            table["Passage"].append("---")
            table["Answer"].append(str(ans["answer"]).upper())
            table["Answer"].append("---")
            table["Probabilities"].append(probabilities)
            table["Probabilities"].append("---")
            table["Source"].append(names)
            table["Source"].append("---")
        else:
            table["Passage"].append(passage)
            table["Passage"].append("---")
            table["Answer"].append("no_answer_calculated")
            table["Answer"].append("---")
            table["Probabilities"].append("P(p|q): {}".format(round(probs[i],5)))
            table["Probabilities"].append("---")
            table["Source"].append(names)
            table["Source"].append("---")
    df = pd.DataFrame(table)
    print("time: "+ str(time.time()-start))
    
    with open("HISTORY.txt","a", encoding = "utf-8") as f:
        f.write(hist)
        f.write(" " + str(current_time))
        f.write("\n")
        f.close()
    df.to_csv("{}.csv".format(st_hashed), index=False)
    
    return df

iface = gr.Interface(examples = [
        ["How high is the highest mountain?","China.pdf"], 
        ["Where does UK prime minister live?","London.pdf"]
    ],
    
    fn =predict,
    inputs = [gr.inputs.Textbox(default="What is Open-domain question answering?"),
              gr.inputs.File(),
    ],
    outputs = [
        gr.outputs.Dataframe(),
            ],
    description=description,
    title = title,
allow_flagging ="manual",flagging_options = ["correct","wrong"],
                     allow_screenshot=False)

iface.launch(share = True,enable_queue=True, show_error =True)