import json

import gradio
import datasets

import numpy
import pandas
import sentence_transformers
import faiss

model = sentence_transformers.SentenceTransformer('allenai-specter')

full_data = datasets.load_dataset("ccm/publications")['train'].to_pandas()

filter = ["\"abstract\": null" in json.dumps(bibdict) for bibdict in full_data['bib_dict'].values]
data = full_data[~pandas.Series(filter)]
data.reset_index(inplace=True)

dimensionality = len(data['embedding'][0])
index = faiss.IndexFlatL2(dimensionality) 

vectors = numpy.stack(data['embedding'].to_list(), axis=0)

index.add(vectors)

def search(query, k):
    query = numpy.expand_dims(model.encode(query), axis=0)
    _, I = index.search(query, k)
    top_five = data.loc[I[0]]
    search_results = ""

    for i in range(k):
        search_results += '### ' + top_five["bib_dict"].values[i]["title"] + '\n'
        search_results += top_five["bib_dict"].values[i]["citation"]
        if top_five["pub_url"].values[i] is not None:
            search_results += ", [Paper](" + top_five["pub_url"].values[i] + ")"
        search_results += "\t\n```\n"
        search_results += json.dumps(top_five["bibtex"].values[i], indent=4).replace('\\n', '\n').replace('\\t', '\t').strip("\"")
        search_results += "\t\n```\n"
    return search_results


with gradio.Blocks() as demo:
    with gradio.Group():
        query = gradio.Textbox(placeholder="Enter search terms...", show_label=False, lines=1, max_lines=1)
        with gradio.Accordion("Settings", open=False):
            k = gradio.Number(5.0, label="Number of results", precision=0)
    results = gradio.Markdown()
    query.change(fn=search, inputs=[query, k], outputs=results)
    k.change(fn=search, inputs=[query, k], outputs=results)

demo.launch(debug=True)