
# Image configuration for PettingZoo environments.
pz_encoded_images_args = {
    "time_limit": 2,
    "render_mode": "human",
    "image_encoder": "ResNet18",
    "image_encoder_use_cuda": False,
    "image_encoder_batch_size": 1,
    "trainable_cnn": False,
    "kwargs": "",
    "seed": 2024,
}

pz_raw_images_args = {
    "time_limit": 2,
    "render_mode": "human",
    "image_encoder": "ResNet18",
    "image_encoder_use_cuda": False,
    "image_encoder_batch_size": 1,
    "trainable_cnn": True,
    "kwargs": "",
    "seed": 2024,
}

# Environment key lists.
pettingzoo_fully_coop_env_keys = [
    "pistonball_v6",
    "cooperative_pong_v5",
    "entombed_cooperative_v3",
    "space_invaders_v2"
]

pettingzoo_fully_coop_partial_obs_env_keys = [
    "entombed_cooperative_v3",
    "space_invaders_v2"
]

pettingzoo_non_fully_coop_env_atari_keys = [
    "basketball_pong_v3",
    "boxing_v2",
    "combat_plane_v2",
    "combat_tank_v2",
    "double_dunk_v3",
    "entombed_competitive_v3",
    "flag_capture_v2",
    "foozpong_v3",
    "ice_hockey_v2",
    "joust_v3",
    "mario_bros_v3",
    "maze_craze_v3",
    "othello_v3",
    "pong_v3",
    "quadrapong_v4",
    "space_war_v2",
    "surround_v2",
    "tennis_v3",
    "video_checkers_v4",
    "volleyball_pong_v3",
    "warlords_v3",
    "wizard_of_wor_v3"
]

pettingzoo_non_fully_coop_env_butterfly_mpe_sisl_keys = [
    "knights_archers_zombies_v10",
    "simple_v3",
    "simple_adversary_v3",
    "simple_crypto_v3",
    "simple_push_v3",
    "simple_reference_v3",
    "simple_speaker_listener_v4",
    "simple_spread_v3",
    "simple_tag_v3",
    "simple_world_comm_v3",
    "multiwalker_v9",
    "pursuit_v4",
    "waterworld_v4",
]

pettingzoo_classic_keys = [
    "chess_v6",
    "connect_four_v3",
    "gin_rummy_v4",
    "go_v5",
    "hanabi_v5",
    "leduc_holdem_v4",
    "rps_v2",
    "texas_holdem_no_limit_v6",
    "texas_holdem_v4",
    "tictactoe_v3",
]

overcooked_env_keys = [
    "coordination_ring",
    "asymmetric_advantages",
    "cramped_room",
    "counter_circuit",
    "random3",
    "random0",
    "unident",
    "forced_coordination",
    "soup_coordination",
    "small_corridor",
    "simple_tomato",
    "simple_o_t",
    "simple_o",
    "schelling_s",
    "schelling",
    "m_shaped_s",
    "long_cook_time",
    "large_room",
    "forced_coordination_tomato",
    "cramped_room_tomato",
    "cramped_room_o_3orders",
    "asymmetric_advantages_tomato",
    "bottleneck",
    "cramped_corridor",
    "counter_circuit_o_1order",
]

pressureplate_env_keys = [
    "pressureplate-linear-4p-v0",
    "pressureplate-linear-5p-v0",
    "pressureplate-linear-6p-v0"
]

lbf_env_keys = [
    "lbforaging:Foraging-4s-11x11-3p-2f-coop-v2",
    "lbforaging:Foraging-2s-11x11-3p-2f-coop-v2",
    "lbforaging:Foraging-2s-8x8-3p-2f-coop-v2",
    "lbforaging:Foraging-2s-9x9-3p-2f-coop-v2",
    "lbforaging:Foraging-7s-20x20-5p-3f-coop-v2",
    "lbforaging:Foraging-2s-12x12-2p-2f-coop-v2",
    "lbforaging:Foraging-8s-25x25-8p-5f-coop-v2",
    "lbforaging:Foraging-7s-30x30-7p-4f-coop-v2",
    "lbforaging:Foraging-4s-11x11-3p-2f-coop-v3",
    "lbforaging:Foraging-2s-11x11-3p-2f-coop-v3",
    "lbforaging:Foraging-2s-8x8-3p-2f-coop-v3",
    "lbforaging:Foraging-2s-9x9-3p-2f-coop-v3",
    "lbforaging:Foraging-7s-20x20-5p-3f-coop-v3",
    "lbforaging:Foraging-2s-12x12-2p-2f-coop-v3",
    "lbforaging:Foraging-8s-25x25-8p-5f-coop-v3",
    "lbforaging:Foraging-7s-30x30-7p-4f-coop-v3",
]

