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

# Determine the absolute path to preferences.csv
current_dir = os.path.dirname(os.path.abspath(__file__))
csv_path = os.path.join(current_dir, 'preferences.csv')

# Check if preferences.csv exists
if not os.path.isfile(csv_path):
    print(f"Error: 'preferences.csv' not found at {csv_path}")
    sys.exit(1)

# Load metadata
try:
    df = pd.read_csv(csv_path)
except Exception as e:
    print(f"Error reading 'preferences.csv': {e}")
    sys.exit(1)

# Function to display image pairs
def show_image_pair(pair_id):
    try:
        record = df[df['pair_id'] == pair_id].iloc[0]
    except IndexError:
        return None, None, "Invalid Pair ID.", "", "", "", "", "", ""
    
    img1_path = os.path.join(current_dir, record['image1'])
    img2_path = os.path.join(current_dir, record['image2'])
    
    # Check if images exist
    if not os.path.isfile(img1_path):
        return None, None, "Image 1 not found.", "", "", "", "", "", ""
    if not os.path.isfile(img2_path):
        return None, None, "Image 2 not found.", "", "", "", "", "", ""
    
    try:
        img1 = Image.open(img1_path)
    except Exception as e:
        img1 = None
        print(f"Error opening Image 1: {e}")
    
    try:
        img2 = Image.open(img2_path)
    except Exception as e:
        img2 = None
        print(f"Error opening Image 2: {e}")
    
    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")
    
    # Pair Selection Row
    with gr.Row():
        with gr.Column(scale=1):
            pair_id_input = gr.Dropdown(label="Select Pair ID", choices=pair_ids, value=pair_ids[0])
        with gr.Column(scale=1):
            show_button = gr.Button("Show Image Pair")
    
    # Images Row
    with gr.Row():
        with gr.Column(scale=1):
            img1 = gr.Image(label="Image 1")
        with gr.Column(scale=1):
            img2 = gr.Image(label="Image 2")
    
    # Label 1 and Label 1 Score Row
    with gr.Row():
        with gr.Column(scale=2):
            label1 = gr.Textbox(label="Label 1", interactive=False, lines=1)
        with gr.Column(scale=1):
            label1_score = gr.Textbox(label="Label 1 Score", interactive=False, lines=1)
    
    # Label 2 and Label 2 Score Row
    with gr.Row():
        with gr.Column(scale=2):
            label2 = gr.Textbox(label="Label 2", interactive=False, lines=1)
        with gr.Column(scale=1):
            label2_score = gr.Textbox(label="Label 2 Score", interactive=False, lines=1)
    
    # Label 3 and Label 3 Score Row
    with gr.Row():
        with gr.Column(scale=2):
            label3 = gr.Textbox(label="Label 3", interactive=False, lines=1)
        with gr.Column(scale=1):
            label3_score = gr.Textbox(label="Label 3 Score", interactive=False, lines=1)
    
    # Prompt Row
    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(label="Prompt", interactive=False, lines=2)
    
    # Define Button Click Action
    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()