From 7dc4210f711e0cb92ac668b9a217561731b2e83c Mon Sep 17 00:00:00 2001
From: Daniel Pochert <dp0chdev@gmail.com>
Date: Sun, 17 Nov 2024 00:33:14 +0100
Subject: [PATCH] [Bug] fix training session encounter not applying nature
 correctly (#4801)

* fix training session encounter not applying nature correctly

* implement pokemon.setCustomNature + nature unlock utility method(s)

* update doc

---------

Co-authored-by: Moka <54149968+MokaStitcher@users.noreply.github.com>
---
 .../shady-vitamin-dealer-encounter.ts         |  2 +-
 .../encounters/training-session-encounter.ts  |  4 +--
 src/field/pokemon.ts                          |  5 ++++
 src/modifier/modifier.ts                      | 12 +++------
 src/system/game-data.ts                       | 26 +++++++++++++++++++
 5 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/src/data/mystery-encounters/encounters/shady-vitamin-dealer-encounter.ts b/src/data/mystery-encounters/encounters/shady-vitamin-dealer-encounter.ts
index 798e85868f6..3f1ace47b0f 100644
--- a/src/data/mystery-encounters/encounters/shady-vitamin-dealer-encounter.ts
+++ b/src/data/mystery-encounters/encounters/shady-vitamin-dealer-encounter.ts
@@ -138,7 +138,7 @@ export const ShadyVitaminDealerEncounter: MysteryEncounter =
             newNature = randSeedInt(25) as Nature;
           }
 
-          chosenPokemon.customPokemonData.nature = newNature;
+          chosenPokemon.setCustomNature(newNature);
           encounter.setDialogueToken("newNature", getNatureName(newNature));
           queueEncounterMessage(scene, `${namespace}:cheap_side_effects`);
           setEncounterExp(scene, [ chosenPokemon.id ], 100);
diff --git a/src/data/mystery-encounters/encounters/training-session-encounter.ts b/src/data/mystery-encounters/encounters/training-session-encounter.ts
index 4a0a91bf751..725c4ba79eb 100644
--- a/src/data/mystery-encounters/encounters/training-session-encounter.ts
+++ b/src/data/mystery-encounters/encounters/training-session-encounter.ts
@@ -226,8 +226,8 @@ export const TrainingSessionEncounter: MysteryEncounter =
           const onBeforeRewardsPhase = () => {
             queueEncounterMessage(scene, `${namespace}:option.2.finished`);
             // Add the pokemon back to party with Nature change
-            playerPokemon.setNature(encounter.misc.chosenNature);
-            scene.gameData.setPokemonCaught(playerPokemon, false);
+            playerPokemon.setCustomNature(encounter.misc.chosenNature);
+            scene.gameData.unlockSpeciesNature(playerPokemon.species, encounter.misc.chosenNature);
 
             // Add pokemon and modifiers back
             scene.getPlayerParty().push(playerPokemon);
diff --git a/src/field/pokemon.ts b/src/field/pokemon.ts
index 07a958fe27b..f9e7d7d1cad 100644
--- a/src/field/pokemon.ts
+++ b/src/field/pokemon.ts
@@ -1072,6 +1072,11 @@ export default abstract class Pokemon extends Phaser.GameObjects.Container {
     this.calculateStats();
   }
 
+  setCustomNature(nature: Nature): void {
+    this.customPokemonData.nature = nature;
+    this.calculateStats();
+  }
+
   generateNature(naturePool?: Nature[]): void {
     if (naturePool === undefined) {
       naturePool = Utils.getEnumValues(Nature);
diff --git a/src/modifier/modifier.ts b/src/modifier/modifier.ts
index 4e59d6c01f0..bbd1763cd84 100644
--- a/src/modifier/modifier.ts
+++ b/src/modifier/modifier.ts
@@ -1,5 +1,5 @@
 import type BattleScene from "#app/battle-scene";
-import { FusionSpeciesFormEvolution, pokemonEvolutions, pokemonPrevolutions } from "#app/data/balance/pokemon-evolutions";
+import { FusionSpeciesFormEvolution, pokemonEvolutions } from "#app/data/balance/pokemon-evolutions";
 import { getBerryEffectFunc, getBerryPredicate } from "#app/data/berry";
 import { getLevelTotalExp } from "#app/data/exp";
 import { allMoves } from "#app/data/move";
@@ -2197,14 +2197,8 @@ export class PokemonNatureChangeModifier extends ConsumablePokemonModifier {
    * @returns
    */
   override apply(playerPokemon: PlayerPokemon): boolean {
-    playerPokemon.customPokemonData.nature = this.nature;
-    let speciesId = playerPokemon.species.speciesId;
-    playerPokemon.scene.gameData.dexData[speciesId].natureAttr |= 1 << (this.nature + 1);
-
-    while (pokemonPrevolutions.hasOwnProperty(speciesId)) {
-      speciesId = pokemonPrevolutions[speciesId];
-      playerPokemon.scene.gameData.dexData[speciesId].natureAttr |= 1 << (this.nature + 1);
-    }
+    playerPokemon.setCustomNature(this.nature);
+    playerPokemon.scene.gameData.unlockSpeciesNature(playerPokemon.species, this.nature);
 
     return true;
   }
diff --git a/src/system/game-data.ts b/src/system/game-data.ts
index 70b0fadf0b0..3e3a6ce8f8b 100644
--- a/src/system/game-data.ts
+++ b/src/system/game-data.ts
@@ -1791,6 +1791,32 @@ export class GameData {
     });
   }
 
+  /**
+   * Checks whether the root species of a given {@PokemonSpecies} has been unlocked in the dex
+   */
+  isRootSpeciesUnlocked(species: PokemonSpecies): boolean {
+    return !!this.dexData[species.getRootSpeciesId()]?.caughtAttr;
+  }
+
+  /**
+   * Unlocks the given {@linkcode Nature} for a {@linkcode PokemonSpecies} and its prevolutions.
+   * Will fail silently if root species has not been unlocked
+   */
+  unlockSpeciesNature(species: PokemonSpecies, nature: Nature): void {
+    if (!this.isRootSpeciesUnlocked(species)) {
+      return;
+    }
+
+    //recursively unlock nature for species and prevolutions
+    const _unlockSpeciesNature = (speciesId: Species) => {
+      this.dexData[speciesId].natureAttr |= 1 << (nature + 1);
+      if (pokemonPrevolutions.hasOwnProperty(speciesId)) {
+        _unlockSpeciesNature(pokemonPrevolutions[speciesId]);
+      }
+    };
+    _unlockSpeciesNature(species.speciesId);
+  }
+
   updateSpeciesDexIvs(speciesId: Species, ivs: integer[]): void {
     let dexEntry: DexEntry;
     do {