import os

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

import functools
from PIL import Image, ImageDraw
import gradio as gr

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


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


CHECKPOINTS = {
    "LayoutLMv1 🦉": "impira/layoutlm-document-qa",
    "Donut 🍩": "naver-clova-ix/donut-base-finetuned-docvqa",
}

PIPELINES = {}


def construct_pipeline(model):
    global PIPELINES
    if model in PIPELINES:
        return PIPELINES[model]

    device = "cuda" if torch.cuda.is_available() else "cpu"
    ret = get_pipeline(checkpoint=CHECKPOINTS[model], device=device)
    PIPELINES[model] = ret
    return ret


@functools.lru_cache(1024)
def run_pipeline(model, question, document, top_k):
    pipeline = construct_pipeline(model)
    return pipeline(question=question, **document.context, top_k=top_k)


# TODO: Move into docquery
# TODO: Support words past the first page (or window?)
def lift_word_boxes(document):
    return document.context["image"][0][1]


def expand_bbox(word_boxes):
    if len(word_boxes) == 0:
        return None

    min_x, min_y, max_x, max_y = zip(*[x[1] for x in word_boxes])
    min_x, min_y, max_x, max_y = [min(min_x), min(min_y), max(max_x), max(max_y)]
    return [min_x, min_y, max_x, max_y]


# LayoutLM boxes are normalized to 0, 1000
def normalize_bbox(box, width, height, padding=0.005):
    min_x, min_y, max_x, max_y = [c / 1000 for c in box]
    if padding != 0:
        min_x = max(0, min_x - padding)
        min_y = max(0, min_y - padding)
        max_x = min(max_x + padding, 1)
        max_y = min(max_y + padding, 1)
    return [min_x * width, min_y * height, max_x * width, max_y * height]


examples = [
    [
        "invoice.png",
        "What is the invoice number?",
    ],
    [
        "contract.jpeg",
        "What is the purchase amount?",
    ],
    [
        "statement.png",
        "What are net sales for 2020?",
    ],
]


def process_path(path):
    if path:
        try:
            document = load_document(path)
            return document, document.preview, None
        except Exception:
            pass
    return None, None, None


def process_upload(file):
    if file:
        return process_path(file.name)
    else:
        return None, None, None


colors = ["#64A087", "green", "black"]


def process_question(question, document, model=list(CHECKPOINTS.keys())[0]):
    if document is None:
        return None, None

    predictions = run_pipeline(model, question, document, 3)
    image = document.preview.copy()
    draw = ImageDraw.Draw(image, "RGBA")
    for i, p in enumerate(ensure_list(predictions)):
        if i > 0:
            # Keep the code around to produce multiple boxes, but only show the top
            # prediction for now
            break

        if "start" in p and "end" in p and p.get("page") == 0:
            x1, y1, x2, y2 = normalize_bbox(
                expand_bbox(lift_word_boxes(document)[p["start"] : p["end"] + 1]),
                image.width,
                image.height,
            )
            draw.rectangle(((x1, y1), (x2, y2)), fill=(0, 255, 0, int(0.4 * 255)))

    return image, predictions


def load_example_document(img, question, model):
    document = ImageDocument(Image.fromarray(img))
    preview, answer = process_question(question, document, model)
    return document, question, preview, answer


CSS = """
#short-upload-box .w-full {
    min-height: 10rem !important;
}
#question input {
    font-size: 16px;
}
"""

with gr.Blocks(css=CSS) as demo:
    gr.Markdown("# DocQuery: Query Documents w/ NLP")
    document = gr.Variable()
    example_question = gr.Textbox(visible=False)
    example_image = gr.Image(visible=False)

    gr.Markdown("## 1. Upload a file or select an example")
    with gr.Row(equal_height=True):
        with gr.Column():
            upload = gr.File(
                label="Upload a file", interactive=True, elem_id="short-upload-box"
            )
            url = gr.Textbox(label="... or a URL", interactive=True)
        gr.Examples(
            examples=examples,
            inputs=[example_image, example_question],
        )

    gr.Markdown("## 2. Ask a question")

    with gr.Row(equal_height=True):
        question = gr.Textbox(
            label="Question",
            placeholder="e.g. What is the invoice number?",
            lines=1,
            max_lines=1,
            elem_id="question",
        )
        model = gr.Radio(
            choices=list(CHECKPOINTS.keys()),
            value=list(CHECKPOINTS.keys())[0],
            label="Model",
        )

    with gr.Row():
        clear_button = gr.Button("Clear", variant="secondary")
        submit_button = gr.Button("Submit", variant="primary", elem_id="submit-button")

    with gr.Row():
        image = gr.Image(visible=True)
        with gr.Column():
            output = gr.JSON(label="Output")

    clear_button.click(
        lambda _: (None, None, None, None),
        inputs=clear_button,
        outputs=[image, document, question, output],
    )
    upload.change(fn=process_upload, inputs=[upload], outputs=[document, image, output])
    url.change(fn=process_path, inputs=[url], outputs=[document, image, output])

    question.submit(
        fn=process_question,
        inputs=[question, document, model],
        outputs=[image, output],
    )

    submit_button.click(
        process_question,
        inputs=[question, document, model],
        outputs=[image, output],
    )

    model.change(
        process_question, inputs=[question, document, model], outputs=[image, output]
    )

    example_image.change(
        fn=load_example_document,
        inputs=[example_image, example_question, model],
        outputs=[document, question, image, output],
    )

    gr.Markdown("### More Info")
    gr.Markdown(
        "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."
    )
    gr.Markdown("[Github Repo](https://github.com/impira/docquery)")

if __name__ == "__main__":
    demo.launch()