eval_utils.py

from __future__ import annotations

import os
import time
import json
from openai import OpenAI
import random
from tqdm import tqdm
import math
import argparse

random.seed = 42

from gem5 import simulator
from config import PathConfig
from config import utilsConfig

expConfig = PathConfig()
utilConfig = utilsConfig()

"""Core utility Functions for generating and evaluating programs given output JSON files."""

def check_generated_programs_on_gem5(generated_programs_path: str, experiment_prefix: str, experiment_output_folder:str, include_src_tgt: bool = False) -> None:
    
    """
    Accept a JSON file in the format (unique_key, list of programs to evaluate) and evaluate each program on gem5.
    """
    
    TEST_CASE_DICT_FILE = utilConfig.test_case_dict
    ALL_TEST_INFO_FILE = utilConfig.all_test_info
    
    with open(generated_programs_path, "r") as f:
        generated_programs_dict = json.load(f)
    
    with open(TEST_CASE_DICT_FILE, "r") as f:
        test_case_dict = json.load(f)
    
    with open(ALL_TEST_INFO_FILE, "r") as f:
        all_test_info = [json.loads(line) for line in f]
    
    env = simulator.make(timeout_seconds_gem5=120, verbose=True, use_logical_cpus=True, port=80, workers=80, exit_early_on_fail=True)
    code_list = []
    test_cases_list = []
    problem_id_list = []
    test_ids_successfully_added_list = []
    
    for test_set_pair in all_test_info:
        src_prog = test_set_pair["src_prog"]
        tgt_prog = test_set_pair["tgt_prog"]
        problem_id = test_set_pair["problem_id"]
        unique_key = test_set_pair["unique_key"]
        test_cases = test_case_dict[problem_id]
        
        if include_src_tgt:
            code_list.append(src_prog)
            problem_id_list.append(problem_id)
            test_cases_list.append(test_cases)
            test_ids_successfully_added_list.append([unique_key, "src", src_prog])
            
            code_list.append(tgt_prog)
            problem_id_list.append(problem_id)
            test_cases_list.append(test_cases)
            test_ids_successfully_added_list.append([unique_key, "tgt", tgt_prog])

        for prog_ind in range(len(generated_programs_dict.get(unique_key, []))):
            prog = generated_programs_dict[unique_key][prog_ind]
            if len(prog) == 0:
                continue
            code_list.append(prog)
            problem_id_list.append(problem_id)
            test_cases_list.append(test_cases)
            test_ids_successfully_added_list.append([unique_key, prog_ind, prog])
    
    print(f"Total number of programs to evaluate on gem5: {len(code_list)}")
    print([x for x in code_list if len(x) == 0])
    results = env.submit_multiple_single_submissions(code_list, test_cases_list, problem_id_list, "gem5")
    timestamp = time.strftime("%Y%m%d-%H%M%S")
    
    with open(f"{experiment_output_folder}/{experiment_prefix}/{timestamp}_gem5_results.txt", "w") as f1:
        f1.write(str(results))
    
    with open(f"{experiment_output_folder}/{experiment_prefix}/{timestamp}_gem5_results.json", "w") as f:
        dict_results = [result.to_dict() for result in results]
        json.dump(dict_results, f)        
    
    with open(f"{experiment_output_folder}/{experiment_prefix}/{timestamp}_test_ids_successfully_added_list.json", "w") as f2:
        json.dump(test_ids_successfully_added_list, f2, indent=4)  

