# -*- coding: utf-8 -*-

import json
import os
from pathlib import Path

import gradio as gr
import numpy as np
from doc_ufcn import models
from doc_ufcn.main import DocUFCN
from PIL import Image, ImageDraw

from config import parse_configurations

# Load the config
config = parse_configurations(Path("config.json"))

# Download the model
model_path, parameters = models.download_model(name=config["model_name"])

# Store classes_colors list
classes_colors = config["classes_colors"]

# Store classes
classes = parameters["classes"]

# Check that the number of colors is equal to the number of classes -1
assert len(classes) - 1 == len(
    classes_colors
), f"The parameter classes_colors was filled with the wrong number of colors. {len(classes)-1} colors are expected instead of {len(classes_colors)}."

# Check that the paths of the examples are valid
for example in config["examples"]:
    assert os.path.exists(example), f"The path of the image '{example}' does not exist."

# Load the model
model = DocUFCN(
    no_of_classes=len(classes),
    model_input_size=parameters["input_size"],
    device="cpu",
)
model.load(model_path=model_path, mean=parameters["mean"], std=parameters["std"])


def query_image(image):
    """
    Draws the predicted polygons with the color provided by the model on an image

    :param image: An image to predict
    :return: Image and dict, an image with the predictions and a
        dictionary mapping an object idx (starting from 1) to a dictionary describing the detected object:
        - `polygon` key : list, the coordinates of the points of the polygon,
        - `confidence` key : float, confidence of the model,
        - `channel` key : str, the name of the predicted class.
    """

    # Make a prediction with the model
    detected_polygons, probabilities, mask, overlap = model.predict(
        input_image=image, raw_output=True, mask_output=True, overlap_output=True
    )

    # Load image
    image = Image.fromarray(image)

    # Make a copy of the image to keep the source and also to be able to use Pillow's blend method
    img2 = image.copy()

    # Initialize the dictionary which will display the json on the application
    predict = []

    # Create the polygons on the copy of the image for each class with the corresponding color
    # We do not draw polygons of the background channel (channel 0)
    for channel in range(1, len(classes)):
        for i, polygon in enumerate(detected_polygons[channel]):
            # Draw the polygons on the image copy.
            # Loop through the class_colors list (channel 1 has color 0)
            ImageDraw.Draw(img2).polygon(
                polygon["polygon"], fill=classes_colors[channel - 1]
            )

            # Build the dictionary
            # Add an index to dictionary keys to differentiate predictions of the same class
            predict.append(
                {
                    "polygon": np.asarray(polygon["polygon"])
                    .astype(int)
                    .tolist(),  # The list of coordinates of the points of the polygon
                    "confidence": polygon[
                        "confidence"
                    ],  # Confidence that the model predicts the polygon in the right place
                    "channel": classes[
                        channel
                    ],  # The channel on which the polygon is predicted
                }
            )

    # Return the blend of the images and the dictionary formatted in json
    return Image.blend(image, img2, 0.5), json.dumps(predict, indent=20)


with gr.Blocks() as process_image:

    # Create app title
    gr.Markdown(f"<h1 align='center'>{config['title']}</h1>")

    # Create app description
    gr.Markdown(config["description"])

    # Create a first row of blocks
    with gr.Row():

        # Create a column on the left
        with gr.Column():

            # Generates an image that can be uploaded by a user
            image = gr.Image()

            # Create a row under the image
            with gr.Row():

                # Generate a button to clear the inputs and outputs
                clear_button = gr.Button("Clear", variant="secondary")

                # Generates a button to submit the prediction
                submit_button = gr.Button("Submit", variant="primary")

            # Create a row under the buttons
            with gr.Row():

                # Generate example images that can be used as input image
                examples = gr.Examples(inputs=image, examples=config["examples"])

        # Create a column on the right
        with gr.Column():

            # Generates an output image that does not support upload
            image_output = gr.Image(interactive=False)

            # Create a row under the predicted image
            with gr.Row():

                # Create a column so that the JSON output doesn't take the full size of the page
                with gr.Column():

                    # Create a collapsible region
                    with gr.Accordion("JSON"):

                        # Generates a json with the model predictions
                        json_output = gr.JSON()

    # Create the button to clear the inputs and outputs
    clear_button.click(
        lambda x, y, z: (None, None, None),
        inputs=[image, image_output, json_output],
        outputs=[image, image_output, json_output],
    )

    # Create the button to submit the prediction
    submit_button.click(query_image, inputs=image, outputs=[image_output, json_output])

# Launch the application
process_image.launch()