import streamlit as st  #Web App
import urllib
from lxml import html
import requests
import re
import os
from stqdm import stqdm
import time
import shutil
from PIL import Image

import pickle
docs = None
api_key = ' '

st.set_page_config(layout="wide")

image = Image.open('arxiv_decode.png')
st.image(image, width=1000)


#title
st.title("Answering questions from scientific papers")
st.markdown("##### This tool will allow you to ask questions and get based on scientific papers. It uses OpenAI's GPT models, and you must have your own API key. Each query is about 10k tokens, which costs about $0.20.")
st.markdown("##### Current version searches on ArXiv papers only. 🚧Under development🚧")
st.markdown("Used libraries:\n * [PaperQA](https://github.com/whitead/paper-qa) \n* [langchain](https://github.com/hwchase17/langchain)")

api_key_url = 'https://help.openai.com/en/articles/4936850-where-do-i-find-my-secret-api-key'

api_key = st.text_input('OpenAI API Key',
         placeholder='sk-...', 
         help=f"['What is that?']({api_key_url})",
         type="password")     

os.environ["OPENAI_API_KEY"] = f"{api_key}" # 
if len(api_key) != 51:
    st.warning('Please enter a valid OpenAI API key.', icon="⚠️")



def call_arXiv_API(search_query, search_by='all', sort_by='relevance', max_results='10', folder_name='arxiv-dl'):
    '''
      Scraps the arXiv's html to get data from each entry in a search. Entries has the following formatting:
      <entry>\n    
      <id>http://arxiv.org/abs/2008.04584v2</id>\n    
      <updated>2021-05-11T12:00:24Z</updated>\n    
      <published>2020-08-11T08:47:06Z</published>\n    
      <title>Bayesian Selective Inference: Non-informative Priors</title>\n    
      <summary>  We discuss Bayesian inference for parameters selected using the data. First,\nwe provide a critical analysis of the existing positions in the literature\nregarding the correct Bayesian approach under selection. Second, we propose two\ntypes of non-informative priors for selection models. These priors may be\nemployed to produce a posterior distribution in the absence of prior\ninformation as well as to provide well-calibrated frequentist inference for the\nselected parameter. We test the proposed priors empirically in several\nscenarios.\n</summary>\n    
      <author>\n      <name>Daniel G. Rasines</name>\n    </author>\n    <author>\n      <name>G. Alastair Young</name>\n    </author>\n    
      <arxiv:comment xmlns:arxiv="http://arxiv.org/schemas/atom">24 pages, 7 figures</arxiv:comment>\n    
      <link href="http://arxiv.org/abs/2008.04584v2" rel="alternate" type="text/html"/>\n    
      <link title="pdf" href="http://arxiv.org/pdf/2008.04584v2" rel="related" type="application/pdf"/>\n    
      <arxiv:primary_category xmlns:arxiv="http://arxiv.org/schemas/atom" term="math.ST" scheme="http://arxiv.org/schemas/atom"/>\n    
      <category term="math.ST" scheme="http://arxiv.org/schemas/atom"/>\n    
      <category term="stat.TH" scheme="http://arxiv.org/schemas/atom"/>\n  
      </entry>\n  
    '''

    # Remove space in seach query
    search_query=search_query.strip().replace(" ", "+")
    # Call arXiv API
    arXiv_url=f'http://export.arxiv.org/api/query?search_query={search_by}:{search_query}&sortBy={sort_by}&start=0&max_results={max_results}'
    with urllib.request.urlopen(arXiv_url) as url:
        s = url.read()
    
    # Parse the xml data
    root = html.fromstring(s)
    # Fetch relevant pdf information
    pdf_entries = root.xpath("entry")

    pdf_titles   = []
    pdf_authors  = []
    pdf_urls     = []
    pdf_categories = []
    folder_names = []
    pdf_citation = []
    pdf_years = []

    for i, pdf in enumerate(pdf_entries):
      # print(pdf.xpath('updated/text()')[0][:4])
      # xpath return a list with every ocurrence of the html path. Since we're getting each entry individually, we'll take the first element to avoid an unecessary list
      pdf_titles.append(re.sub('[^a-zA-Z0-9]', ' ', pdf.xpath("title/text()")[0]))
      pdf_authors.append(pdf.xpath("author/name/text()"))
      pdf_urls.append(pdf.xpath("link[@title='pdf']/@href")[0])
      pdf_categories.append(pdf.xpath("category/@term"))
      folder_names.append(folder_name)
      pdf_years.append(pdf.xpath('updated/text()')[0][:4])
      pdf_citation.append(f"{', '.join(pdf_authors[i])}, {pdf_titles[i]}. arXiv [{pdf_categories[i][0]}] ({pdf_years[i]}), (available at {pdf_urls[i]}).")

      

    pdf_info=list(zip(pdf_titles, pdf_urls, pdf_authors, pdf_categories, folder_names, pdf_citation))
    
    # Check number of available files
    # print('Requesting {max_results} files'.format(max_results=max_results))
    if len(pdf_urls)<int(max_results):
        matching_pdf_num=len(pdf_urls)
        # print('Only {matching_pdf_num} files available'.format(matching_pdf_num=matching_pdf_num))
    return pdf_info, pdf_citation