def compute_gem5_statistics(experiment_output_folder: str, experiment_prefix: str, default_val: float = 1.0, generate_dataset: bool = False, original_program_prompt_dict_file: str | None = None, original_prev_gen_file: str | None = None, new_generated_programs_file: str | None = None, include_src_tgt: bool = False) -> None:
    
    """
    Compute statistics from the gem5 results and save them multiple JSON files.
    
    Args:
        experiment_output_folder: Folder where the experiment results are stored
        experiment_prefix: Prefix for the experiment results files
        default_val: Default value to use for relative performance calculations (default: 1.0)
        generate_dataset: Whether to generate a new dataset of best performing programs based on the results (default: False).
        original_program_prompt_dict_file: If generate_dataset is True, the JSON file containing the original program prompts and generations to use for identifying the generations corresponding to the gem5 results (default: None).
        original_prev_gen_file: If generate_dataset is True, the JSON file containing the previous generation indices for each test ID to use for updating with the new best generation indices (default: None).
        new_generated_programs_file: If generate_dataset is True, the JSON file containing the newly generated programs corresponding to the generations in original_program_prompt_dict_file to use for creating the best program dataset (default: None).
        include_src_tgt: Whether to include the source and target programs in the evaluation (default: False).
    Returns:
        None
    """
    FOLDER_PATH = f"{experiment_output_folder}/{experiment_prefix}"
    
    all_files = os.listdir(FOLDER_PATH)
    
    gem5_result_json_file = f"{FOLDER_PATH}/{[file for file in all_files if file.endswith('_gem5_results.json')][0]}"
    
    successfully_added_test_ids_file = f"{FOLDER_PATH}/{[file for file in all_files if file.endswith('_test_ids_successfully_added_list.json')][0]}"
    
    successfully_compiled_test_ids_file = successfully_added_test_ids_file.replace("_added_list", "_compiled_list")
    
    runtime_file = gem5_result_json_file.replace("_gem5_results.json", "_runtime.json")
    
    scaled_perf_file = gem5_result_json_file.replace("_gem5_results.json", "_scaled_perf.json")
    
    REFERENCE_RUNTIME_PATH = utilConfig.reference_runtime_path 

    if include_src_tgt:
        #Assume reference runtime file has not yet been created.
        reference_runtime = {}
    else:
        #Load reference runtime file
        reference_runtime = json.load(open(REFERENCE_RUNTIME_PATH, "r"))
    
    test_set_result_dict = {}
    
    with open(successfully_added_test_ids_file) as f:
        test_ids_successfully_added_list = json.load(f)
        
    with open(gem5_result_json_file) as f1:
        results = json.load(f1)
    
    if generate_dataset:
        assert original_program_prompt_dict_file is not None, "original_program_prompt_dict must be provided if generate_dataset is True."
        all_prev_generations_dict = {}
        with open(f"{FOLDER_PATH}/{original_program_prompt_dict_file}", "r") as f:
            original_program_prompt_dict = json.load(f)
            for key in original_program_prompt_dict.keys():
                split_ids = key.split("_")
                custom_id = "_".join(split_ids[:-1])
                generation_id = int(split_ids[-1])
                if custom_id not in all_prev_generations_dict:
                    all_prev_generations_dict[custom_id] = [generation_id]
                else:
                    all_prev_generations_dict[custom_id].append(generation_id)
        
        assert new_generated_programs_file is not None, "new_generated_programs_file must be provided if generate_dataset is True."
        with open(f"{FOLDER_PATH}/{new_generated_programs_file}", "r") as f:
            new_generated_programs_dict = json.load(f)            
        
    #Set up the result dictionary with initial values, using reference values if needed.
    generations_list = ["src", "tgt", "file"]
    generation_ids = {}
    for success_test in test_ids_successfully_added_list:
        test_set_result_dict[success_test[0]] = {}
        generation_ids[success_test[0]] = []
        for generation_type in generations_list:
            if generation_type in ["src", "tgt"] and not include_src_tgt:
                #Use reference runtime
                test_set_result_dict[success_test[0]][generation_type] = [reference_runtime[success_test[0]][generation_type]]
            else:
                test_set_result_dict[success_test[0]][generation_type] = []
            
    successfully_compiled_test_ids = []
    success_compiled_dict = {}
    
    for test_id in test_set_result_dict.keys():
        success_compiled_dict[test_id] = []
    
    for result_ind in tqdm(range(len(results))):
        runtime = results[result_ind]["agg_runtime"]
        correctness = results[result_ind]["mean_acc"]
        
        # Filter out any generations that were not completely correct.
        if correctness == 1.0:
            test_id = test_ids_successfully_added_list[result_ind][0]
            if test_ids_successfully_added_list[result_ind][1] in ["src", "tgt"]:
                generation_type = test_ids_successfully_added_list[result_ind][1]
                test_set_result_dict[test_id][generation_type] = [runtime]
            else:
                generation_type = "file"
                test_set_result_dict[test_id][generation_type].append(runtime)
                if generate_dataset:
                    # Identify the generation from which the current result was obtained.
                    current_generation_index = test_ids_successfully_added_list[result_ind][1]
                    generation_ids[test_id].append(all_prev_generations_dict[test_id][current_generation_index])
                    success_compiled_dict[test_id].append(current_generation_index)
            successfully_compiled_test_ids.append(test_ids_successfully_added_list[result_ind] + [runtime])
        else:
            print(f"Test ID {test_ids_successfully_added_list[result_ind][0]} with generation number {test_ids_successfully_added_list[result_ind][1]} failed with correctness {correctness} and runtime {runtime}.")
    
    #Save the runtime and successfully compiled test IDs to a file. 

    with open(runtime_file, "w") as f2:
        json.dump(test_set_result_dict, f2)
    
    with open(successfully_compiled_test_ids_file, "w") as f3:
        json.dump(successfully_compiled_test_ids, f3)

    if generate_dataset:
        best_program_dataset_dict = {}
        
        if original_prev_gen_file is not None:
            with open(f"{experiment_output_folder}/{original_prev_gen_file}", "r") as f:
                original_prev_gen_dict = json.load(f)
        else:
            original_prev_gen_dict = {}
            for test_id in generation_ids.keys():
                original_prev_gen_dict[test_id] = []
                
        for test_id in test_set_result_dict.keys():
            if test_id not in original_prev_gen_dict:
                original_prev_gen_dict[test_id] = []
            #compute index of minimum runtime in the file list
            if len(test_set_result_dict[test_id]["file"]) > 0:
                fastest_runtime = min(test_set_result_dict[test_id]["file"])
                if fastest_runtime != math.inf:
                    assert len(test_set_result_dict[test_id]["file"]) == len(success_compiled_dict[test_id]), f"Mismatch in lengths for test ID {test_id}: {len(test_set_result_dict[test_id]['file'])} vs {len(success_compiled_dict[test_id])}"
                    min_runtime_index = test_set_result_dict[test_id]["file"].index(fastest_runtime)
                    original_prev_gen_dict[test_id].append(generation_ids[test_id][min_runtime_index])
                    best_program_dataset_dict[test_id] = new_generated_programs_dict[test_id][success_compiled_dict[test_id][min_runtime_index]]
                else:
                    any_gen_id = None
                    for key in original_program_prompt_dict.keys():
                        if test_id in key:
                            any_gen_id = key
                    if any_gen_id is None:
                        print(f"Warning: No generation found for test ID {test_id}. Using original program prompt.")
                        raise ValueError(f"No generation found for test ID {test_id}.")
                    best_program_dataset_dict[test_id] = original_program_prompt_dict[any_gen_id][1]["content"].split("\n # slower version: \n")[-1].split("\n# optimized version of the same code:\n")[0].strip()
            else:
                any_gen_id = None
                for key in original_program_prompt_dict.keys():
                    if test_id in key:
                        any_gen_id = key
                if any_gen_id is None:
                    print(f"Warning: No generation found for test ID {test_id}. Using original program prompt.")
                    raise ValueError(f"No generation found for test ID {test_id}.")
                best_program_dataset_dict[test_id] = original_program_prompt_dict[any_gen_id][1]["content"].split("\n # slower version: \n")[-1].split("\n# optimized version of the same code:\n")[0].strip()

            with open(f"{FOLDER_PATH}/best_program_dataset.json", "w") as f:
                json.dump(best_program_dataset_dict, f, indent=4)
            
            with open(f"{FOLDER_PATH}/prog_desc_tasks.json", "w") as f:
                prog_desc_tasks = {}
                with open(expConfig.prompts.prog_desc, "r") as f2:
                    system_prompt = f2.read()
                for test_id in best_program_dataset_dict.keys():
                    prog_desc_tasks[test_id] = [
                    {
                        "role": "system",
                        "content": system_prompt
                    },
                    {
                        "role": "user",
                        "content": " Source Program: \n" + str(best_program_dataset_dict[test_id])
                    }  
                    ]
                json.dump(prog_desc_tasks, f, indent=4)
            with open(f"{FOLDER_PATH}/prev_gen_indices.json", "w") as f:
                json.dump(original_prev_gen_dict, f, indent=4)
                
    for test_id in tqdm(test_set_result_dict.keys()):
        try:
            base_runtime = test_set_result_dict[test_id]["src"][0]
            if base_runtime == 0:
                print(f"Base runtime for {test_id} is 0.") 
            if base_runtime == math.inf:
                print(f"Base runtime for {test_id} is Infinity.")   
                continue      
            for generation_type in test_set_result_dict[test_id].keys():
                test_set_result_dict[test_id][generation_type] = [float("{:.2f}".format(float(base_runtime)/float(runtime))) if (runtime != math.inf and base_runtime!=math.inf) else default_val for runtime in test_set_result_dict[test_id][generation_type]]
        except Exception as e:
            print(f"Error in calculating relative performance for {test_id} with error: {e}")
            continue
    
    test_set_result_final_dict = {}
    for test_id in test_set_result_dict.keys():
        if test_set_result_dict[test_id]["src"][0] != math.inf:
            test_set_result_final_dict[test_id] = test_set_result_dict[test_id]

    with open(scaled_perf_file, "w") as f4:
        json.dump(test_set_result_final_dict, f4, indent=4)
    
    sums_dict = {}
    
    for generation_type in test_set_result_dict[list(test_set_result_dict.keys())[0]].keys():
        sums_dict[generation_type] = 0
    
    ignored_test_set_examples = 0
    for test_id in test_set_result_dict.keys():
        
        inf_counter = False
        for key in test_set_result_dict[test_id].keys():
            for perf in test_set_result_dict[test_id][key]:
                if perf == math.inf:
                    inf_counter = True
                    break
            if key=="src" and len(test_set_result_dict[test_id][key]) == 0:
                inf_counter = True
                break
       
        #If any of the performance values is infinity, ignore the test set example.
        
        if inf_counter:
            ignored_test_set_examples += 1
            continue
        
        for generation_type in test_set_result_dict[test_id].keys():
            if len(test_set_result_dict[test_id][generation_type]) != 0:
                max_val = max(test_set_result_dict[test_id][generation_type] + [default_val])
            else: 
                max_val = default_val
            sums_dict[generation_type] += max_val
            
    total_num_test_examples = len(test_set_result_dict.keys()) - ignored_test_set_examples
    print(f"Total number of test set examples: {len(test_set_result_dict.keys())}, Number of ignored test set examples: {ignored_test_set_examples}")
    #print(f"Sums_dict is: {sums_dict}")
    print("Mean performance for the different types of generations is:" + " ".join([f"{generation_type}: {sums_dict[generation_type]/total_num_test_examples}," for generation_type in sums_dict.keys()]))

