import gradio as gr
import cv2
import numpy as np
from ColorTransfer import ImgProcess

def initializeImgProcess(img1 = None, img2 = None):
    return ImgProcess(img1, img2)

class GradioInterface:
    def __init__(self):
        self.imgProcess_obj = None
        self.imgProcess_obj = self.getImgProcess_obj()
        height = 300
        self.fun_percent = 1.0
        with gr.Blocks() as self.app:
            gr.Markdown("<h2 style='text-align: center;'>Image Style Transfer App</h2>")
            with gr.Row():
                with gr.Column(scale = 1):
                    self.img1 = gr.Image(label="Upload First Image", type="pil", height = height)
                    self.img2 = gr.Image(label="Upload Second Image", type="pil", height = height)

                with gr.Column(scale = 1):
                    self.output_img1 = gr.Image(label="Transformed Image1 - Color Transfer", height = height)
                    self.output_img2 = gr.Image(label="Transformed Image2 - Histogram Match", height = height)

            with gr.Row():
                with gr.Column(scale = 1):
                    percent_slider = gr.Slider(label="Slide for fun", minimum=0.1, maximum=1.0, step=0.1, value = 1.0)
                    percent_slider.change(self.updateFunPercent, inputs = percent_slider)
                with gr.Column(scale = 1):
                    transform_button = gr.Button("Transfer Style")            
                    transform_button.click(self._transfer_style, inputs=[self.img1, self.img2], outputs = [self.output_img1, self.output_img2])
                    swap_button = gr.Button("Swap Input Images") 
                    swap_button.click(self.__swapInputImages, inputs=[self.img1, self.img2], outputs = [self.img1, self.img2])

    def getImgProcess_obj(self):
        if self.imgProcess_obj != None:
            return self.imgProcess_obj
        
        else:
            return initializeImgProcess()

    def _transfer_style(self, img1, img2):
        return self.imgProcess_obj.loadAndTransfer(img1, img2, self.fun_percent)
    
    def updateFunPercent(self, val):
        self.fun_percent = val

    def __swapInputImages(self, img1, img2):
        return img2, img1
    
    def launch_app(self):
        self.app.launch(share=True)
    
if __name__ == "__main__":
    print("Starting")

    app_obj = GradioInterface()

    app_obj.launch_app()