"""Document token classification demo using LayoutLMv3.

LayoutLMv2 depended on Detectron2, which breaks on modern Pillow/Python.
LayoutLMv3 keeps the same FUNSD form-understanding demo without Detectron2.
"""

from pathlib import Path

import gradio as gr
import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont
from transformers import LayoutLMv3ForTokenClassification, LayoutLMv3Processor

DATA_DIR = Path(__file__).parent
EXAMPLE_IMAGE = DATA_DIR / "document.png"
SOURCE_INVOICE = DATA_DIR / "invoice.jpeg"

# Apply OCR and token classification once at startup.
processor = LayoutLMv3Processor.from_pretrained(
    "microsoft/layoutlmv3-base", apply_ocr=True
)
model = LayoutLMv3ForTokenClassification.from_pretrained(
    "nielsr/layoutlmv3-finetuned-funsd"
)
model.eval()

# Official FUNSD IOB labels used by the fine-tuned checkpoint.
labels = [
    "O",
    "B-HEADER",
    "I-HEADER",
    "B-QUESTION",
    "I-QUESTION",
    "B-ANSWER",
    "I-ANSWER",
]
id2label = {index: label for index, label in enumerate(labels)}
label2color = {
    "question": "blue",
    "answer": "green",
    "header": "orange",
    "other": "violet",
}

# Keep a stable example image for the Gradio UI.
if SOURCE_INVOICE.exists():
    Image.open(SOURCE_INVOICE).convert("RGB").save(EXAMPLE_IMAGE)


def unnormalize_box(bbox, width, height):
    """Convert LayoutLM 0-1000 boxes back to pixel coordinates."""
    return [
        width * (bbox[0] / 1000),
        height * (bbox[1] / 1000),
        width * (bbox[2] / 1000),
        height * (bbox[3] / 1000),
    ]


def iob_to_label(label):
    """Strip IOB prefixes such as B- and I-."""
    label = label[2:]
    return label.lower() if label else "other"


def process_image(image):
    """Run OCR + LayoutLMv3 and draw predicted entity boxes."""
    if image is None:
        raise gr.Error("Upload a document image first.")

    image = image.convert("RGB")
    width, height = image.size

    encoding = processor(
        image,
        truncation=True,
        return_offsets_mapping=True,
        return_tensors="pt",
    )
    offset_mapping = encoding.pop("offset_mapping")

    with torch.no_grad():
        outputs = model(**encoding)

    predictions = outputs.logits.argmax(-1).squeeze().tolist()
    token_boxes = encoding.bbox.squeeze().tolist()
    is_subword = np.array(offset_mapping.squeeze().tolist())[:, 0] != 0

    true_predictions = [
        id2label[prediction]
        for index, prediction in enumerate(predictions)
        if not is_subword[index]
    ]
    true_boxes = [
        unnormalize_box(box, width, height)
        for index, box in enumerate(token_boxes)
        if not is_subword[index]
    ]

    draw = ImageDraw.Draw(image)
    font = ImageFont.load_default()
    for prediction, box in zip(true_predictions, true_boxes):
        predicted_label = iob_to_label(prediction)
        color = label2color.get(predicted_label, "violet")
        draw.rectangle(box, outline=color)
        draw.text((box[0] + 10, box[1] - 10), predicted_label, fill=color, font=font)

    return image


title = "Interactive demo: LayoutLMv3"
description = (
    "Document understanding demo inspired by Niels Rogge's Spaces. "
    "This Space uses LayoutLMv3 fine-tuned on FUNSD so it runs without Detectron2."
)
article = (
    "<p style='text-align: center'>"
    "<a href='https://arxiv.org/abs/2204.08387' target='_blank'>"
    "LayoutLMv3: Pre-training for Document AI with Unified Text and Image Masking"
    "</a> | "
    "<a href='https://github.com/microsoft/unilm' target='_blank'>Github Repo</a>"
    "</p>"
)

demo = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil", label="Document image"),
    outputs=gr.Image(type="pil", label="Annotated image"),
    title=title,
    description=description,
    article=article,
    examples=[[str(EXAMPLE_IMAGE)]],
    css=".output-image, .input-image {height: 40rem !important; width: 100% !important;}",
)

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