import json
import os
from pathlib import Path

def process_json_files(directory, suffix="_updated"):
    # Iterate through all JSON files in the directory
    for filename in os.listdir(directory):
        if filename.endswith(".json") and "USACO" in filename:
            file_path = os.path.join(directory, filename)
            
            # Read the JSON file
            with open(file_path, 'r') as f:
                data = json.load(f)
            
            # Extract sdict from raw_eval_results
            sdict = data['raw_eval_results']['sdict']
            
            # Calculate successful_tasks and failed_tasks
            successful_tasks = [key for key in sdict if float(sdict[key][0]['result']['fraction_passed']) == 1]
            failed_tasks = [key for key in sdict if float(sdict[key][0]['result']['fraction_passed']) < 1]
            
            # Add new key-value pairs to the results
            data['results']['successful_tasks'] = successful_tasks
            data['results']['failed_tasks'] = failed_tasks
            
            # Create new filename with suffix
            new_filename = f"{Path(filename).stem}{suffix}{Path(filename).suffix}"
            new_file_path = os.path.join(directory, new_filename)
            
            # Write updated data to new file
            with open(new_file_path, 'w') as f:
                json.dump(data, f, indent=4)
            
            print(f"Processed {filename} and saved as {new_filename}")

# Usage
directory_path = "/Users/benediktstroebl/Documents/GitHub/leaderboard/evals_live"
process_json_files(directory_path)