def compute_gem5_statistics_no_contextual(experiment_output_folder: str, experiment_prefix: str, default_val: float = 1.0, generate_dataset: bool = False, original_program_prompt_dict_file: str | None = None, new_generated_programs_file: str | None = None, include_src_tgt: bool = False) -> None:
    """
    Compute statistics from the gem5 results and save them multiple JSON files.
    
    Args:
        experiment_output_folder: Folder where the experiment results are stored
        experiment_prefix: Prefix for the experiment results files
        default_val: Default value to use for relative performance calculations (default: 1.0)
        generate_dataset: Whether to generate a new dataset of best performing programs based on the results (default: False).
        original_program_prompt_dict_file: If generate_dataset is True, the JSON file containing the original program prompts and generations to use for identifying the generations corresponding to the gem5 results (default: None).
        new_generated_programs_file: If generate_dataset is True, the JSON file containing the newly generated programs corresponding to the generations in original_program_prompt_dict_file to use for creating the best program dataset (default: None).
        include_src_tgt: Whether to include the source and target programs in the evaluation (default: False).
    
    Returns:
        None
    """
    FOLDER_PATH = f"{experiment_output_folder}/{experiment_prefix}"
    
    all_files = os.listdir(FOLDER_PATH)
    
    gem5_result_json_file = f"{FOLDER_PATH}/{[file for file in all_files if file.endswith('_gem5_results.json')][0]}"
    
    successfully_added_test_ids_file = f"{FOLDER_PATH}/{[file for file in all_files if file.endswith('_test_ids_successfully_added_list.json')][0]}"
    
    successfully_compiled_test_ids_file = successfully_added_test_ids_file.replace("_added_list", "_compiled_list")
    
    runtime_file = gem5_result_json_file.replace("_gem5_results.json", "_runtime.json")
    
    scaled_perf_file = gem5_result_json_file.replace("_gem5_results.json", "_scaled_perf.json")
    
    REFERENCE_RUNTIME_PATH = utilConfig.reference_runtime_path

    if include_src_tgt:
        #Assume reference runtime file has not yet been created.
        reference_runtime = {}
    else:
        #Load reference runtime file
        reference_runtime = json.load(open(REFERENCE_RUNTIME_PATH, "r"))
    
    test_set_result_dict = {}
    
    with open(successfully_added_test_ids_file) as f:
        test_ids_successfully_added_list = json.load(f)
        
    with open(gem5_result_json_file) as f1:
        results = json.load(f1)
    
    if generate_dataset:
        assert original_program_prompt_dict_file is not None, "original_program_prompt_dict must be provided if generate_dataset is True."
        with open(f"{original_program_prompt_dict_file}", "r") as f:
            original_program_prompt_dict = json.load(f)
        
        assert new_generated_programs_file is not None, "new_generated_programs_file must be provided if generate_dataset is True."
        with open(f"{experiment_output_folder}/{new_generated_programs_file}", "r") as f:
            new_generated_programs_dict = json.load(f)            
        
    #Set up the result dictionary with initial values, using reference values if needed.
    generations_list = ["src", "tgt", "file"]
    for success_test in test_ids_successfully_added_list:
        test_set_result_dict[success_test[0]] = {}
        for generation_type in generations_list:
            if generation_type in ["src", "tgt"] and not include_src_tgt:
                #Use reference runtime
                test_set_result_dict[success_test[0]][generation_type] = [reference_runtime[success_test[0]][generation_type]]
            else:
                test_set_result_dict[success_test[0]][generation_type] = []
            
    successfully_compiled_test_ids = []
    success_compiled_dict = {}
    
    for test_id in test_set_result_dict.keys():
        success_compiled_dict[test_id] = []
    
    for result_ind in tqdm(range(len(results))):
        runtime = results[result_ind]["agg_runtime"]
        correctness = results[result_ind]["mean_acc"]
        
        # Filter out any generations that were not completely correct.
        if correctness == 1.0:
            test_id = test_ids_successfully_added_list[result_ind][0]
            if test_ids_successfully_added_list[result_ind][1] in ["src", "tgt"]:
                generation_type = test_ids_successfully_added_list[result_ind][1]
                test_set_result_dict[test_id][generation_type] = [runtime]
            else:
                generation_type = "file"
                test_set_result_dict[test_id][generation_type].append(runtime)
                if generate_dataset:
                    # Identify the generation from which the current result was obtained.
                    current_generation_index = test_ids_successfully_added_list[result_ind][1]
                    success_compiled_dict[test_id].append(current_generation_index)
            successfully_compiled_test_ids.append(test_ids_successfully_added_list[result_ind] + [runtime])
        else:
            print(f"Test ID {test_ids_successfully_added_list[result_ind][0]} with generation number {test_ids_successfully_added_list[result_ind][1]} failed with correctness {correctness} and runtime {runtime}.")
    
    #Save the runtime and successfully compiled test IDs to a file.
    with open(runtime_file, "w") as f2:
        json.dump(test_set_result_dict, f2)
    
    with open(successfully_compiled_test_ids_file, "w") as f3:
        json.dump(successfully_compiled_test_ids, f3)
    
    if generate_dataset:
        best_program_dataset_dict = {}
                
        for test_id in test_set_result_dict.keys():
            
            #compute index of minimum runtime in the file list
            if len(test_set_result_dict[test_id]["file"]) > 0:
                assert len(test_set_result_dict[test_id]["file"]) == len(success_compiled_dict[test_id]), f"Mismatch in lengths for test ID {test_id}: {len(test_set_result_dict[test_id]['file'])} vs {len(success_compiled_dict[test_id])}"
                fastest_runtime = min(test_set_result_dict[test_id]["file"])
                if fastest_runtime != math.inf:
                    min_runtime_index = test_set_result_dict[test_id]["file"].index(fastest_runtime)
                    best_program_dataset_dict[test_id] = new_generated_programs_dict[test_id][success_compiled_dict[test_id][min_runtime_index]]
                else:
                    any_gen_id = None
                    for key in original_program_prompt_dict.keys():
                        if test_id in key:
                            any_gen_id = key
                    if any_gen_id is None:
                        raise ValueError(f"No generation found for test ID {test_id}.")
                    # best_program_dataset_dict[test_id] = original_program_prompt_dict[any_gen_id][1]["content"].split("\n # slower version: \n")[-1].split("\n# optimized version of the same code:\n")[0].strip()
                    best_program_dataset_dict[test_id] = original_program_prompt_dict[any_gen_id][1]["content"].split("# slower version:")[-1].split("# optimized version of the same code:")[0].strip()
            else:
                any_gen_id = None
                for key in original_program_prompt_dict.keys():
                    if test_id in key:
                        any_gen_id = key
                if any_gen_id is None:
                    print(f"Warning: No generation found for test ID {test_id}. Using original program prompt.")
                    raise ValueError(f"No generation found for test ID {test_id}.")
                # best_program_dataset_dict[test_id] = original_program_prompt_dict[any_gen_id][1]["content"].split("# slower version:\n\n")[-1].split("\n# optimized version of the same code:")[0].strip()
                best_program_dataset_dict[test_id] = original_program_prompt_dict[any_gen_id][1]["content"].split("# slower version:")[-1].split("# optimized version of the same code:")[0].strip()

            with open(f"{FOLDER_PATH}/best_program_dataset.json", "w") as f:
                json.dump(best_program_dataset_dict, f, indent=4)
                
    for test_id in tqdm(test_set_result_dict.keys()):
        try:
            base_runtime = test_set_result_dict[test_id]["src"][0]
            if base_runtime == 0:
                print(f"Base runtime for {test_id} is 0.") 
            if base_runtime == math.inf:
                print(f"Base runtime for {test_id} is Infinity.")   
                continue      
            for generation_type in test_set_result_dict[test_id].keys():
                test_set_result_dict[test_id][generation_type] = [float("{:.2f}".format(float(base_runtime)/float(runtime))) if (runtime != math.inf and base_runtime!=math.inf) else default_val for runtime in test_set_result_dict[test_id][generation_type]]
        except Exception as e:
            print(f"Error in calculating relative performance for {test_id} with error: {e}")
            continue
    
    test_set_result_final_dict = {}
    for test_id in test_set_result_dict.keys():
        if test_set_result_dict[test_id]["src"][0] != math.inf:
            test_set_result_final_dict[test_id] = test_set_result_dict[test_id]

    with open(scaled_perf_file, "w") as f4:
        json.dump(test_set_result_final_dict, f4, indent=4)
    
    sums_dict = {}
    
    for generation_type in test_set_result_dict[list(test_set_result_dict.keys())[0]].keys():
        sums_dict[generation_type] = 0
    
    ignored_test_set_examples = 0
    for test_id in test_set_result_dict.keys():
        
        inf_counter = False
        for key in test_set_result_dict[test_id].keys():
            for perf in test_set_result_dict[test_id][key]:
                if perf == math.inf:
                    inf_counter = True
                    break
            if key=="src" and len(test_set_result_dict[test_id][key]) == 0:
                inf_counter = True
                break
       
        #If any of the performance values is infinity, ignore the test set example.
        
        if inf_counter:
            ignored_test_set_examples += 1
            continue
        
        for generation_type in test_set_result_dict[test_id].keys():
            if len(test_set_result_dict[test_id][generation_type]) != 0:
                max_val = max(test_set_result_dict[test_id][generation_type] + [default_val])
            else: 
                max_val = default_val
            sums_dict[generation_type] += max_val
            
    total_num_test_examples = len(test_set_result_dict.keys()) - ignored_test_set_examples
    print(f"Total number of test set examples: {len(test_set_result_dict.keys())}, Number of ignored test set examples: {ignored_test_set_examples}")
    #print(f"Sums_dict is: {sums_dict}")
    print("Mean performance for the different types of generations is:" + " ".join([f"{generation_type}: {sums_dict[generation_type]/total_num_test_examples}," for generation_type in sums_dict.keys()]))

