import os
import gradio as gr
import util
from random import randint
import sys
from subprocess import call
from PIL import Image

is_colab = util.is_google_colab()

def run_cmd(command):
    try:
        print(command)
        call(command, shell=True)
    except KeyboardInterrupt:
        print("Process interrupted")
        sys.exit(1)

def inference(img, size, type):
    _id = randint(1, 10000)
    INPUT_DIR = "/tmp/input_image" + str(_id) + "/"
    OUTPUT_DIR = "/tmp/output_image" + str(_id) + "/"
    run_cmd("rm -rf " + INPUT_DIR)
    run_cmd("rm -rf " + OUTPUT_DIR)
    run_cmd("mkdir " + INPUT_DIR)
    run_cmd("mkdir " + OUTPUT_DIR)
    img.save(INPUT_DIR + "1.jpg", "JPEG")

    run_cmd("python inference_manga.py "+ os.path.join(INPUT_DIR, "1.jpg") + " " + os.path.join(OUTPUT_DIR, "1_out.jpg"))
    img_out = Image.open(os.path.join(OUTPUT_DIR, "1_out.jpg"))
    if size is "x2":
        img_out = img_out.resize((s//2 for s in img_out.size))
    return [img_out]

input_image = gr.Image(type="pil", label="Input")
upscale_type = gr.Radio(["Manga", "Anime", "General"], label="Select the type of picture you want to upscale:", value="Manga")
upscale_size = gr.Radio(["x4", "x2"], label="Upscale by:", value="x4")
output_image = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery").style(grid=[2], height="auto")

demo = gr.Interface(
    inference,
    inputs=[input_image, upscale_size, upscale_type],
    outputs=[output_image]
)

demo.launch(debug=is_colab, share=is_colab, inline=is_colab)