from PIL import Image, UnidentifiedImageError
from utils import *


def get_evaluation_data(ds):
    evaluation_data = []
    for i in range(len(ds)):
        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://visionlmsftw-vibe-testing-images.hf.space/image/" + str(i),
            "prompt": ds[i]["prompt"],
            "category": ds[i]["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