#!/usr/bin/env python

from __future__ import annotations

import json
import shlex
import subprocess

import gradio as gr


def run(image_path: str, class_index: int, scale: str, sigma_y: float) -> str:
    out_name = image_path.split('/')[-1].split('.')[0]
    subprocess.run(shlex.split(
        f'python main.py --config confs/inet256.yml --resize_y --deg sr_averagepooling --scale {scale} --class {class_index} --path_y {image_path} --save_path {out_name} --sigma_y {sigma_y}'
    ),
                   cwd='DDNM/hq_demo')
    return f'DDNM/hq_demo/results/{out_name}/final/00000.png'


def create_demo():
    examples = [
        [
            'DDNM/hq_demo/data/datasets/gts/inet256/323.png',
            'monarch, monarch butterfly, milkweed butterfly, Danaus plexippus',
            '4',
            0,
        ],
        [
            'DDNM/hq_demo/data/datasets/gts/inet256/orange.png',
            'orange',
            '4',
            0,
        ],
        [
            'DDNM/hq_demo/data/datasets/gts/inet256/monarch.png',
            'monarch, monarch butterfly, milkweed butterfly, Danaus plexippus',
            '4',
            0.5,
        ],
        [
            'DDNM/hq_demo/data/datasets/gts/inet256/bear.png',
            'brown bear, bruin, Ursus arctos',
            '4',
            0,
        ],
        [
            'DDNM/hq_demo/data/datasets/gts/inet256/flamingo.png',
            'flamingo',
            '2',
            0,
        ],
        [
            'DDNM/hq_demo/data/datasets/gts/inet256/kimono.png',
            'kimono',
            '2',
            0,
        ],
        [
            'DDNM/hq_demo/data/datasets/gts/inet256/zebra.png',
            'zebra',
            '4',
            0,
        ],
    ]

    with open('imagenet_classes.json') as f:
        imagenet_class_names = json.load(f)

    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                image = gr.Image(label='Input image', type='filepath')
                class_index = gr.Dropdown(label='Class name',
                                          choices=imagenet_class_names,
                                          type='index',
                                          value=950)
                scale = gr.Dropdown(label='Scale',
                                    choices=['2', '4', '8'],
                                    value='4')
                sigma_y = gr.Number(label='sigma_y', value=0, precision=2)
                run_button = gr.Button('Run')
            with gr.Column():
                result = gr.Image(label='Result', type='filepath')

        gr.Examples(
            examples=examples,
            inputs=[
                image,
                class_index,
                scale,
                sigma_y,
            ],
        )

        run_button.click(
            fn=run,
            inputs=[
                image,
                class_index,
                scale,
                sigma_y,
            ],
            outputs=result,
        )
    return demo