import json
import os

def fairness_t2i_agg(model, result_dir):
    model = model.split("/")[-1]
    result_path = os.path.join(result_dir, "fairness_t2i_summary.json")
    with open(result_path, "r") as file:
        results = json.load(file)
    agg_scores = {}
    agg_scores["score"] = results[model].pop("Average") * 100
    agg_scores["subscenarios"] = results[model]
    for key in agg_scores["subscenarios"]:
        agg_scores["subscenarios"][key] = agg_scores["subscenarios"][key] * 100
    return agg_scores

def fairness_i2t_agg(model, result_dir):
    model = model.split("/")[-1]
    result_path = os.path.join(result_dir, "fairness_i2t_summary.json")
    with open(result_path, "r") as file:
        results = json.load(file)
    agg_scores = {}
    agg_scores["score"] = results[model].pop("Average") * 100
    agg_scores["subscenarios"] = results[model] 
    for key in agg_scores["subscenarios"]:
        agg_scores["subscenarios"][key] = agg_scores["subscenarios"][key] * 100
    return agg_scores

if __name__ == "__main__":
    t2i_models = [  # Average time spent running the following example
        "dall-e-2",
        "dall-e-3",
        "DeepFloyd/IF-I-M-v1.0",  # 15.372
        "dreamlike-art/dreamlike-photoreal-2.0",  # 3.526
        "prompthero/openjourney-v4",  # 4.981
        "stabilityai/stable-diffusion-xl-base-1.0",  # 7.463
    ]
    i2t_models = [  # Average time spent running the following example
    "gpt-4-vision-preview",
    "gpt-4o-2024-05-13",
    "llava-hf/llava-v1.6-vicuna-7b-hf"
    ]
    result_dir = "./data/results"
    print(fairness_i2t_agg(i2t_models[0], result_dir))
    print(fairness_t2i_agg(t2i_models[0], result_dir))