from PIL import Image, UnidentifiedImageError
from utils import *


def get_evaluation_data(ds):
    evaluation_data = []
    for i in range(len(ds)):
        try:
            img = ds[i]["image"]
            thumbnail_img = img.copy()
            thumbnail_img.thumbnail((256, 256))
            evaluation_data.append({
                "id": ds[i]["ex_id"],
                "image_thumbnail": image_to_base64(thumbnail_img),
                "image_full": image_to_base64(img),
                "image_full_url": "https://huggingface.co/", # Dummy
                "prompt": ds[i]["prompt"],
                "category": ds[i]["category"]
            })
        except (UnidentifiedImageError, OSError, ValueError): # To handle .heic images -> can be removed when dataset is fixed
            img = Image.new("RGB", (256, 256), color="white")
            evaluation_data.append({
                "id": i,
                "image_thumbnail": image_to_base64(img),
                "image_full": image_to_base64(img),
                "image_full_url": "https://huggingface.co/", # Dummy
                "prompt": "Dummy prompt",
                "category": "Dummy category"
            })
    return evaluation_data

def get_model_names(ds_results):
    models = list(set(ds_results['model_id']))
    return models

def get_responses(ds_results):
    responses = {}
    
    for model in set(ds_results['model_id']):
        model_responses = [
            row["model_response"]
            for row in ds_results
            if row["model_id"] == model
        ]
        responses[model] = {i: resp for i, resp in enumerate(model_responses)}
    
    return responses