def download_pdf(pdf_info):
   
    # if len(os.listdir(f'./{folder_name}') ) != 0:
            # check folder is empty to avoid using papers from old runs:
            # os.remove(f'./{folder_name}/*')
    all_reference_text = []
    for i,p in enumerate(stqdm(pdf_info, desc='Searching and downloading papers')):

        pdf_title=p[0]
        pdf_url=p[1]
        pdf_author=p[2]
        pdf_category=p[3]
        folder_name=p[4]
        pdf_citation=p[5]
        r = requests.get(pdf_url, allow_redirects=True)
        if  i == 0:
            if not os.path.exists(f'{folder_name}'):
                os.makedirs(f"{folder_name}")
            else:
                shutil.rmtree(f'{folder_name}') 
                os.makedirs(f"{folder_name}")
        with open(f'{folder_name}/{pdf_title}.pdf', 'wb') as currP:
            currP.write(r.content)
        if i == 0:
            st.markdown("###### Papers found:")
        st.markdown(f"{i+1}. {pdf_citation}")
        time.sleep(0.15)
        all_reference_text.append(f"{i+1}. {pdf_citation}\n")
    if 'all_reference_text' not in st.session_state:
        st.session_state.key = 'all_reference_text'
    st.session_state['all_reference_text'] = ' '.join(all_reference_text)

    # print(all_reference_text)
    


max_results_current = 5
max_results = max_results_current
# pdf_info = ''
# pdf_citation = ''
def search_click_callback(search_query, max_results):
    global pdf_info, pdf_citation
    pdf_info, pdf_citation = call_arXiv_API(f'{search_query}', max_results=max_results)
    download_pdf(pdf_info)
    return pdf_info




with st.form(key='columns_in_form', clear_on_submit = False):
    c1, c2 = st.columns([8,1])
    with c1:
        search_query = st.text_input("Input search query here:", placeholder='Keywords for most relevant search...', value=''
                                       )#search_query, max_results_current))

    with c2:
        max_results = st.text_input("Max papers", value=max_results_current)
        max_results_current = max_results_current
    searchButton = st.form_submit_button(label = 'Search')
    #         search_click(search_query, max_results_default)

if searchButton:
    global pdf_info
    pdf_info =  search_click_callback(search_query, max_results)
    if 'pdf_info' not in st.session_state:
        st.session_state.key = 'pdf_info'
    st.session_state['pdf_info'] = pdf_info
    # print(f'This is PDF info from search:{pdf_info}')


# def tokenize_callback():

    # return docs

# tokenization_form = st.form(key='tokenization-form')
# tokenization_form.markdown(f"Happy with your paper search results? ")
# toknizeButton = tokenization_form.form_submit_button(label = "Yes! Let's tokenize.", on_click=tokenize_callback())
# tokenization_form.markdown("If not, change keywords and search again. [This step costs!](https://openai.com/api/pricing/)")



# submitButton = form.form_submit_button('Submit')
# with st.form(key='tokenization_form', clear_on_submit = False):
#     st.markdown(f"Happy with your paper search results? If not, change keywords and search again. [This step costs!](https://openai.com/api/pricing/)")
#     # st.text_input("Input search query here:", placeholder='Keywords for most relevant search...'
#     #                                )#search_query, max_results_current))
#     toknizeButton = st.form_submit_button(label = "Yes! Let's tokenize.")

# if toknizeButton:
#     tokenize_callback()

# tokenize_callback()




def answer_callback(question_query):
    import paperqa
    global docs
    # global pdf_info
    progress_text = "Please wait..."
    # my_bar = st.progress(0, text = progress_text)
    st.info('Please wait...', icon="🔥")
    if docs is None:
        # my_bar.progress(0.2, "Please wait...")
        pdf_info = st.session_state['pdf_info']
        # print('buliding docs')
        docs = paperqa.Docs()
        pdf_paths = [f"{p[4]}/{p[0]}.pdf" for p in pdf_info]
        pdf_citations = [p[5] for p in pdf_info]
        print(list(zip(pdf_paths, pdf_citations)))

        for d, c in zip(pdf_paths, pdf_citations):
            # print(d,c)
            docs.add(d, c)
    # docs._build_faiss_index()
    answer = docs.query(question_query)
    # print(answer.formatted_answer)
    # my_bar.progress(1.0, "Done!")
    st.success('Done!')
    return answer.formatted_answer



form = st.form(key='question_form')
question_query = form.text_input("What do you wanna know from these papers?", placeholder='Input questions here...',
                value='')
submitButton = form.form_submit_button('Submit')

if submitButton:
    with st.expander("Found papers:", expanded=True):
        st.write(f"{st.session_state['all_reference_text']}")
    st.text_area("Answer:", answer_callback(question_query), height=600)

# with st.form(key='question_form', clear_on_submit = False):
#     question_query = st.text_input("What do you wanna know from these papers?", placeholder='Input questions here')
#     # st.text_input("Input search query here:", placeholder='Keywords for most relevant search...'
#     #                                )#search_query, max_results_current))
#     submitButton = form.form_submit_button(label = "Submit", on_click=answer_callback(question_query))


# Simulation-based inference bayesian model selection





# test = "<ul> \
#         <li>List item here</li> \
#         <li>List item here</li> \
#         <li>List item here</li> \
#         <li>List item here</li> \
#         </ul>"
# test = "'''It was the best of times, it was the worst of times, it was \
#             the age of wisdom, it was the age of foolishness, it was \
#             the epoch of belief, it was the epoch of incredulity, it \
#             was the season of Light, it was the season of Darkness, it\
#             was the spring of hope, it was the winter of despair, (...)'''"

# citation_text = st.text_area('Papers found:',test, height=300) # f'{pdf_citation}'


# for i, cite in enumerate(pdf_citation):
#     st.markdown(f'{i+1}. {cite}')
#     time.sleep(1)


# def make_clickable('link',text):
#     return f'<a target="_blank" href="{link}">{text}'