Resolve "controller hotplugging"

Closes #118 (closed)

cooperative_cuisine/pygame_2d_vis/gui.py

@@ -67,7 +67,7 @@ class PlayerKeySet:
         pickup_key: pygame.key,
         switch_key: pygame.key,
         players: list[str],
-        joystick: int,
+        joystick: int | None,
     ):
         """Creates a player key set which contains information about which keyboard keys control the player.
 
@@ -109,6 +109,8 @@ class PlayerKeySet:
                     return
         self.current_player = self.controlled_players[self.current_idx]
 
+    def __repr__(self) -> str:
+        return f"Keyset(current={self.current_player}, players={self.controlled_players}, joy={self.joystick})"
 
 class PyGameGUI:
     """Visualisation of the overcooked environment and reading keyboard inputs using pygame."""
@@ -232,7 +234,7 @@ class PyGameGUI:
                 pickup_key=pygame.K_e,
                 switch_key=pygame.K_SPACE,
                 players=players,
-                joystick=0,
+                joystick=None,
             )
             key_set2 = PlayerKeySet(
                 move_keys=[pygame.K_LEFT, pygame.K_RIGHT, pygame.K_UP, pygame.K_DOWN],
@@ -240,10 +242,16 @@ class PyGameGUI:
                 pickup_key=pygame.K_o,
                 switch_key=pygame.K_p,
                 players=players,
-                joystick=1,
+                joystick=None,
             )
             key_sets = [key_set1, key_set2]
 
+            if self.joysticks:
+                for idx, key in enumerate(self.joysticks.keys()):
+                    if idx >= len(key_sets):
+                        break
+                    key_sets[idx].joystick = key
+
             if disjunct:
                 key_set1.set_controlled_players(players[::2])
                 key_set2.set_controlled_players(players[1::2])
@@ -1450,7 +1458,7 @@ class PyGameGUI:
         if env_info.status_code == 403:
             raise ValueError(f"Forbidden Request: {env_info.json()['detail']}")
         elif env_info.status_code == 409:
-            print("CONFLICT")
+            log.warning("CONFLICT")
         env_info = env_info.json()
         assert isinstance(env_info, dict), "Env info must be a dictionary"
         self.current_env_id = env_info["env_id"]
@@ -1584,7 +1592,6 @@ class PyGameGUI:
             if answer.status_code == 200:
                 answer_json = answer.json()
                 self.player_info = answer_json["player_info"]["0"]
-                print("TUTORIAL PLAYER INFO", self.player_info)
                 self.level_info = answer_json["level_info"]
                 self.player_info = {self.player_info["player_id"]: self.player_info}
             else:
@@ -1605,7 +1612,6 @@ class PyGameGUI:
             if answer.status_code == 200:
                 answer_json = answer.json()
                 self.player_info = answer_json["player_info"]
-                print("GAME PLAYER INFO", self.player_info)
 
                 self.level_info = answer_json["level_info"]
                 self.last_level = self.level_info["last_level"]
@@ -1859,7 +1865,7 @@ class PyGameGUI:
             self.get_game_connection()
         else:
             self.menu_state = MenuStates.Start
-            print(
+            log.warning(
                 "COULD NOT START STUDY; Response:",
                 answer.status_code,
                 answer.json()["detail"],
@@ -2018,21 +2024,28 @@ class PyGameGUI:
                         # joystick, filling up the list without needing to create them manually.
                         joy = pygame.joystick.Joystick(event.device_index)
                         self.joysticks[joy.get_instance_id()] = joy
-                        print(f"Joystick {joy.get_instance_id()} connected")
+
+                        for key_set in self.key_sets:
+                            if key_set.joystick is None:
+                                key_set.joystick = joy.get_instance_id()
+                                break
+                        log.debug(f"Joystick {joy.get_instance_id()} connected")
 
                     # disconnect joystick
                     if event.type == pygame.JOYDEVICEREMOVED:
                         del self.joysticks[event.instance_id]
-                        print(f"Joystick {event.instance_id} disconnected")
-                        print("Number of joysticks:", pygame.joystick.get_count())
+                        for key_set in self.key_sets:
+                            if key_set.joystick == event.instance_id:
+                                key_set.joystick = None
+                        log.debug(f"Joystick {event.instance_id} disconnected")
+                        log.debug(f"Number of joysticks:"+str(pygame.joystick.get_count()))
 
                     # Press enter key or controller start button instead of mouse button press
                     if (
                         event.type == pygame.JOYBUTTONDOWN
                         and any(
                             [
-                                self.joysticks and self.joysticks[i].get_button(7)
-                                for i in range(len(self.joysticks))
+                                joy.get_button(7) for joy in self.joysticks.values()
                             ]
                         )
                         or (
@@ -2157,4 +2170,4 @@ if __name__ == "__main__":
         args.manager_ids,
         CONNECT_WITH_STUDY_SERVER=True,
         debug=args.do_study,
-    )
+    )
\ No newline at end of file