From b5a71858412e9d7937aecc381d24587ef6b8e09d Mon Sep 17 00:00:00 2001
From: Pawkkie <61265402+Pawkkie@users.noreply.github.com>
Date: Mon, 15 Jul 2024 03:19:41 -0400
Subject: [PATCH] Fix U-Turn With Ace Mon (#4748) (#4972)

* Fix U-Turn With Ace Mon

* ASSUME for test
---
 src/battle_ai_switch_items.c |  9 ++++-----
 test/battle/ai.c             | 13 +++++++++++++
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/src/battle_ai_switch_items.c b/src/battle_ai_switch_items.c
index c50057d7a5..50aef125e0 100644
--- a/src/battle_ai_switch_items.c
+++ b/src/battle_ai_switch_items.c
@@ -1901,12 +1901,11 @@ static u32 GetBestMonIntegrated(struct Pokemon *party, int firstId, int lastId,
 
         else if (batonPassId != PARTY_SIZE)
             return batonPassId;
-
-        // If ace mon is the last available Pokemon and U-Turn/Volt Switch was used - switch to the mon.
-        else if (aceMonId != PARTY_SIZE
-            && (gMovesInfo[gLastUsedMove].effect == EFFECT_HIT_ESCAPE || gMovesInfo[gLastUsedMove].effect == EFFECT_PARTING_SHOT))
-            return aceMonId;
     }
+    // If ace mon is the last available Pokemon and U-Turn/Volt Switch was used - switch to the mon.
+    if (aceMonId != PARTY_SIZE && (gMovesInfo[gLastUsedMove].effect == EFFECT_HIT_ESCAPE || gMovesInfo[gLastUsedMove].effect == EFFECT_PARTING_SHOT))
+        return aceMonId;
+        
     return PARTY_SIZE;
 }
 
diff --git a/test/battle/ai.c b/test/battle/ai.c
index e98cb7b08e..578ca345ff 100644
--- a/test/battle/ai.c
+++ b/test/battle/ai.c
@@ -699,6 +699,19 @@ AI_SINGLE_BATTLE_TEST("AI_FLAG_SMART_MON_CHOICES: Post-KO switches prioritize of
     }
 }
 
+AI_SINGLE_BATTLE_TEST("AI_FLAG_SMART_MON_CHOICES: U-Turn will send out Ace Mon if it's the only one remaining")
+{
+    GIVEN {
+        ASSUME(gMovesInfo[MOVE_U_TURN].effect == EFFECT_HIT_ESCAPE);
+        AI_FLAGS(AI_FLAG_CHECK_BAD_MOVE | AI_FLAG_CHECK_VIABILITY | AI_FLAG_TRY_TO_FAINT | AI_FLAG_SMART_MON_CHOICES | AI_FLAG_ACE_POKEMON);
+        PLAYER(SPECIES_WOBBUFFET);
+        OPPONENT(SPECIES_WOBBUFFET) { Moves(MOVE_U_TURN); }
+        OPPONENT(SPECIES_WOBBUFFET);
+    } WHEN {
+        TURN { EXPECT_MOVE(opponent, MOVE_U_TURN); EXPECT_SEND_OUT(opponent, 1); }
+    }
+}
+
 AI_SINGLE_BATTLE_TEST("AI_FLAG_SMART_SWITCHING: AI switches out after sufficient stat drops")
 {
     GIVEN {