import os

import cv2
import gradio as gr
import numpy as np
from PIL import Image
from cv2.ximgproc import guidedFilter
from imgutils.data import load_image
from imgutils.restore import restore_with_nafnet, restore_with_scunet


def clean_adverse(
        input_image: Image.Image,
        diameter: float = 5,
        sigma_color: float = 8,
        sigma_space: float = 8,
        radius: float = 4,
        eps: float = 16,
) -> Image.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 clean(
        image: Image.Image,

        diameter: float = 5,
        sigma_color: float = 8,
        sigma_space: float = 8,
        radius: float = 4,
        eps: float = 16,

        use_scunet_clean: bool = False,
        use_nafnet_clean: bool = False
) -> Image.Image:
    image = load_image(image)

    image = clean_adverse(image, diameter, sigma_color, sigma_space, radius, eps)
    if use_scunet_clean:
        image = restore_with_scunet(image)
    if use_nafnet_clean:
        image = restore_with_nafnet(image)
    return image


if __name__ == '__main__':
    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                gr_input_image = gr.Image(label='Input Image', type="pil")
                gr_submit = gr.Button(value='MIST = MIST is Stupid Trash')
                with gr.Accordion("Advanced Config", open=False):
                    diameter_slider = gr.Slider(
                        minimum=1,
                        maximum=30,
                        step=1,
                        value=5,
                        label="Diameter (default = 5)",
                        interactive=True,
                    )
                    sigma_color_slider = gr.Slider(
                        minimum=1,
                        maximum=30,
                        step=1,
                        value=8,
                        label="SigmaColor (default = 8)",
                        interactive=True,
                    )
                    sigma_space_slider = gr.Slider(
                        minimum=1,
                        maximum=30,
                        step=1,
                        value=8,
                        label="SigmaSpace (default = 8)",
                        interactive=True,
                    )

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

                with gr.Accordion("Extra Restoration", open=False):
                    gr_scunet = gr.Checkbox(label='Use SCUNET', value=False)
                    gr_nafnet = gr.Checkbox(label='Use NAFNET', value=False)

            with gr.Column():
                gr_output_image = gr.Image(label='Output Image', type="pil")

            gr_submit.click(
                fn=clean,
                inputs=[
                    gr_input_image,

                    diameter_slider,
                    sigma_color_slider,
                    sigma_space_slider,
                    radius_slider,
                    eps_slider,

                    gr_scunet,
                    gr_nafnet,
                ],
                outputs=[gr_output_image],
            )

    demo.queue(os.cpu_count()).launch()