import torch
from PIL import ImageDraw, ImageFont, Image
from transformers import AutoModelForTokenClassification, AutoProcessor
import fitz  # PyMuPDF
import io
import os

MODEL_KEY = os.getenv("MODEL_KEY")


def extract_data_from_pdf(pdf_path, page_number=0):
    """
    Extracts image, words, and bounding boxes from a specified page of a PDF.

    Args:
    - pdf_path (str): Path to the PDF file.
    - page_number (int): Page number to extract data from (0-indexed).

    Returns:
    - image: An image of the specified page.
    - words: A list of words found on the page.
    - boxes: A list of bounding boxes corresponding to the words.
    """
    # Open the PDF
    doc = fitz.open(pdf_path)
    page = doc.load_page(page_number)

    # Extract image of the page
    pix = page.get_pixmap()
    image_bytes = pix.tobytes("png")
    image = Image.open(io.BytesIO(image_bytes))

    # Extract words and their bounding boxes
    words = []
    boxes = []
    for word in page.get_text("words"):
        words.append(word[4])
        boxes.append(word[:4])  # (x0, y0, x1, y1)

    doc.close()
    return image, words, boxes


def merge_pairs_v2(pairs):
    if not pairs:
        return []

    merged = [pairs[0]]
    for current in pairs[1:]:
        last = merged[-1]
        if last[0] == current[0]:
            # Merge 'y' values (as strings) if 'x' values are the same
            merged[-1] = [last[0], last[1] + " " + current[1]]
        else:
            merged.append(current)

    return merged


def create_pretty_table(data):
    table = "<div>"
    for row in data:
        color = (
            "blue"
            if row[0] == "Heder"
            else "green"
            if row[0] == "Section"
            else "red"
            if row[0] == "Number"
            else "black"
        )
        table += "<p style='color:{};'>---{}---</p>{}".format(
            color, row[0], row[1]
        )
    table += "</div>"
    return table


# When using this function in Gradio, set the output type to 'html'


def interference(example, page_number=0):
    try:
        image, words, boxes = extract_data_from_pdf(example, page_number)
        boxes = [list(map(int, box)) for box in boxes]

        # Process the image and words
        model = AutoModelForTokenClassification.from_pretrained(
            "karida/LayoutLMv3_RFP", use_auth_token=MODEL_KEY
        )
        processor = AutoProcessor.from_pretrained(
            "microsoft/layoutlmv3-base", apply_ocr=False
        )
        encoding = processor(image, words, boxes=boxes, return_tensors="pt")

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

        logits = outputs.logits
        predictions = logits.argmax(-1).squeeze().tolist()
        model_words = encoding.word_ids()

        # Process predictions
        token_boxes = encoding.bbox.squeeze().tolist()
        width, height = image.size

        true_predictions = [model.config.id2label[pred] for pred in predictions]
        true_boxes = token_boxes
        # Draw annotations on the image
        draw = ImageDraw.Draw(image)
        font = ImageFont.load_default()

        def iob_to_label(label):
            label = label[2:]
            return "other" if not label else label.lower()

        label2color = {
            "question": "blue",
            "answer": "green",
            "header": "red",
            "other": "violet",
        }

        # print(len(true_predictions), len(true_boxes), len(model_words))

        table = []
        ids = set()

        for prediction, box, model_word in zip(
            true_predictions, true_boxes, model_words
        ):
            predicted_label = iob_to_label(prediction)
            draw.rectangle(box, outline=label2color[predicted_label], width=2)
            # draw.text((box[0] + 10, box[1] - 10), text=predicted_label, fill=label2color[predicted_label], font=font)
            if (
                model_word
                and model_word not in ids
                and predicted_label != "other"
            ):
                ids.add(model_word)
                table.append([predicted_label[0], words[model_word]])

        values = merge_pairs_v2(table)
        # values = [
        #     ["Heder", x[1]] if x[0] == "q" else ["Section", x[1]]
        #     for x in values
        # ]
        new_values = []
        for x in values:
            if x[0] == "q":
                new_values.append(["Heder", x[1]])
            elif x[0] == "a":
                new_values.append(["Section", x[1]])
            elif x[0] == "h":
                new_values.append(["Number", x[1]])

        table = create_pretty_table(new_values)
        print(table)
        return image, table

    except IndexError as e:
        image = Image.open("error_image.png")
        # Return a custom HTML-styled error message if an IndexError occurs
        gr.Error(
            "Error: in the current version of the model, the maximum number of words per page is 512."
        )
        return (
            image,
            f"<div style='color: red; font-weight: bold;'>IndexError: {str(e)}</div>",
        )

    except Exception as e:
        image = Image.open("error_image.png")
        # Handle other exceptions
        gr.Error(f"An error occurred: {e}")
        return (
            image,
            f"<div style='color: red; font-weight: bold;'>An error occurred: {str(e)}</div>",
        )


import gradio as gr

description_text = """
<p>
    Information Retrieval for Request For Proposal documents with use of the LayoutLM for Token Classification.<br>
    Classified tokens:<br>
    -    Number - <span style='color: red;'>shown in red</span><br>
    -    Header - <span style='color: blue;'>shown in blue</span><br>
    -    Section - <span style='color: green;'>shown in green</span><br>
    -    other (ignored) - <span style='color: violet;'>shown in violet</span>
</p>
"""

flagging_options = ["great example", "bad example"]


iface = gr.Interface(
    fn=interference,
    inputs=["file", "number"],
    outputs=["image", "html"],
    # examples=[["output.pdf", 1]],
    description=description_text,
    flagging_options=flagging_options,
)


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