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


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)
    # Convert HSV to BGR
    result_image = cv2.cvtColor(new_HSV_image, cv2.COLOR_HSV2BGR)

    # 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)

    hue_plots = np.concatenate((fig_1_cv, fig_2_cv), axis=1)

    cv2.imwrite('hue.jpg', hue_plots)
    cv2.imwrite('result_image.jpg', result_image)


    return ['result_image.jpg', 'hue.jpg']

title = 'Color Harmonization'
description = 'Compute Color Harmonization with different templates based on Color Harmonization paper by Daniel Cohen-Or et al. More metrics for user interfaces on https://interfacemetrics.aalto.fi'
article = "<p style='text-align: center'></p>"
examples = [['./examples/aim_interface.png', "V", 25], ['./examples/esfahan_unsplash.jpeg', "I", 352]]
# css = ".output_image, .input_image {height: 40rem !important; width: 100% !important;}"

gr.Interface(
    inference,
    [gr.inputs.Image(type='file', label='Original Image'),
     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='Harmonized Image'),
     gr.outputs.Image(type='file', label='Hue')],
    title=title,
    description=description,
    article=article,
    examples=examples,
    # css=css,
).launch(debug=False, enable_queue=True)