[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>

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);

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);

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);

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;
   }

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 {