import json
import streamlit as st
import requests as req

# TODO: improve layout (columns, sidebar, forms)
# st.set_page_config(layout='wide')


st.title('Question Answering example')


##########################################################
st.subheader('1. A simple question (extractive, closed domain)')
##########################################################


WIKI_URL = 'https://en.wikipedia.org/w/api.php'
WIKI_QUERY = "?format=json&action=query&prop=extracts&explaintext=1"
WIKI_BERT = "&titles=BERT_(language_model)"
WIKI_METHOD = 'GET'

response = req.request(WIKI_METHOD, f'{WIKI_URL}{WIKI_QUERY}{WIKI_BERT}')
resp_json = json.loads(response.content.decode("utf-8"))
wiki_bert = resp_json['query']['pages']['62026514']['extract']
paragraph =  wiki_bert

written_passage = st.text_area(
    'Paragraph used for QA (you can also edit, or copy/paste new content)', 
    paragraph, 
    height=250
)
if written_passage:
    paragraph = written_passage

# question = 'How many languages does bert understand?'
question = 'How many attention heads does Bert have?'

written_question = st.text_input(
    'Question used for QA (you can also edit, and experiment with the answers)', 
    question
)
if written_question:
    question = written_question

QA_URL = "https://api-inference.huggingface.co/models/deepset/roberta-base-squad2"
QA_METHOD = 'POST'


if st.button('Run QA inference (get answer prediction)'):
    if paragraph and question:
        inputs = {'question': question, 'context': paragraph}
        payload = json.dumps(inputs)
        prediction = req.request(QA_METHOD, QA_URL, data=payload)
        answer = json.loads(prediction.content.decode("utf-8"))
        # >>> answer structure:
        # {
        #   "score": 0.24088488519191742,
        #   "start": 3595,
        #   "end": 3602,
        #   "answer": "over 70"
        # }
        answer_dict = dict(answer)
        # st.write(answer_dict)
        answer_span = answer_dict["answer"]
        answer_score = answer_dict["score"]
        st.write(f'Answer: **{answer_span}**')
        start_par = max(0, answer_dict["start"]-86)
        stop_para = min(answer_dict["end"]+90, len(paragraph))
        answer_context = paragraph[start_par:stop_para].replace(answer_span, f'**{answer_span}**')
        st.write(f'Answer context (and score): ... _{answer_context}_ ... (score: {format(answer_score, ".3f")})')
        st.write(f'Answer JSON: ')
        st.write(answer)
    else:
        st.write('Write some passage of text and a question')
        st.stop()