import csv
import itertools

if __name__ == '__main__':
    run_name = "IndivQLearningVarNearbyObs1"

    map_names = ["Pentagon", "ISR", "SUNY", "MIT"]
    randomize_starts = [True, False]
    shields = ["centralized", "decentralized", "none"]
    grid_world_obs_types = ["NearbyObsDiscrete", "FullObsDiscrete"]
    learner_anneal_eps = [(0.2, 0.2), (0.2, 0.01)]
    punish_unsafe_orig_actions = [(True, 0), (True, -10)]
    num_runs = 10

    with open(f"../../parallel_configs/{run_name}.csv", "w") as file:
        writer = csv.DictWriter(file, ["run_name", "shield", "punish_unsafe_orig_action",
                                       "punish_unsafe_orig_action_modifier", "randomize_starts", "map_type",
                                       "grid_world_map_name",
                                       "grid_world_obs_type", "learner_type", "learner_anneal_eps_start",
                                       "learner_anneal_eps_finish", "max_total_steps", "seed"])
        writer.writeheader()
        for run_type_idx, (map_name, random_start, shield, obs_type, (eps_anneal_start, eps_anneal_finish),
                           (punish_unsafe_action, unsafe_action_rew_modifier)) in enumerate(
            itertools.product(map_names, randomize_starts, shields, grid_world_obs_types, learner_anneal_eps,
                              punish_unsafe_orig_actions)):

            for run_num_of_same_type in range(num_runs):
                global_run_idx = run_type_idx * num_runs + run_num_of_same_type

                concat_run_name = run_name + "/" + str(global_run_idx) + "_" + str(run_type_idx) + "_" + str(
                    run_num_of_same_type)

                writer.writerow({
                    "run_name": concat_run_name,
                    "shield": shield,
                    "punish_unsafe_orig_action": punish_unsafe_action,
                    "punish_unsafe_orig_action_modifier": unsafe_action_rew_modifier,
                    "randomize_starts": random_start,
                    "map_type": "GridWorld",
                    "grid_world_map_name": map_name,
                    "grid_world_obs_type": obs_type,
                    "learner_type": "Individual_Q",
                    "learner_anneal_eps_start": eps_anneal_start,
                    "learner_anneal_eps_finish": eps_anneal_finish,
                    "max_total_steps": int(1e6),
                    "seed": run_num_of_same_type
                })
