import numpy as np
import cv2
from PIL import Image
from cv2.ximgproc import guidedFilter
import gradio as gr


def clean_image(
    input_image: Image,
    diameter: float = 5,
    sigma_color: float = 8,
    sigma_space: float = 8,
    radius: float = 4,
    eps: float = 16,
) -> Image:
    img = np.array(input_image).astype(np.float32)
    y = img.copy()

    for _ in range(64):
        y = cv2.bilateralFilter(y, diameter, sigma_color, sigma_space)

    for _ in range(4):
        y = guidedFilter(img, y, radius, eps)

    output_image = Image.fromarray(y.clip(0, 255).astype(np.uint8))
    return output_image


def example(_image):
    pass


def send_to_input(output):
    return output


def ui():
    with gr.Blocks() as app:
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(type="pil", label="Input Image")
                start_btn = gr.Button(value="Start", variant="primary")

                with gr.Accordion("Advanced Config", open=False):
                    # ref: https://github.com/gogodr/AdverseCleanerExtension
                    gr.Markdown("#### Bilateral Filter")
                    diameter_slider = gr.Slider(
                        minimum=1,
                        maximum=30,
                        step=1,
                        value=5,
                        label="Diameter",
                        interactive=True,
                    )
                    sigma_color_slider = gr.Slider(
                        minimum=1,
                        maximum=30,
                        step=1,
                        value=8,
                        label="SigmaColor",
                        interactive=True,
                    )
                    sigma_space_slider = gr.Slider(
                        minimum=1,
                        maximum=30,
                        step=1,
                        value=8,
                        label="SigmaSpace",
                        interactive=True,
                    )

                    gr.Markdown("#### Guided Filter")
                    radius_slider = gr.Slider(
                        minimum=1,
                        maximum=30,
                        step=1,
                        value=4,
                        label="Radius",
                        interactive=True,
                    )
                    eps_slider = gr.Slider(
                        minimum=1,
                        maximum=30,
                        step=1,
                        value=16,
                        label="Accuracy",
                        interactive=True,
                    )

                gr.Examples(
                    examples=[
                        ["./examples/sample1.jpg"],
                        ["./examples/sample2.jpg"],
                        ["./examples/sample3.jpg"],
                    ],
                    inputs=[
                        input_image,
                        diameter_slider,
                        sigma_color_slider,
                        sigma_space_slider,
                        radius_slider,
                        eps_slider,
                    ],
                    outputs=[],
                    fn=example,
                    cache_examples=True,
                )

            with gr.Column():
                output_image = gr.Image(
                    type="pil", label="Output Image", interactive=False
                )
                send_to_input_btn = gr.Button(value="Use as input", variant="secondary")

        gr.Markdown(
            "The lllyasviel's original repo is [here](https://github.com/lllyasviel/AdverseCleaner/tree/main)."
        )

        start_btn.click(fn=clean_image, inputs=[input_image], outputs=[output_image])
        send_to_input_btn.click(
            fn=send_to_input, inputs=[output_image], outputs=[input_image]
        )

    app.launch()


if __name__ == "__main__":
    ui()