import gradio as gr

import kornia as K
from kornia.core import Tensor


def load_img(file):
    # load the image using the rust backend
    img_bgr: Tensor = K.io.load_image(file.name, K.io.ImageLoadType.RGB32)
    img_bgr = img_bgr[None, ...].float() / 255.0

    img_rgb: Tensor = K.color.bgr_to_rgb(img_bgr)
    img_gray = K.color.rgb_to_grayscale(img_rgb)

    return img_gray


def canny_edge_detector(file):
    x_gray = load_img(file)
    x_canny: Tensor = K.filters.canny(x_gray)[0]
    img_out = 1.0 - x_canny.clamp(0.0, 1.0)
    return K.utils.tensor_to_image(img_out)


def sobel_edge_detector(file):
    x_gray = load_img(file)
    x_sobel: Tensor = K.filters.sobel(x_gray)
    img_out = 1.0 - x_sobel
    return K.utils.tensor_to_image(img_out)


def simple_edge_detector(file, order=1, dir="x"):
    x_gray = load_img(file)
    grads: Tensor = K.filters.spatial_gradient(x_gray, order=2)  # BxCx2xHxW
    grads_x = grads[:, :, 0]
    grads_y = grads[:, :, 1]
    if dir == "x":
        img_out = 1.0 - grads_x.clamp(0.0, 1.0)
    else:
        img_out = 1.0 - grads_y.clamp(0.0, 1.0)
    return K.utils.tensor_to_image(img_out)


def laplacian_edge_detector(file, kernel_size=5):
    x_gray = load_img(file)
    x_canny: Tensor = K.filters.canny(x_gray)[0]
    img_out = 1.0 - x_canny.clamp(0.0, 1.0)
    return K.utils.tensor_to_image(img_out)


examples = [
    ["examples/doraemon.png"],
]

title = "Kornia Edge Detector"
description = "<p style='text-align: center'>This is a Gradio demo for Kornia's Edge Detector.</p><p style='text-align: center'>To use it, simply upload your image, or click one of the examples to load them, and use the sliders to enhance! Read more at the links at the bottom.</p>"
article = "<p style='text-align: center'><a href='https://kornia.readthedocs.io/en/latest/' target='_blank'>Kornia Docs</a> | <a href='https://github.com/kornia/kornia' target='_blank'>Kornia Github Repo</a> | <a href='https://kornia-tutorials.readthedocs.io/en/latest/image_enhancement.html' target='_blank'>Kornia Enhancements Tutorial</a></p>"


def change_kernel(choice):
    if choice == "Laplacian":
        return gr.update(visible=True)
    else:
        return gr.update(visible=False)


def change_order(choice):
    if choice == "simple":
        return gr.update(visible=True)
    else:
        return gr.update(visible=False)


def change_button(choice):
    if choice == "simple":
        return gr.update("Detect").click(
            canny_edge_detector, inputs=image_input, outputs=image_output
        )
    elif choice == "sobel":
        return gr.update("Detect").click(
            sobel_edge_detector, inputs=image_input, outputs=image_output
        )
    elif choice == "laplacian":
        return gr.update("Detect").click(
            laplacian_edge_detector, inputs=image_input, outputs=image_output
        )
    else:
        return gr.update("Detect").click(
            simple_edge_detector, inputs=image_input, outputs=image_output
        )


with gr.Blocks() as demo:
    with gr.Row():
        image_input = gr.Image()
        image_output = gr.Image()
        radio = gr.Radio(
            ["canny", "simple", "sobel", "laplacian"],
            label="Type of Edge Detector",
        )
    with gr.Row():
        kernel = gr.Slider(
            minimum=1,
            maximum=6,
            step=1,
            default=5,
            label="kernel_size",
            visible=False,
        )
        order = gr.Slider(
            minimum=1,
            maximum=2,
            step=1,
            default=1,
            label="Derivative order",
            visible=False,
        )
        Button = gr.Button("Detect")

        radio.change(fn=change_kernel, inputs=radio, outputs=kernel)
        radio.change(fn=change_order, inputs=radio, outputs=order)
        radio.change(fn=change_button, inputs=radio, outputs=Button)
    gr.Examples(examples, inputs=[image_input])


demo.launch()