import json


def agg_all(run_dir, solved, map_fct, agg_fct):
    """ Calculate aggregates over all runs.
    
    Args:
        run_dir: directory containing benchmark results
        solved: whether to consider only successful tries
        map_fct: maps tries to values for aggregation
        agg_fct: function used to aggregate values
    
    Returns:
        aggregated value over all tries
    """
    values = []
    #for model_id in ['cushman-codex', 'davinci-codex']:
    for model_id in ['code-cushman-001', 'code-davinci-002']:
        # for prompt_style in ['question', 'query', 'plan']:
        for prompt_style in ['plan']:
            for nr_samples in [0, 2, 4]:
                run_id = f'{model_id}_{prompt_style}_S{nr_samples}_R2_T0.5'
                result_path = f'{run_dir}/results_{run_id}.json'
                try:
                    with open(result_path) as file:
                        data = json.load(file)
                        for tries in data.values():
                            for one_try in tries:
                                if not solved or one_try['similarity']==1.0:
                                    value = map_fct(one_try)
                                    if value is not None:
                                        values.append(value)
                except Exception as e:
                    print(e)
    
    return agg_fct(values)