import os
import gradio as gr
from datasets import load_dataset
import random

auth_token = os.environ.get("auth_token")
iiw_400 = load_dataset('google/imageinwords', token=auth_token, name="IIW-400")

def display_iiw_data(index):
    data = iiw_400['test'][index]
    image_html = f'<img src="{data["image/url"]}" style="width:100%; max-width:800px; height:auto;">'
    iiw_text = f"<h2>IIW Human Descriptions:</h2><p style='font-size: 16px'>{data['IIW']}</p>"
    iiw_p5b_text = f"<h2>IIW PaLI 5B Predictions:</h2><p style='font-size: 16px'>{data['IIW-P5B']}</p>"
    ratings = "<h2>Ratings:</h2>"
    if data['iiw-human-sxs-iiw-p5b'] is not None:
        for key, value in data['iiw-human-sxs-iiw-p5b'].items():
            key = key.split("metrics/")[-1]
            emoji = ""
            if key == "Comprehensiveness":
                emoji = "📚"  # Book
            elif key == "Specificity":
                emoji = "🎯"  # Bullseye
            elif key == "Hallucination":
                emoji = "👻"  # Ghost
            elif key == "First few line(s) as tldr":
                emoji = "🔍"  # Magnifying Glass Tilted Left
            elif key == "Human Like":
                emoji = "👤"  # Bust in Silhouette
            ratings += f"<p style='font-size: 16px'>{emoji} <strong>{key}</strong>: {value}</p>"
    return image_html, iiw_text, iiw_p5b_text, ratings

def random_index():
    while True:
        index = random.randint(0, len(iiw_400['test']) - 1)
        if iiw_400['test'][index]['iiw-human-sxs-iiw-p5b'] is not None:
            return index

demo = gr.Blocks()

with demo:
    gr.Markdown("# Slide across the slider to see various examples from IIW-400")

    with gr.Column():
        slider = gr.Slider(minimum=0, maximum=400)
        with gr.Row():
            index = random_index()
            with gr.Column():
                image_output = gr.HTML(display_iiw_data(index)[0])
            with gr.Column():
                iiw_text_output = gr.HTML(display_iiw_data(index)[1])
                iiw_p5b_text_output = gr.HTML(display_iiw_data(index)[2])
                ratings_output = gr.HTML(display_iiw_data(index)[3])

    slider.change(display_iiw_data, inputs=[slider], outputs=[image_output, iiw_text_output, iiw_p5b_text_output, ratings_output])

demo.launch(debug=True)