import gradio as gr
from torchvision.utils import save_image
from attack import Attacker
import argparse


def do_attack(img, eps, step_size, steps, progress=gr.Progress()):
    args=argparse.Namespace()
    args.out_dir='./'
    args.target='auto'
    args.eps=eps
    args.step_size=step_size
    args.steps=steps
    args.test_atk=False

    step = progress.tqdm(range(steps))

    def pdg_prog(ori_images, images, labels):
        step.update(1)

    attacker = Attacker(args, pgd_callback=pdg_prog)
    atk_img, noise = attacker.attack_(img)
    attacker.save_image(atk_img, noise, 'out.png')
    return 'out.png'

with gr.Blocks(title="Anime AI Detect Fucker Demo", theme="dark") as demo:
    gr.HTML('<a href="https://github.com/7eu7d7/anime-ai-detect-fucker">github repo</a>')

    with gr.Row():
        eps = gr.Slider(label="eps (Noise intensity)", minimum=1, maximum=16, step=1, value=1)
        step_size = gr.Slider(label="Noise step size", minimum=0.001, maximum=16, step=0.001, value=0.136)
    with gr.Row():
        steps = gr.Slider(label="step count", minimum=1, maximum=100, step=1, value=20)
        model_name = gr.Dropdown(label="attack target",
                                 choices=["auto", "human", "ai"],
                                 value="auto", show_label=True)

    input_image = gr.Image(label="Clean Image", type="pil")

    atk_btn = gr.Button("Attack")

    with gr.Column():
        output_image = gr.Image(label="Attacked Image")

    atk_btn.click(fn=do_attack,
                      inputs=[input_image, eps, step_size, steps],
                      outputs=output_image)

demo.launch()