import gradio as gr
import pandas as pd
from PIL import Image
import os

# Load metadata
df = pd.read_csv('preferences.csv')

# Function to display image pairs
def show_image_pair(pair_id):
    record = df[df['pair_id'] == pair_id].iloc[0]
    img1_path = record['image1']
    img2_path = record['image2']
    img1 = Image.open(img1_path)
    img2 = Image.open(img2_path)
    return img1, img2, record['prompt'], record['label1'], record['label1_score'], record['label2'], record['label2_score'], record['label3'], record['label3_score']

# Create dropdown options
pair_ids = df['pair_id'].tolist()

# Define Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# MID-Space Dataset Viewer")
    with gr.Row():
        with gr.Column():
            pair_id_input = gr.Dropdown(label="Select Pair ID", choices=pair_ids, value=pair_ids[0])
            show_button = gr.Button("Show Image Pair")
        with gr.Column():
            img1 = gr.Image(label="Image 1")
            img2 = gr.Image(label="Image 2")
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", interactive=False)
    with gr.Row():
        label1 = gr.Textbox(label="Label 1", interactive=False)
        label1_score = gr.Number(label="Label 1 Score", interactive=False)
    with gr.Row():
        label2 = gr.Textbox(label="Label 2", interactive=False)
        label2_score = gr.Number(label="Label 2 Score", interactive=False)
    with gr.Row():
        label3 = gr.Textbox(label="Label 3", interactive=False)
        label3_score = gr.Number(label="Label 3 Score", interactive=False)
    
    show_button.click(
        show_image_pair, 
        inputs=[pair_id_input], 
        outputs=[img1, img2, prompt, label1, label1_score, label2, label2_score, label3, label3_score]
    )

if __name__ == "__main__":
    demo.launch()