import gradio as gr
import torch
from utils import *
import cv2
import numpy as np
import matplotlib.pyplot as plt
import io

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)

    HSV_image = cv2.cvtColor(color_image, cv2.COLOR_BGR2HSV)
    selected_harmomic_scheme = HarmonicScheme(str(template), int(angel))
    new_HSV_image = selected_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, selected_harmomic_scheme, "Source Hue")
    fig_1_cv = get_img_from_fig(fig1)
    fig2 = plothis(histo_2, selected_harmomic_scheme, "Target Hue")
    fig_2_cv = get_img_from_fig(fig2)
    # Convert HSV to BGR
    result_image = cv2.cvtColor(new_HSV_image, cv2.COLOR_HSV2BGR)

    cv2.imwrite('fig_1_cv.jpg', fig_1_cv)
    cv2.imwrite('fig_2_cv.jpg', fig_2_cv)
    cv2.imwrite('result_image.jpg', result_image)


    return ['result_image.jpg', 'fig_1_cv.jpg', 'fig_2_cv.jpg']

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'),
     gr.outputs.Image(type='file', label='Source Hue'),
     gr.outputs.Image(type='file', label='Target Hue'),],
    title=title,
    description=description,
    article=article,
    examples=examples,
    # css=css,
).launch(debug=True, enable_queue=True)