From ea4434db2a24f948a8784b92b7f4891881d914c0 Mon Sep 17 00:00:00 2001
From: Pawkkie <61265402+Pawkkie@users.noreply.github.com>
Date: Tue, 21 Jan 2025 18:07:15 -0500
Subject: [PATCH] Switch trapping AI will consider Trace (#6059)

---
 include/battle_ai_util.h      |  1 +
 src/battle_ai_switch_items.c  |  5 +++--
 test/battle/ai/ai_switching.c | 19 ++++++++++++++++---
 3 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/include/battle_ai_util.h b/include/battle_ai_util.h
index f0adcd8af0..552faa1306 100644
--- a/include/battle_ai_util.h
+++ b/include/battle_ai_util.h
@@ -80,6 +80,7 @@ bool32 IsAbilityOfRating(u32 ability, s8 rating);
 bool32 AI_IsAbilityOnSide(u32 battlerId, u32 ability);
 bool32 AI_MoveMakesContact(u32 ability, u32 holdEffect, u32 move);
 bool32 ShouldUseZMove(u32 battlerAtk, u32 battlerDef, u32 chosenMove);
+u32 AI_GetBattlerAbility(u32 battler);
 
 // stat stage checks
 bool32 AnyStatIsRaised(u32 battlerId);
diff --git a/src/battle_ai_switch_items.c b/src/battle_ai_switch_items.c
index 7087a7dc20..24fafaf8bd 100644
--- a/src/battle_ai_switch_items.c
+++ b/src/battle_ai_switch_items.c
@@ -482,7 +482,7 @@ static bool32 ShouldSwitchIfTrapperInParty(u32 battler)
 
         monAbility = GetMonAbility(&party[i]);
 
-        if (CanAbilityTrapOpponent(monAbility, opposingBattler))
+        if (CanAbilityTrapOpponent(monAbility, opposingBattler) || (CanAbilityTrapOpponent(AI_GetBattlerAbility(opposingBattler), opposingBattler) && monAbility == ABILITY_TRACE))
         {
             // If mon in slot i is the most suitable switchin candidate, then it's a trapper than wins 1v1
             if (i == AI_DATA->mostSuitableMonId[battler])
@@ -1915,7 +1915,8 @@ static u32 GetBestMonIntegrated(struct Pokemon *party, int firstId, int lastId,
                 }
 
                 // If mon can trap
-                if (CanAbilityTrapOpponent(AI_DATA->switchinCandidate.battleMon.ability, opposingBattler)
+                if ((CanAbilityTrapOpponent(AI_DATA->switchinCandidate.battleMon.ability, opposingBattler)
+                    || (CanAbilityTrapOpponent(AI_GetBattlerAbility(opposingBattler), opposingBattler) && AI_DATA->switchinCandidate.battleMon.ability == ABILITY_TRACE))
                     && CountUsablePartyMons(opposingBattler) > 0
                     && canSwitchinWin1v1)
                     trapperId = i;
diff --git a/test/battle/ai/ai_switching.c b/test/battle/ai/ai_switching.c
index 26d111ba50..f86dd2b5f9 100644
--- a/test/battle/ai/ai_switching.c
+++ b/test/battle/ai/ai_switching.c
@@ -422,7 +422,7 @@ AI_SINGLE_BATTLE_TEST("AI_FLAG_SMART_MON_CHOICES: AI will switch in trapping mon
     u32 aiSmartMonChoicesFlag = 0; // Enables trapping behaviour after KOs
     PARAMETRIZE { aiSmartMonChoicesFlag = 0; } // No trapping behaviour
     PARAMETRIZE { aiSmartMonChoicesFlag = AI_FLAG_SMART_MON_CHOICES; } // Traps with mid battle switches
-    GIVEN{
+    GIVEN {
         ASSUME(gSpeciesInfo[SPECIES_MAWILE].types[0] == TYPE_STEEL);
         AI_FLAGS(AI_FLAG_CHECK_BAD_MOVE | AI_FLAG_CHECK_VIABILITY | AI_FLAG_TRY_TO_FAINT | aiSmartMonChoicesFlag);
         PLAYER(SPECIES_MAWILE) { Speed(2); Moves(MOVE_PROTECT, MOVE_TACKLE); }
@@ -443,7 +443,7 @@ AI_SINGLE_BATTLE_TEST("AI won't use trapping behaviour if player only has 1 mon
     u32 aiSmartMonChoicesFlag = 0; // Enables trapping behaviour after KOs
     PARAMETRIZE { aiSmartMonChoicesFlag = 0; } // No trapping behaviour
     PARAMETRIZE { aiSmartMonChoicesFlag = AI_FLAG_SMART_MON_CHOICES; } // Traps with mid battle switches
-    GIVEN{
+    GIVEN {
         ASSUME(gSpeciesInfo[SPECIES_MAWILE].types[0] == TYPE_STEEL);
         AI_FLAGS(AI_FLAG_CHECK_BAD_MOVE | AI_FLAG_CHECK_VIABILITY | AI_FLAG_TRY_TO_FAINT | aiSmartMonChoicesFlag);
         PLAYER(SPECIES_MAWILE) { Speed(2); Moves(MOVE_PROTECT, MOVE_TACKLE); }
@@ -451,7 +451,20 @@ AI_SINGLE_BATTLE_TEST("AI won't use trapping behaviour if player only has 1 mon
         OPPONENT(SPECIES_MAGNEZONE) { Speed(1); Ability(ABILITY_MAGNET_PULL); Moves(MOVE_SHOCK_WAVE); }
         OPPONENT(SPECIES_MEGANIUM) { Speed(3); Moves(MOVE_EARTH_POWER); }
     } WHEN {
-            TURN{ MOVE(player, MOVE_PROTECT); EXPECT_MOVE(opponent, MOVE_SELF_DESTRUCT); EXPECT_SEND_OUT(opponent, 2); }
+        TURN{ MOVE(player, MOVE_PROTECT); EXPECT_MOVE(opponent, MOVE_SELF_DESTRUCT); EXPECT_SEND_OUT(opponent, 2); }
+    }
+}
+
+AI_SINGLE_BATTLE_TEST("AI will trap player using Trace if player has a trapper")
+{
+    GIVEN {
+        AI_FLAGS(AI_FLAG_CHECK_BAD_MOVE | AI_FLAG_CHECK_VIABILITY | AI_FLAG_TRY_TO_FAINT | AI_FLAG_OMNISCIENT | AI_FLAG_SMART_SWITCHING);
+        PLAYER(SPECIES_DUGTRIO) { Ability(ABILITY_ARENA_TRAP); Moves(MOVE_ROCK_TOMB); }
+        PLAYER(SPECIES_DUGTRIO);
+        OPPONENT(SPECIES_GENGAR);
+        OPPONENT(SPECIES_PORYGON2) { Ability(ABILITY_TRACE); Item(ITEM_EVIOLITE); Moves(MOVE_ICE_BEAM); }
+    } WHEN {
+        TURN { MOVE(player, MOVE_ROCK_TOMB); EXPECT_SWITCH(opponent, 1); }
     }
 }