from notebooks.analysis_utils import * 
import gymnasium as gym
from multiprocessing import Pool

env_name = "molecular-generation" # change this to your env name

env = gym.make(f'gfn_challenges/{env_name}')
model_list = dir_to_names(f"../output_full/{env_name}")


def process_model(model):
    try:
        # change the num_of_samples to match your env, force_update would recalculate the val errs when the results exist
        # In the paper, I used 200000 for discrete envs, and 50000 fo r continuous envs
        gen_val_errs(model[0], env, num_of_samples=200000, force_update = True) 
    except Exception as e:
        print(f"Error processing model {model[0]}: {e}")
        pass

# Set the number of worker processes
process_pool_size = 8

if __name__ == "__main__":
    # gen_val_errs(model_list[2][0], env, num_of_samples=200000, force_update = False) 
    with Pool(process_pool_size) as pool:
        pool.map(process_model, model_list)