import gradio as gr
import torch
from utils import *

torch.hub.download_url_to_file(
    'https://github.com/aalto-ui/aim/raw/aim2/backend/data/tests/input_values/wikipedia.org_website.png',
    'wikipedia.org_website.png')
torch.hub.download_url_to_file(
    'https://github.com/aalto-ui/aim/raw/aim2/backend/data/tests/input_values/aalto.fi_website.png',
    'aalto.fi_website.png')


def inference(img, template, angel):
    color_image = cv2.imread(img.name, cv2.IMREAD_COLOR)
    height, width, _ = color_image.shape

    # Resize if it is bigeer than 960 * 800
    if width > height:
        if width > 960:  # 3/4 * 1280
            coef_div = width / 960.0
            color_image = cv2.resize(color_image, dsize=(int(width / coef_div), int(height / coef_div)),
                                     interpolation=cv2.INTER_CUBIC)
    else:
        if height > 800:  # 800
            coef_div = height / 800.0
            color_image = cv2.resize(color_image, dsize=(int(width / coef_div), int(height / coef_div)),
                                     interpolation=cv2.INTER_CUBIC)

    HSV_image = cv2.cvtColor(color_image, cv2.COLOR_BGR2HSV)
    selected_harmomic_scheme = HarmonicScheme(str(template), int(angel))
    new_HSV_image = best_harmomic_scheme.hue_shifted(HSV_image, num_superpixels=-1)

    # Compute shifted histogram
    histo_1 = count_hue_histogram(HSV_image)
    histo_2 = count_hue_histogram(new_HSV_image)

    # Create Hue Plots
    fig1 = plothis(histo_1, best_harmomic_scheme, "Source Hue")
    fig_1_cv = get_img_from_fig(fig1)
    fig2 = plothis(histo_2, best_harmomic_scheme, "Target Hue")
    fig_2_cv = get_img_from_fig(fig2)

    # Stack Hue Plots
    vis = np.concatenate((fig_1_cv, fig_2_cv), axis=0)
    # Convert HSV to BGR
    result_image = cv2.cvtColor(new_HSV_image, cv2.COLOR_HSV2BGR)

    # Final output
    canvas = np.full((800, 960, 3), (255, 255, 255), dtype=np.uint8)
    # compute center offset
    x_center = (960 - width) // 2
    y_center = (800 - height) // 2
    # copy img image into center of result image
    canvas[y_center:y_center + height, x_center:x_center + width] = result_image

    # Combine
    output = np.concatenate((vis, canvas), axis=1)
    cv2.imwrite('output.png', output)

    return ['output.png']


title = 'Color Harmonization'
description = 'Compute Color Harmonization with Different Templates'
article = "<p style='text-align: center'></p>"
examples = [['wikipedia.org_website.png'], ['aalto.fi_website.png']]
css = ".output_image, .input_image {height: 40rem !important; width: 100% !important;}"

gr.Interface(
    inference,
    [gr.inputs.Image(type='file', label='Input'),
     gr.inputs.Dropdown(["X", "Y", "T", "I", "mirror_L", "L", "V", "i"],
                        default="X",
                        label="Template"),
     gr.inputs.Slider(0, 359, label="Angle")],
    [gr.outputs.Image(type='file', label='Color Harmonization of Output Image')],
    title=title,
    description=description,
    article=article,
    examples=examples,
    css=css,
).launch(debug=True, enable_queue=True)