"""Experiments using open-source LLMs."""

def experiment_analysis(model_output_folder: str, experiment_result_file:str, experiment_prefix: str, generate_dataset: bool, original_program_prompt_dict_file: str | None = None, original_prev_gen_file: str | None = None, new_generated_programs_file: str | None = None, no_contextual: bool = False, include_src_tgt: bool = False) -> None:
    
    """
    Run the none experiment with the specified model and output folder.
    
    Args:
        model: The model to use for the experiment
        model_output_folder: The folder where the experiment results will be saved
        experiment_result_file: The JSON file containing the generated programs to evaluate on gem5.
        experiment_prefix: The prefix to use for the experiment results files
        generate_dataset: Whether to generate a new dataset of best performing programs based on the results (default: False).
        original_program_prompt_dict_file: If generate_dataset is True, the JSON file containing the original program prompts and generations to use for identifying the generations corresponding to the gem5 results (default: None).
        original_prev_gen_file: If generate_dataset is True, the JSON file containing the previous generation indices for each test ID to use for updating with the new best generation indices (default: None).
        new_generated_programs_file: If generate_dataset is True, the JSON file containing the newly generated programs corresponding to the generations in original_program_prompt_dict_file to use for creating the best program dataset (default: None).
        no_contextual: Whether to compute statistics for RAS/AEGIS or for no_contextual ablations (default: False). 
        include_src_tgt: Whether to include source and target programs in the analysis (default: False).
    Returns:
        None
    """
    
    check_generated_programs_on_gem5(
        generated_programs_path=f"{model_output_folder}/{experiment_result_file}",
        experiment_prefix=experiment_prefix,
        experiment_output_folder=model_output_folder,
        include_src_tgt=include_src_tgt
    )
    
    if no_contextual:
        compute_gem5_statistics_no_contextual(
            experiment_output_folder=model_output_folder,
            experiment_prefix=experiment_prefix,
            default_val=1.0,
            generate_dataset=generate_dataset,
            original_program_prompt_dict_file=original_program_prompt_dict_file,
            new_generated_programs_file=experiment_result_file,
            include_src_tgt=include_src_tgt
        )      
    else:   
        if not generate_dataset:
            compute_gem5_statistics(
                experiment_output_folder=model_output_folder,
                experiment_prefix=experiment_prefix,
                default_val=1.0,
                include_src_tgt=include_src_tgt
            )
        else:
            compute_gem5_statistics(
                experiment_output_folder=model_output_folder,
                experiment_prefix=experiment_prefix,
                default_val=1.0,
                generate_dataset=generate_dataset,
                original_program_prompt_dict_file=original_program_prompt_dict_file,
                original_prev_gen_file=original_prev_gen_file,
                new_generated_programs_file=new_generated_programs_file,
                include_src_tgt=include_src_tgt
            )

