import gradio as gr
import numpy as np
from PIL import Image
import torch
import pandas as pd
from transformers import AutoImageProcessor, AutoModelForObjectDetection, AutoProcessor, Pix2StructForConditionalGeneration
import torch
from io import StringIO

device="cpu"

MAX_PATCHES = 1024
MAX_NEW_TOKENS = 1024
TABLE_THRESHOLD = 0.9
TABLE_PADDING = 5

# Detection related
table_detr_processor = AutoImageProcessor.from_pretrained("microsoft/table-transformer-detection")
table_detr_model = AutoModelForObjectDetection.from_pretrained("microsoft/table-transformer-detection", revision="no_timm")
table_detr_model.to(device)
table_detr_model.eval()

no_table_found = Image.open("app_assets/no_table_found.png")

# Recognition related
table_recog_processor = AutoProcessor.from_pretrained("KennethTM/pix2struct-base-table2html")
table_recog_model = Pix2StructForConditionalGeneration.from_pretrained("KennethTM/pix2struct-base-table2html")
table_recog_model.to(device)
table_recog_model.eval()

def table_detection(image, threshold=TABLE_THRESHOLD):

    inputs = table_detr_processor(images=image, return_tensors="pt")
    inputs = {k: v.to(device) for k, v in inputs.items()}

    with torch.inference_mode():

        outputs = table_detr_model(**inputs)

    target_sizes = torch.tensor([image.size[::-1]])
    results = table_detr_processor.post_process_object_detection(outputs, threshold=threshold, target_sizes=target_sizes)
    table_boxes = [i for i in results[0]["boxes"]]

    tables = []
    if len(table_boxes) == 0:
        tables.append(no_table_found)
    else:
        padding = TABLE_PADDING
        for box in table_boxes:
            box = [int(i) for i in box]
            box[0] = max(0, box[0]-padding)
            box[1] = max(0, box[1]-padding)
            box[2] = min(image.width, box[2]+padding)
            box[3] = min(image.height, box[3]+padding)
            tables.append(image.crop(box))

    return tables

def table_recognition(image, max_new_tokens = MAX_NEW_TOKENS):

    encoding = table_recog_processor(image, return_tensors="pt", max_patches=MAX_PATCHES)

    with torch.inference_mode():
        flattened_patches = encoding.pop("flattened_patches").to(device)
        attention_mask = encoding.pop("attention_mask").to(device)
        predictions = table_recog_model.generate(flattened_patches=flattened_patches, attention_mask=attention_mask, max_new_tokens=max_new_tokens)

    predictions_decoded = table_recog_processor.tokenizer.batch_decode(predictions, skip_special_tokens=True)
    table_html = predictions_decoded[0]
    
    return table_html

def table_recognition_outputs(image):
    # Table to HTML
    table_html = table_recognition(image)

    # Write HTML to files
    with open("table.html", "w") as file:
        file.write(table_html)

    df = pd.read_html(StringIO(table_html))[0]
    df.to_csv("table.csv", index=False)

    return [table_html,
            gr.DownloadButton("Download HTML", value="table.html", visible=True),
            gr.DownloadButton("Download CSV", value="table.csv", visible=True)]

demo_detection = [
    "app_assets/example_one_table.jpg",
    "app_assets/example_two_tables.jpg",
]

demo_recognition = [
    "app_assets/example_recog_1.jpg",
    "app_assets/example_recog_2.jpg",
]

with gr.Blocks() as demo:

    with gr.Tab("Recognition"):
        gr.Markdown("# Table recognition")
        gr.Markdown("This model ([KennethTM/pix2struct-base-table2html](https://huggingface.co/KennethTM/pix2struct-base-table2html)) converts an image of a table to HTML format and is finetuned from [Pix2Struct base model](https://huggingface.co/google/pix2struct-base).")
        gr.Markdown("The model expects an image containing only a table. If the table is embedded in a document, first use the detection model in the 'Detection' tab.")
        gr.Markdown("*Note that recognition model inference is slow on CPU (a few minutes), please be patient*")
        with gr.Row():
            with gr.Column():
                input_table = gr.Image(type="pil", label="Table", show_label=True, scale=1)
                
            with gr.Column():
                output_html = gr.HTML(label="Table (HTML format)", show_label=False)
                
                with gr.Row():
                    download_html = gr.DownloadButton(visible=False)
                    download_csv = gr.DownloadButton(visible=False)

        with gr.Row():
            examples = gr.Examples(demo_recognition, input_table, cache_examples=False, label="Example tables (MMTab dataset)")

        input_table.change(fn=table_recognition_outputs, inputs=input_table, outputs=[output_html, download_html, download_csv])

    with gr.Tab("Detection"):
        gr.Markdown("# Table detection")
        gr.Markdown("This model detect tables in a document image with [Microsoft's Table Transformer model](https://huggingface.co/microsoft/table-transformer-detection).")
        gr.Markdown("Use the detection to find tables, download the results and use as input for table recognition in the 'Recognition' tab.")
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(type="pil", label="Document", show_label=True, scale=1)
            
            with gr.Column():
                output_gallery = gr.Gallery(type="pil", label="Tables", show_label=True, scale=1, format="png")
        
        with gr.Row():
            examples = gr.Examples(demo_detection, input_image, cache_examples=False, label="Example documents (PubTabNet dataset)")

        input_image.change(fn=table_detection, inputs=input_image, outputs=output_gallery)

demo.launch()