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,
        use_adverse_clean: bool = True,
        use_scunet_clean: bool = True,
        use_nafnet_clean: bool = False
) -> Image.Image:
    image = load_image(image)
    if use_adverse_clean:
        image = clean_adverse(image)
    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='FUCK YOU MIST!')
                gr_adverse = gr.Checkbox(label='Use Adverse Clean', value=True)
                gr_scunet = gr.Checkbox(label='Use SCUNET', value=True)
                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, gr_adverse, gr_scunet, gr_nafnet],
                outputs=[gr_output_image],
            )

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