def final_results_table(perf_files: list[str] = []):
    
    """
    Compute and print the final results used in the paper for the different experiments based on the scaled_perf JSON files generated from compute_gem5_statistics.
    
    Args:
        perf_files: List of JSON files containing the scaled performance results for the different experiments to include.
    
    Returns:
        None
    """
    
    for perf_file in perf_files:
        perf_src = 0
        perf_tgt = 0
        perf_best_mean = 0
        programs_considered = 0
        programs_optimized = 0
        tgt_programs_optimized = 0
        with open(perf_file) as f:
            data = json.load(f)
            for problem_id in data.keys():
                programs_considered += 1
                if len(data[problem_id].keys()) > 3:
                    print(f"Issue in number of keys for problem {problem_id}, has keys {data[problem_id].keys()}")
                for key in data[problem_id].keys():
                    if len(data[problem_id][key]) > 0:
                        if key == "src":
                            perf_src += data[problem_id][key][0]
                        elif key == "tgt":
                            tgt_improvement = data[problem_id][key][0]
                            perf_tgt += tgt_improvement
                            if tgt_improvement >= 1.1:
                                tgt_programs_optimized += 1
                        else:
                            max_perf_improvement = max([1.0] + data[problem_id][key])
                            perf_best_mean += max_perf_improvement
                            if max_perf_improvement >= 1.1:
                                programs_optimized += 1
                    else:
                        if key == "src":
                            perf_src += 1.0
                        elif key == "tgt":
                            perf_tgt += 1.0
                        else:
                            perf_best_mean += 1.0
        print(f"For file {perf_file}, number of test set programs considered were {programs_considered}, and mean performance was {perf_best_mean/programs_considered}, src mean performance was {perf_src/programs_considered}, and tgt mean performance was {perf_tgt/programs_considered}. Percentage of programs optimized is: {programs_optimized/programs_considered}. Percentage of programs optimized by the corresponding target is:  {tgt_programs_optimized/programs_considered} \n")

