From dba28cf93f623064e986267de745c196f637efc7 Mon Sep 17 00:00:00 2001
From: kittenchilly <mario.pikmin.zelda@gmail.com>
Date: Mon, 14 Aug 2023 19:10:15 -0400
Subject: [PATCH] Allow Z-Moves to be selected after Mega Evolution (#3233)

* Allow Z-Moves to be selected after Mega Evolution

* You cannot do Wish Mega Evolution at all if you hold a z crystal

* Rebase to master
---
 src/battle_util.c   | 26 +++++++++++++++-----------
 src/battle_z_move.c | 10 ----------
 2 files changed, 15 insertions(+), 21 deletions(-)

diff --git a/src/battle_util.c b/src/battle_util.c
index b5fb8732d..91df461f0 100644
--- a/src/battle_util.c
+++ b/src/battle_util.c
@@ -10076,19 +10076,19 @@ bool32 CanMegaEvolve(u8 battlerId)
     species = GetMonData(mon, MON_DATA_SPECIES);
     itemId = GetMonData(mon, MON_DATA_HELD_ITEM);
 
+#if DEBUG_BATTLE_MENU == TRUE
+    if (gBattleStruct->debugHoldEffects[battlerId])
+        holdEffect = gBattleStruct->debugHoldEffects[battlerId];
+    else
+#endif
+    if (itemId == ITEM_ENIGMA_BERRY_E_READER)
+        holdEffect = gEnigmaBerries[battlerId].holdEffect;
+    else
+        holdEffect = ItemId_GetHoldEffect(itemId);
+
     // Check if there is an entry in the evolution table for regular Mega Evolution.
     if (GetBattleFormChangeTargetSpecies(battlerId, FORM_CHANGE_BATTLE_MEGA_EVOLUTION_ITEM) != SPECIES_NONE)
     {
-    #if DEBUG_BATTLE_MENU == TRUE
-        if (gBattleStruct->debugHoldEffects[battlerId])
-            holdEffect = gBattleStruct->debugHoldEffects[battlerId];
-        else
-    #endif
-        if (itemId == ITEM_ENIGMA_BERRY_E_READER)
-            holdEffect = gEnigmaBerries[battlerId].holdEffect;
-        else
-            holdEffect = ItemId_GetHoldEffect(itemId);
-
         // Can Mega Evolve via Mega Stone.
         if (holdEffect == HOLD_EFFECT_MEGA_STONE)
             return TRUE;
@@ -10096,7 +10096,11 @@ bool32 CanMegaEvolve(u8 battlerId)
 
     // Check if there is an entry in the evolution table for Wish Mega Evolution.
     if (GetBattleFormChangeTargetSpecies(battlerId, FORM_CHANGE_BATTLE_MEGA_EVOLUTION_MOVE) != SPECIES_NONE)
-        return TRUE;
+    {
+        // Can't Wish Mega Evolve if holding a Z Crystal.
+        if (holdEffect != HOLD_EFFECT_Z_CRYSTAL)
+            return TRUE;
+    }
 
     // No checks passed, the mon CAN'T mega evolve.
     return FALSE;
diff --git a/src/battle_z_move.c b/src/battle_z_move.c
index dec3701c0..a24ad8f38 100644
--- a/src/battle_z_move.c
+++ b/src/battle_z_move.c
@@ -159,7 +159,6 @@ void QueueZMove(u8 battlerId, u16 baseMove)
 bool32 IsViableZMove(u8 battlerId, u16 move)
 {
     struct Pokemon *mon;
-    struct MegaEvolutionData *mega = &(((struct ChooseMoveStruct *)(&gBattleResources->bufferA[gActiveBattler][4]))->mega);
     u8 battlerPosition = GetBattlerPosition(battlerId);
     u8 partnerPosition = GetBattlerPosition(BATTLE_PARTNER(battlerId));
     u32 item;
@@ -185,15 +184,6 @@ bool32 IsViableZMove(u8 battlerId, u16 move)
     if ((GetBattlerPosition(battlerId) == B_POSITION_PLAYER_LEFT || (!(gBattleTypeFlags & BATTLE_TYPE_MULTI) && GetBattlerPosition(battlerId) == B_POSITION_PLAYER_RIGHT)) && !CheckBagHasItem(ITEM_Z_POWER_RING, 1))
         return FALSE;
 
-    if (mega->alreadyEvolved[battlerPosition])
-        return FALSE;   // Trainer has mega evolved
-
-    if (gBattleTypeFlags & BATTLE_TYPE_DOUBLE)
-    {
-        if (IsPartnerMonFromSameTrainer(battlerId) && (mega->alreadyEvolved[partnerPosition] || (mega->toEvolve & gBitTable[BATTLE_PARTNER(battlerId)])))
-            return FALSE;   // Partner has mega evolved or is about to mega evolve
-    }
-
 #if DEBUG_BATTLE_MENU == TRUE
     if (gBattleStruct->debugHoldEffects[battlerId])
         holdEffect = gBattleStruct->debugHoldEffects[battlerId];