import os

os.environ["TOKENIZERS_PARALLELISM"] = "false"

import streamlit as st

import torch
from docquery.pipeline import get_pipeline
from docquery.document import load_bytes, load_document


def ensure_list(x):
    if isinstance(x, list):
        return x
    else:
        return [x]


@st.experimental_singleton
def construct_pipeline():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    ret = get_pipeline(device=device)
    return ret


@st.cache
def run_pipeline(question, document):
    return construct_pipeline()(question=question, **document.context)


st.markdown("# DocQuery: Query Documents w/ NLP")

if "document" not in st.session_state:
    st.session_state["document"] = None

input_type = st.radio("Pick an input type", ["Upload", "URL"], horizontal=True)


def load_file_cb():
    if st.session_state.file_input is None:
        return

    file = st.session_state.file_input
    with loading_placeholder:
        with st.spinner("Processing..."):
            document = load_bytes(file, file.name)
            _ = document.context
            st.session_state.document = document


def load_url(url):
    if st.session_state.url_input is None:
        return

    url = st.session_state.url_input
    with loading_placeholder:
        with st.spinner("Downloading..."):
            document = load_document(url)
        with st.spinner("Processing..."):
            _ = document.context
        st.session_state.document = document


if input_type == "Upload":
    file = st.file_uploader(
        "Upload a PDF or Image document", key="file_input", on_change=load_file_cb
    )

elif input_type == "URL":
    # url = st.text_input("URL", "", on_change=load_url_callback, key="url_input")
    url = st.text_input("URL", "", key="url_input", on_change=load_url_cb)

question = st.text_input("QUESTION", "")

document = st.session_state.document
loading_placeholder = st.empty()
if document is not None:
    col1, col2 = st.columns(2)
    col1.image(document.preview, use_column_width=True)

if document is not None and question is not None and len(question) > 0:
    predictions = run_pipeline(question=question, document=document)

    col2.header("Answers")
    for p in ensure_list(predictions):
        col2.subheader(f"{ p['answer'] }: ({round(p['score'] * 100, 1)}%)")


"DocQuery uses LayoutLMv1 fine-tuned on DocVQA, a document visual question answering dataset, as well as SQuAD, which boosts its English-language comprehension. To use it, simply upload an image or PDF, type a question, and click 'submit', or click one of the examples to load them."

"[Github Repo](https://github.com/impira/docquery)"