def plot_performance_improvement_over_time(model="qwen_3_coder", perf_files_for_plotting: dict = {}, measure: str = "Mean Best Speedup"): 
    """
    Plot the performance improvement over beam search steps for the different experiments based on the scaled_perf JSON files generated from compute_gem5_statistics.
    
    Args:
        model: The model for which the experiments were conducted, used for naming the saved plot file (default: "qwen_3_coder").
        perf_files_for_plotting: A dictionary mapping experiment names to lists of JSON files containing the scaled performance results for the different beam search steps to include in the plot.
        measure: The performance measure to plot, either "Mean Best Speedup" or "% Optimized" (default: "Mean Best Speedup").
    """
       
    import matplotlib.pyplot as plt
    final_averages_dict = {}
    starting_value = 1.0 if measure == "Mean Best Speedup" else 0.0
    for experiment in perf_files_for_plotting.keys():
        final_averages_dict[experiment] = []
        for perf_file in perf_files_for_plotting[experiment]:
            perf_src = 0
            perf_tgt = 0
            perf_best_mean = 0
            programs_considered = 0
            programs_optimized = 0
            tgt_programs_optimized = 0
            with open(perf_file) as f:
                data = json.load(f)
                for problem_id in data.keys():
                    programs_considered += 1
                    if len(data[problem_id].keys()) > 3:
                        print(f"Issue in number of keys for problem {problem_id}, has keys {data[problem_id].keys()}")
                    for key in data[problem_id].keys():
                        if len(data[problem_id][key]) > 0:
                            if key == "src":
                                perf_src += data[problem_id][key][0]
                            elif key == "tgt":
                                tgt_improvement = data[problem_id][key][0]
                                perf_tgt += tgt_improvement
                                if tgt_improvement >= 1.1:
                                    tgt_programs_optimized += 1
                            else:
                                max_perf_improvement = max([1.0] + data[problem_id][key])
                                perf_best_mean += max_perf_improvement
                                if max_perf_improvement >= 1.1:
                                    programs_optimized += 1
                        else:
                            if key == "src":
                                perf_src += 1.0
                            elif key == "tgt":
                                perf_tgt += 1.0
                            else:
                                perf_best_mean += 1.0
            print(f"Number of programs considered for {experiment} is {programs_considered}")
            if measure == "Mean Best Speedup":
                final_averages_dict[experiment].append(perf_best_mean/programs_considered)
            elif measure == "% Optimized":
                final_averages_dict[experiment].append(programs_optimized/programs_considered)
    print(f"Final averages dict is: {final_averages_dict}")
    plt.xticks(range(0,4))
    plt.grid(True)
    for experiment in final_averages_dict.keys():
        x = [i for i in range(len(final_averages_dict[experiment]) + 1)]
        y = [starting_value] + final_averages_dict[experiment]
        plt.plot(x, y, label=experiment)
        plt.xlabel("Number of Beam Search Steps")
        plt.ylabel(measure)
        plt.title(f"{measure} Over Beam Search Steps")
    plt.legend()    
    #plt.savefig(f"{self.nd_folder}/{self.model}/paper_experiments/{measure}_over_beam_search_steps.png")
    plt.savefig(f"edit_distance_analysis/{model}_images/{measure}_over_beam_search_steps.pdf")
    plt.close()

