import csv
import itertools

if __name__ == '__main__':
    run_name = "ReplicateIngyFigure9-3"

    map_names = ["Pentagon", "ISR", "SUNY", "MIT"]
    randomize_starts = [False]
    shields = ["none"]
    grid_world_obs_types = ["FullObsDiscrete"]
    learner_anneal_eps = [(1.0, 0.05)]
    punish_unsafe_orig_actions = [(False, 0)]
    num_runs = 10
    collision_penalties = [0, -25, -50, -75, -100]

    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", "learner_evaluation_epsilon", "seed",
                                       "grid_world_collision_penalty", "grid_world_agents_bounce", "max_num_episodes"])
        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), collision_penalty) in enumerate(
            itertools.product(map_names, randomize_starts, shields, grid_world_obs_types, learner_anneal_eps,
                              punish_unsafe_orig_actions, collision_penalties)):

            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,
                    "learner_evaluation_epsilon": eps_anneal_finish,
                    "seed": run_num_of_same_type,
                    "grid_world_collision_penalty": collision_penalty,
                    "max_num_episodes": int(1e3),
                    "grid_world_agents_bounce": True
                })
