#!/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, sigma_y: float) -> str:
    out_name = image_path.split("/")[-1].split(".")[0]
    subprocess.run(
        shlex.split(
            f"python main.py --config confs/inet256.yml --deg colorization --scale 1 --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 = [
        [
            "sample_images/monarch_gray.png",
            "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus",
            0,
        ],
        [
            "sample_images/tiger_gray.png",
            "tiger, Panthera tigris",
            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)
                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,
                sigma_y,
            ],
        )

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