if __name__ == "__main__":
    
    parser = argparse.ArgumentParser(description="Run experiment analysis for a given model and experiment output folder.")
    
    parser.add_argument("--model_output_folder", type=str, required=True, help="The folder where the experiment results are saved.")
    parser.add_argument("--experiment_result_file", type=str, required=True, help="The JSON file containing the generated programs to evaluate on gem5, relative to the model_output_folder.")
    parser.add_argument("--experiment_prefix", type=str, required=True, help="The prefix to use for the experiment results files.")
    parser.add_argument("--generate_dataset", action="store_true", help="Whether to generate a new dataset of best performing programs based on the results.")
    parser.add_argument("--original_program_prompt_dict_file", type=str, default=None, help="If generate_dataset is True, the JSON file containing the original program prompts and generations to use for identifying the generations corresponding to the gem5 results, relative to the model_output_folder.")
    parser.add_argument("--original_prev_gen_file", type=str, default=None, help="If generate_dataset is True, the JSON file containing the previous generation indices for each test ID to use for updating with the new best generation indices, relative to the model_output_folder. Used in the RAS/AEGIS settings when retrieving one program pair per prompt to avoid resampling programs that already led to best program.")
    parser.add_argument("--new_generated_programs_file", type=str, default=None, help="If generate_dataset is True, the JSON file containing the newly generated programs corresponding to the generations in original_program_prompt_dict_file to use for creating the best program dataset, relative to the model_output_folder.")
    parser.add_argument("--no_contextual", action="store_true", help="Whether to compute statistics for RAS/AEGIS or for no_contextual ablations.")
    parser.add_argument("--include_src_tgt", action="store_true", help="Whether to include source and target programs in the analysis.")
    
    args = parser.parse_args()
    print(args)
    experiment_analysis(
            model_output_folder=args.model_output_folder,
            experiment_result_file=args.experiment_result_file,
            experiment_prefix=args.experiment_prefix,
            generate_dataset=args.generate_dataset,
            original_program_prompt_dict_file=args.original_program_prompt_dict_file,
            original_prev_gen_file=args.original_prev_gen_file,
            new_generated_programs_file=args.new_generated_programs_file,
            no_contextual=args.no_contextual,
            include_src_tgt=args.include_src_tgt       
        )
    
    """
    # Example Utility Usage for Plotting PIE results
        
    AEGIS_JSON_FILES = [
        "experiments/<model_folder>/aegis/1/<timestamp>_scaled_perf.json",
        "experiments/<model_folder>/aegis/2/<timestamp>_scaled_perf.json",
        "experiments/<model_folder>/aegis/3/<timestamp>_scaled_perf.json",
        "experiments/<model_folder>/aegis/4/<timestamp>_scaled_perf.json",    
    ]
    RAS_FILES = [
        "experiments/<model_folder>/ras/1/<timestamp>_scaled_perf.json",
        "experiments/<model_folder>/ras/2/<timestamp>_scaled_perf.json",
        "experiments/<model_folder>/ras/3/<timestamp>_scaled_perf.json",
        "experiments/<model_folder>/ras/4/<timestamp>_scaled_perf.json"
    ]

    RAS_NO_CONTEXTUAL_FILES = [
        "experiments/<model_folder>/no_contextual/1/<timestamp>_scaled_perf.json",
        "experiments/<model_folder>/no_contextual/2/<timestamp>_scaled_perf.json",
        "experiments/<model_folder>/no_contextual/3/<timestamp>_scaled_perf.json",
        "experiments/<model_folder>/no_contextual/4/<timestamp>_scaled_perf.json"
    ]

    AEGIS_NO_CONTEXTUAL_JSON_FILES = [
        "experiments/<model_folder>/aegis_no_contextual/1/<timestamp>_scaled_perf.json",
        "experiments/<model_folder>/aegis_no_contextual/2/<timestamp>_scaled_perf.json",
        "experiments/<model_folder>/aegis_no_contextual/3/<timestamp>_scaled_perf.json",
        "experiments/<model_folder>/aegis_no_contextual/4/<timestamp>_scaled_perf.json"
    ]
    experiment_list = [AEGIS_NO_CONTEXTUAL_JSON_FILES, AEGIS_JSON_FILES, RAS_NO_CONTEXTUAL_FILES, RAS_FILES]

    #Prevent cascading errors when e.g. retrieved examples exceed context or result in no output. In such cases, use the best generation from previous iterations if available to compute the performance improvements, instead of using an empty file list which would result in a performance of 1.0 and thus no improvement over the base program.
    
    for experiment in experiment_list:
        current_index = len(experiment)-1
        while(current_index > 0):
            final_raw_file = experiment[current_index]
            with open(final_raw_file, "r") as f:
                final_raw_data = json.load(f)
            prev_iterations_data = {}
            for i in range(current_index):
                with open(experiment[i+1], "r") as f:
                    prev_iterations_data[i+1] = json.load(f)
                    
            for key in final_raw_data.keys():
                if len(final_raw_data[key]["file"]) == 0:
                    for j in range(current_index, 0, -1):
                        if key in prev_iterations_data[j]:
                            if len(prev_iterations_data[j][key]["file"]) > 0:
                                final_raw_data[key]["file"] = prev_iterations_data[j][key]["file"]
                                break
                            
            new_filename = final_raw_file.split(".json")[0] + "_no_cascading.json"
            print(new_filename)
            with open(new_filename, "w") as f:
                json.dump(final_raw_data, f, indent=4)
            current_index -= 1

    final_results_table(
        perf_files= [
            "experiments/<model_folder>/instruct_only/<timestamp>_scaled_perf.json",
            "experiments/<model_folder>/ras/4/<timestamp>_scaled_perf_no_cascading.json",
            "experiments/<model_folder>/no_contextual/4/<timestamp>_scaled_perf_no_cascading.json",
            "experiments/<model_folder>/dynamic_retrieval/<timestamp>_scaled_perf.json",
            "experiments/<model_folder>/aegis/4/<timestamp>_scaled_perf_no_cascading.json",
            "experiments/<model_folder>/aegis_no_contextual/4/<timestamp>_scaled_perf_no_cascading.json",
        ]
    )
    
    for l in [RAS_FILES, RAS_NO_CONTEXTUAL_FILES, AEGIS_JSON_FILES, AEGIS_NO_CONTEXTUAL_JSON_FILES]:
        for i, element in enumerate(l):
            if i!=0:
                l[i] = element.split(".json")[0] + "_no_cascading.json"
        
    plot_performance_improvement_over_time(
        model="qwen_3_coder",
        perf_files_for_plotting = {
        "RAS (Contextual Retrieval)": RAS_FILES,
        "RAS (No Contextual)": RAS_NO_CONTEXTUAL_FILES,
        "AEGIS (Contextual Retrieval)": AEGIS_JSON_FILES,
        "AEGIS (No Contextual)": AEGIS_NO_CONTEXTUAL_JSON_FILES,
        }, 
        # measure = "Mean Best Speedup"
        measure = "% Optimized"
    )
    """