rware_env_keys = [
    "rware:rware-small-4ag-hard-v1",
    "rware:rware-tiny-4ag-hard-v1",
    "rware:rware-tiny-2ag-hard-v1",
    "rware:rware-small-4ag-hard-v2",
    "rware:rware-tiny-4ag-hard-v2",
    "rware:rware-tiny-2ag-hard-v2",
]

mpe_env_keys = [
    "mpe:SimpleSpeakerListener-v0",
    "mpe:SimpleSpread-3-v0",
    "mpe:SimpleSpread-4-v0",
    "mpe:SimpleSpread-5-v0",
    "mpe:SimpleSpread-8-v0",
    "mpe:MultiSpeakerListener-v0",
    "mpe:SimpleAdversary-v0",
    "mpe:SimpleCrypto-v0",
    "mpe:SimplePush-v0",
    "mpe:SimpleReference-v0",
    "mpe:SimpleTag-v0",
    "mpe:SimpleWorldComm-v0",
    "mpe:ClimbingSpread-v0",
]

capturetarget_env_keys = ["CaptureTarget-6x6-1t-2a-v0"]

boxpushing_env_keys = ["BoxPushing-6x6-2a-v0"]

algorithms = {
    "QMIX": {
        "config": "qmix",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "target_update_interval_or_tau": 4,
        },
    },
    "QMIX_ns": {
        "config": "qmix_ns",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "target_update_interval_or_tau": 4,
        },
    },
    "CDS": {
        "config": "cds",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "target_update_interval": 4,
        },
    },
    "COMA": {
        "config": "coma",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "batch_size_run": 2,
            "target_update_interval_or_tau": 4,
        },
    },
    "COMA_ns": {
        "config": "coma_ns",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "batch_size_run": 2,
            "target_update_interval_or_tau": 4,
        },
    },
    "EMC": {
        "config": "emc",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "burn_in_period": 2,
            "target_update_interval": 4,
        },
    },
    "EOI": {
        "config": "eoi",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "batch_size_run": 2,
            "target_update_interval_or_tau": 4,
        },
    },
    "HAPPO": {
        "config": "happo",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "batch_size_run": 2,
            "data_chunk_length": 3,
        },
    },
    "IA2C": {
        "config": "ia2c",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "batch_size_run": 2,
        },
    },
    "IA2C_ns": {
        "config": "ia2c_ns",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "batch_size_run": 2,
            "target_update_interval_or_tau": 4,
        },
    },
    "IPPO": {
        "config": "ippo",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "batch_size_run": 2,
        },
    },
    "IPPO_ns": {
        "config": "ippo_ns",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "batch_size_run": 2,
            "target_update_interval_or_tau": 4,
        },
    },
    "IQL": {
        "config": "iql",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "target_update_interval_or_tau": 4,
        },
    },
    "IQL_ns": {
        "config": "iql_ns",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "target_update_interval_or_tau": 4,
        },
    },
    "MAA2C": {
        "config": "maa2c",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "batch_size_run": 2,
            "target_update_interval_or_tau": 4,
        },
    },
    "MAA2C_ns": {
        "config": "maa2c_ns",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "batch_size_run": 2,
            "target_update_interval_or_tau": 4,
        },
    },
    "MADDPG": {
        "config": "maddpg",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "target_update_interval_or_tau": 4,
        },
    },
    "MADDPG_ns": {
        "config": "maddpg_ns",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
        },
    },
    "MAPPO": {
        "config": "mappo",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "batch_size_run": 2,
            "target_update_interval_or_tau": 4,
        },
    },
    "MAPPO_ns": {
        "config": "mappo_ns",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "batch_size_run": 2,
            "target_update_interval_or_tau": 4,
        },
    },
    "MASER": {
        "config": "maser",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "target_update_interval": 4,
        },
    },
    "MAT-DEC": {
        "config": "mat_dec",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "batch_size_run": 2,
            "target_update_interval_or_tau": 4,
        },
    },
    "QPLEX": {
        "config": "qplex",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "target_update_interval": 4,
        },
    },
    "VDN": {
        "config": "vdn",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "target_update_interval_or_tau": 4,
        },
    },
    "VDN_ns": {
        "config": "vdn_ns",
        "algo_args": {
            "test_nepisode": 2,
            "test_interval": 2,
            "t_max": 10,
            "log_interval": 2,
            "runner_log_interval": 2,
            "learner_log_interval": 2,
            "batch_size": 2,
            "buffer_size": 2,
            "target_update_interval_or_tau": 4,
        },
    },
}
