import gradio as gr
from train import main
import spaces
import tensorflow as tf


# Function to calculate the aspect ratio
def calculate_aspect_ratio(width, height):
    return width / height

# Function to update image dimensions based on aspect ratio
def update_dimensions(value, aspect_ratio, dim):
    if dim == "width":
        new_height = int(value / aspect_ratio)
        return gr.update(value=value, visible=True), gr.update(value=new_height, visible=True)
    elif dim == "height":
        new_width = int(value * aspect_ratio)
        return gr.update(value=new_width, visible=True), gr.update(value=value, visible=True)

# Function to handle image size selection
def handle_image_size_selection(img_size, content_img):
    if img_size == "custom size":
        height, width, _ = content_img.shape
        aspect_ratio = calculate_aspect_ratio(width, height)
        return (gr.update(value=width, visible=True),
                gr.update(value=height, visible=True),
                gr.update(visible=True), 
                gr.update(visible=True), 
                aspect_ratio)
    else:
        return (gr.update(value=450, visible=False),
                gr.update(value=300, visible=False),
                gr.update(visible=False), 
                gr.update(visible=False), 
                None)

# Define the function to process images
# @spaces.GPU
def process_images(content_img, style_img, epochs, steps_per_epoch, learning_rate, content_loss_factor, style_loss_factor, img_size, img_width, img_height):
    print("Start processing")
    print(tf.config.list_physical_devices('GPU'))
    # with tf.device("/cuda:0"):
    output_img = main(content_img, style_img, epochs, steps_per_epoch, learning_rate, content_loss_factor, style_loss_factor, img_size, img_width, img_height)
    return output_img

# @spaces.GPU
def create_app():
    with gr.Blocks() as demo:
        aspect_ratio = gr.State(None)
        with gr.Row():
            with gr.Column(scale=1):
                content_img = gr.Image(type="numpy", label="Content Image")
                style_img = gr.Image(type="numpy", label="Style Image")
                process_button = gr.Button("Process")
                with gr.Accordion("Parameters", open=False):
                    epochs = gr.Slider(minimum=1, maximum=100, step=1, label="Epochs", value=20)
                    steps_per_epoch = gr.Slider(minimum=1, maximum=1000, step=1, label="Steps per Epoch", value=100)
                    learning_rate = gr.Slider(minimum=0.0001, maximum=0.1, step=0.0001, label="Learning Rate", value=0.01)
                    content_loss_factor = gr.Slider(minimum=0.1, maximum=10, step=0.1, label="Content Loss Factor", value=1.0)
                    style_loss_factor = gr.Slider(minimum=0.1, maximum=1000, step=0.1, label="Style Loss Factor", value=100.0)
                    img_size = gr.Dropdown(choices=["default size", "custom size"], label="Image Size", value="default size")
                    img_width = gr.Number(label="Image Width", value=450, visible=False)
                    img_height = gr.Number(label="Image Height", value=300, visible=False)
            with gr.Column(scale=1):
                output_img = gr.Image(label="Output Image")

        img_size.change(
            fn=handle_image_size_selection,
            inputs=[img_size, content_img],
            outputs=[img_width, img_height, img_width, img_height, aspect_ratio]
        )

        img_width.change(
            fn=lambda w, ar: update_dimensions(w, ar, "width"),
            inputs=[img_width, aspect_ratio],
            outputs=[img_width, img_height]
        )

        img_height.change(
            fn=lambda h, ar: update_dimensions(h, ar, "height"),
            inputs=[img_height, aspect_ratio],
            outputs=[img_width, img_height]
        )

        process_button.click(
            process_images,
            inputs=[content_img, style_img, epochs, steps_per_epoch, learning_rate, content_loss_factor, style_loss_factor, img_size, img_width, img_height],
            outputs=[output_img]
        )

    demo.launch()

if __name__ == "__main__":
    create_app()