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

# Specific values for percentage and complexity
percentage_values = [10, 30, 50, 70, 100]
complexity_values = [16, 32]

# Function to load and display images based on user selection
def display_images(percentage_idx, complexity_idx):
    # Map the slider index to the actual value
    percentage = percentage_values[percentage_idx]
    complexity = complexity_values[complexity_idx]
    
    # 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
with gr.Blocks() as demo:
    gr.Markdown("# Raw vs. Embeddings Inference Results")
    gr.Markdown("Use the sliders to adjust the percentage of data for training and task complexity.")

    # Layout for vertical side-by-side sliders
    with gr.Row():
        # Column for percentage slider
        with gr.Column():
            percentage_slider = gr.Slider(minimum=0, maximum=4, step=1, value=0, 
                                          label="Percentage of Data for Training (10, 30, 50, 70, 100)", 
                                          interactive=True, orientation="vertical")
        
        # Column for complexity slider
        with gr.Column():
            complexity_slider = gr.Slider(minimum=0, maximum=1, step=1, value=0, 
                                          label="Task Complexity (16, 32)", 
                                          interactive=True, orientation="vertical")

    # Outputs (display the images side by side and set a smaller size for the images)
    with gr.Row():
        raw_img = gr.Image(label="Raw Channels", type="pil", width=300, height=300, interactive=False)  # Smaller image size
        embeddings_img = gr.Image(label="Embeddings", type="pil", width=300, height=300, interactive=False)  # Smaller image size
    
    # Trigger image updates when sliders change
    percentage_slider.change(fn=display_images, inputs=[percentage_slider, complexity_slider], outputs=[raw_img, embeddings_img])
    complexity_slider.change(fn=display_images, inputs=[percentage_slider, complexity_slider], outputs=[raw_img, embeddings_img])

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