Cramorant forme reverts on switch-out

fp/battle_bots/mcts_randbats/main.py

@@ -3,6 +3,8 @@
 
 from poke_engine import MctsResult
 
+import constants
+from data.pkmn_sets import TeamDatasets
 from fp.battle import Battle
 from config import FoulPlayConfig
 
@@ -63,6 +65,44 @@ class BattleBot(Battle):
     def __init__(self, *args, **kwargs):
         super(BattleBot, self).__init__(*args, **kwargs)
 
+    def _num_battles_randombattles(self, revealed_pkmn):
+        return int(FoulPlayConfig.parallelism * PARALLELISM_MULTIPLIER[revealed_pkmn])
+
+    def _search_time_per_battle_randombattles(self, num_battles):
+        return FoulPlayConfig.search_time_ms // (
+            max(num_battles // FoulPlayConfig.parallelism, 1)
+        )
+
+    def _num_battles_battle_factory(self):
+        opponent_active_num_sets = len(
+            TeamDatasets.get_all_remaining_sets(self.opponent.active)
+        )
+        if self.team_preview or (
+            self.opponent.active.hp > 0
+            and opponent_active_num_sets > (FoulPlayConfig.parallelism * 2)
+        ):
+            num_battles = FoulPlayConfig.parallelism * 4
+        elif opponent_active_num_sets > FoulPlayConfig.parallelism:
+            num_battles = FoulPlayConfig.parallelism * 2
+        else:
+            num_battles = FoulPlayConfig.parallelism
+
+        return num_battles
+
+    def _search_time_per_battle_battle_factory(self, num_battles):
+        if (
+            self.team_preview
+            or self.turn < 3
+            or (self.time_remaining is not None and self.time_remaining > 60)
+        ):
+            search_time_per_battle = FoulPlayConfig.search_time_ms
+        else:
+            search_time_per_battle = self._search_time_per_battle_randombattles(
+                num_battles
+            )
+
+        return search_time_per_battle
+
     def find_best_move(self):
         if self.team_preview:
             self.user.active = self.user.reserve.pop(0)
@@ -72,16 +112,28 @@ def find_best_move(self):
         if self.opponent.active is not None:
             revealed_pkmn += 1
 
-        num_battles = int(
-            FoulPlayConfig.parallelism * PARALLELISM_MULTIPLIER[revealed_pkmn]
-        )
+        if self.battle_type == constants.RANDOM_BATTLE:
+            num_battles = self._num_battles_randombattles(revealed_pkmn)
+            search_time_per_battle = self._search_time_per_battle_randombattles(
+                num_battles
+            )
+        elif self.battle_type == constants.BATTLE_FACTORY:
+            num_battles = self._num_battles_battle_factory()
+            search_time_per_battle = self._search_time_per_battle_battle_factory(
+                num_battles
+            )
+        else:
+            raise ValueError("Unsupported battle type: {}".format(self.battle_type))
+
         battles = prepare_random_battles(self, num_battles)
         for b, _ in battles:
             fill_in_opponent_unrevealed_pkmn(b)
 
         logger.info("Searching for a move using MCTS...")
-        search_time_per_battle = FoulPlayConfig.search_time_ms // (
-            max(num_battles // FoulPlayConfig.parallelism, 1)
+        logger.info(
+            "Sampling {} battles at {}ms each".format(
+                num_battles, search_time_per_battle
+            )
         )
         with ProcessPoolExecutor(max_workers=FoulPlayConfig.parallelism) as executor:
             futures = []

fp/battle_modifier.py

@@ -364,6 +364,12 @@ def switch_or_drag(battle, split_msg, switch_or_drag="switch"):
                 )
             )
 
+        if side.active.name in ["cramorantgulping", "cramorantgorging"]:
+            logger.info(
+                "Resetting {} to 'cramorant' on switch out".format(side.active.name)
+            )
+            side.active.name = "cramorant"
+
     if side_name == "user" and user_just_switched_into_zoroark(battle, switch_or_drag):
         logger.info(
             "User switched/dragged into Zoroark - replacing the split_msg pokemon"

tests/test_battle_modifiers.py

@@ -263,6 +263,17 @@ def setUp(self):
         self.battle.opponent.active = self.opponent_active
         self.battle.opponent.reserve = []
 
+    def test_cramorantgulping_reverts_to_cramorant_in_switchout(self):
+        self.battle.opponent.active.name = "cramorantgulping"
+        split_msg = ["", "switch", "p2a: caterpie", "Caterpie, L100, M", "100/100"]
+        switch_or_drag(self.battle, split_msg)
+
+        self.assertEqual("caterpie", self.battle.opponent.active.name)
+        self.assertIn("cramorant", [p.name for p in self.battle.opponent.reserve])
+        self.assertNotIn(
+            "cramorantgulping", [p.name for p in self.battle.opponent.reserve]
+        )
+
     def test_user_switching_in_zaciancrowned_properly_re_initializes_stats(self):
         self.battle.request_json = {
             "active": [],