# https://huggingface.co/spaces/Norod78/ComicsHeroU2Net/blob/main/app.py as template/basics
import os
from pathlib import Path
import subprocess
import sys
import PIL
import cv2 as cv
import torch
from PIL import Image
import gradio as gr
import numpy as np

torch.set_grad_enabled(False)


# https://en.wikipedia.org/wiki/Unsharp_masking
# https://stackoverflow.com/a/55590133/1495606
def unsharp_mask(image, kernel_size=(5, 5), sigma=1.0, amount=2.0, threshold=0):
    """Return a sharpened version of the image, using an unsharp mask."""
    blurred = cv.GaussianBlur(image, kernel_size, sigma)
    sharpened = float(amount + 1) * image - float(amount) * blurred
    sharpened = np.maximum(sharpened, np.zeros(sharpened.shape))
    sharpened = np.minimum(sharpened, 255 * np.ones(sharpened.shape))
    sharpened = sharpened.round()
    if threshold > 0:
        low_contrast_mask = np.absolute(image - blurred) < threshold
        np.copyto(sharpened, image, where=low_contrast_mask)
    return sharpened


def normPRED(d):
    ma = np.max(d)
    mi = np.min(d)

    dn = (d - mi) / (ma - mi)

    return dn


def array_to_np(array_in):
    array_in = normPRED(array_in)
    array_in = np.squeeze(255.0 * array_in)
    array_in = np.transpose(array_in, (1, 2, 0))
    return array_in


def image_as_array(image_in):
    image_in = np.array(image_in, np.float32)
    tmpImg = np.zeros((image_in.shape[0], image_in.shape[1], 3))
    image_in = image_in / np.max(image_in)
    if image_in.shape[2] == 1:
        tmpImg[:, :, 0] = (image_in[:, :, 0] - 0.485) / 0.229
        tmpImg[:, :, 1] = (image_in[:, :, 0] - 0.485) / 0.229
        tmpImg[:, :, 2] = (image_in[:, :, 0] - 0.485) / 0.229
    else:
        tmpImg[:, :, 0] = (image_in[:, :, 0] - 0.485) / 0.229
        tmpImg[:, :, 1] = (image_in[:, :, 1] - 0.456) / 0.224
        tmpImg[:, :, 2] = (image_in[:, :, 2] - 0.406) / 0.225

    tmpImg = tmpImg.transpose((2, 0, 1))
    image_out = np.expand_dims(tmpImg, 0)
    return image_out


def face2hero(model,
              img: Image.Image,
              size: int,
              unsharpen: bool = False
              ) -> Image.Image:
    temp_img = img.resize((size, size))
    input = torch.Tensor(image_as_array(temp_img))
    results = model(input)
    hero_np_image = array_to_np(results[1].detach().numpy())
    if unsharpen:
        hero_np_image = unsharp_mask(hero_np_image)
    hero_image = Image.fromarray(hero_np_image.astype(np.uint8))

    # output = img_concat_h(array_to_image(aligned_img), hero_image)
    del results

    return hero_image


def inference(img: Image.Image, unsharpen, model_path):
    model = torch.jit.load(model_path)
    model.eval()
    out = face2hero(model, img, 320, unsharpen)
    out = out.resize((img.width, img.height))
    return out


title = "Comic StyleTransfer U2Net"
description = "Style transfer a face into one of a \"Comics Style Transfer\""
article = ""
examples = []
print(torch.cuda.is_available())
demo = gr.Interface(
    inference,
    inputs=[gr.Image(type="pil", label="Input"), gr.Checkbox(label="Enable Unsharpen"),
            gr.Dropdown(show_label=True, label="Select Model", choices=list(Path(".").glob(f"*.jit.pt")),
                        max_choices=1)],
    outputs=[gr.Image(type="pil", label="Output")],
    title=title,
    description=description,
    article=article,
    examples=examples,
    allow_flagging="never"
)

demo.queue()
demo.launch()