import os
import random
import textwrap

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 dynamic_clean_adverse(
        input_image: Image.Image,
        diameter_min: int = 4,
        diameter_max: int = 6,
        sigma_color_min: float = 6.0,
        sigma_color_max: float = 10.0,
        sigma_space_min: float = 6.0,
        sigma_space_max: float = 10.0,
        radius_min: int = 3,
        radius_max: int = 6,
        eps_min: float = 16.0,
        eps_max: float = 24.0,
        b_iters: int = 64,
        g_iters: int = 8,
):
    img = np.array(input_image).astype(np.float32)
    y = img.copy()

    for _ in range(b_iters):
        diameter = random.randint(diameter_min, diameter_max)
        sigma_color = random.random() * (sigma_color_max - sigma_color_min) + sigma_color_min
        sigma_space = random.random() * (sigma_space_max - sigma_space_min) + sigma_space_min
        y = cv2.bilateralFilter(y, diameter, sigma_color, sigma_space)

    for _ in range(g_iters):
        radius = random.randint(radius_min, radius_max)
        eps = random.random() * (eps_max - eps_min) + eps_min
        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_min: int = 4,
        diameter_max: int = 6,
        sigma_color_min: float = 6.0,
        sigma_color_max: float = 10.0,
        sigma_space_min: float = 6.0,
        sigma_space_max: float = 10.0,
        radius_min: int = 3,
        radius_max: int = 6,
        eps_min: float = 16.0,
        eps_max: float = 24.0,
        b_iters: int = 64,
        g_iters: int = 8,

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

    image = dynamic_clean_adverse(
        image,
        diameter_min, diameter_max,
        sigma_color_min, sigma_color_max,
        sigma_space_min, sigma_space_max,
        radius_min, radius_max,
        eps_min, eps_max,
        b_iters, g_iters
    )
    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():
            gr_markdown = gr.Markdown(textwrap.dedent("""
                Cleaner for [MIST](https://github.com/mist-project/mist-v2)(**M**IST **I**s **S**tupid **T**rash) noises.

                Inspired by https://github.com/lllyasviel/AdverseCleaner

                * **Update 2023.12.18**, allow random dynamic adversarial clean and iterate steps.
            """).strip())
        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', variant='primary')
                with gr.Accordion("Advanced Config", open=False):
                    with gr.Row():
                        gr_diameter_min = gr.Slider(
                            minimum=1, maximum=30, step=1, value=4,
                            label="Diameter Min (default = 4)", interactive=True,
                        )
                        gr_diameter_max = gr.Slider(
                            minimum=1, maximum=30, step=1, value=6,
                            label="Diameter Max (default = 6)", interactive=True,
                        )

                    with gr.Row():
                        gr_sigma_color_min = gr.Slider(
                            minimum=1, maximum=30, step=1, value=6,
                            label="SigmaColor Min (default = 6)", interactive=True,
                        )
                        gr_sigma_color_max = gr.Slider(
                            minimum=1, maximum=30, step=1, value=10,
                            label="SigmaColor Max (default = 10)", interactive=True,
                        )

                    with gr.Row():
                        gr_sigma_space_min = gr.Slider(
                            minimum=1, maximum=30, step=1, value=6,
                            label="SigmaSpace Min (default = 6)", interactive=True,
                        )
                        gr_sigma_space_max = gr.Slider(
                            minimum=1, maximum=30, step=1, value=10,
                            label="SigmaSpace Max (default = 10)", interactive=True,
                        )

                    with gr.Row():
                        gr_radius_min = gr.Slider(
                            minimum=1, maximum=30, step=1, value=3,
                            label="Radius Min (default = 3)", interactive=True,
                        )
                        gr_radius_max = gr.Slider(
                            minimum=1, maximum=30, step=1, value=6,
                            label="Radius Max (default = 6)", interactive=True,
                        )

                    with gr.Row():
                        gr_eps_min = gr.Slider(
                            minimum=1, maximum=30, step=1, value=16,
                            label="Accuracy Min (default = 16)", interactive=True,
                        )
                        gr_eps_max = gr.Slider(
                            minimum=1, maximum=30, step=1, value=24,
                            label="Accuracy Max (default = 24)", interactive=True,
                        )

                    with gr.Row():
                        gr_b_iters = gr.Slider(
                            minimum=1, maximum=256, step=1, value=64,
                            label="Bilateral Filter Iters (default = 64)", interactive=True,
                        )
                        gr_g_iters = gr.Slider(
                            minimum=1, maximum=32, step=1, value=8,
                            label="Guided Filter Iters (default = 8)", interactive=True,
                        )

                with gr.Accordion("Extra Restoration", open=False):
                    with gr.Row():
                        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,

                    gr_diameter_min,
                    gr_diameter_max,
                    gr_sigma_color_min,
                    gr_sigma_color_max,
                    gr_sigma_space_min,
                    gr_sigma_space_max,
                    gr_radius_min,
                    gr_radius_max,
                    gr_eps_min,
                    gr_eps_max,
                    gr_b_iters,
                    gr_g_iters,

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

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