import gradio as gr
import os
from PIL import Image

# Paths to the images folder
RAW_PATH = os.path.join("images", "raw")
EMBEDDINGS_PATH = os.path.join("images", "embeddings")

# Function to load and display images based on user selection
def display_images(percentage, complexity):
    # Generate the paths to the images
    raw_image_path = os.path.join(RAW_PATH, f"percentage_{percentage}_complexity_{complexity}.png")
    embeddings_image_path = os.path.join(EMBEDDINGS_PATH, f"percentage_{percentage}_complexity_{complexity}.png")
    
    # Load images using PIL
    raw_image = Image.open(raw_image_path)
    embeddings_image = Image.open(embeddings_image_path)
    
    # Return the loaded images
    return raw_image, embeddings_image

# Define the Gradio interface
data_percentage_options = [10, 30, 50, 70, 100]
task_complexity_options = [16, 32]

# Define the layout and appearance of the UI
with gr.Blocks() as demo:
    gr.Markdown("# Raw vs. Embeddings Inference Results")
    gr.Markdown("Select a data percentage and task complexity to view the corresponding inference result for raw channels and embeddings.")
    
    # Inputs (data percentage and task complexity)
    with gr.Row():
        percentage_dropdown = gr.Dropdown(data_percentage_options, label="Percentage of Data for Training", value=10)
        complexity_radio = gr.Radio(task_complexity_options, label="Task Complexity", value=16)
    
    # Outputs (display the images side by side)
    with gr.Row():
        raw_img = gr.Image(label="Raw Channels", type="pil", interactive=False)
        embeddings_img = gr.Image(label="Embeddings", type="pil", interactive=False)
    
    # Update images instantaneously without clicking submit
    percentage_dropdown.change(fn=display_images, inputs=[percentage_dropdown, complexity_radio], outputs=[raw_img, embeddings_img], live=True)
    complexity_radio.change(fn=display_images, inputs=[percentage_dropdown, complexity_radio], outputs=[raw_img, embeddings_img], live=True)

# Launch the app
if __name__ == "__main__":
    demo.launch()