From ee8207c60a7156a1f764a7f149432b9f16988621 Mon Sep 17 00:00:00 2001
From: thatguy11325 <148832074+thatguy11325@users.noreply.github.com>
Date: Fri, 28 Jun 2024 22:14:32 -0400
Subject: [PATCH] Add Party struct

---
 pokemonred_puffer/data/party.py  | 52 ++++++++++++++++++++++++++++++++
 pokemonred_puffer/environment.py |  2 ++
 2 files changed, 54 insertions(+)
 create mode 100644 pokemonred_puffer/data/party.py

diff --git a/pokemonred_puffer/data/party.py b/pokemonred_puffer/data/party.py
new file mode 100644
index 0000000..306501a
--- /dev/null
+++ b/pokemonred_puffer/data/party.py
@@ -0,0 +1,52 @@
+from ctypes import Structure, Union, c_uint16, c_uint8, sizeof
+
+from pyboy import PyBoy
+
+
+class BoxStruct(Structure):
+    _fields_ = [
+        ("Species", c_uint8),
+        ("HP", c_uint16),
+        ("BoxLevel", c_uint8),
+        ("Status", c_uint8),
+        ("Type1", c_uint8),
+        ("Type2", c_uint8),
+        ("CatchRate", c_uint8),
+        ("Moves", 4 * c_uint8),
+        ("OTID", c_uint16),
+        ("Exp", 3 * c_uint8),
+        ("HPExp", c_uint16),
+        ("AttackExp", c_uint16),
+        ("DefenseExp", c_uint16),
+        ("SpeedExp", c_uint16),
+        ("SpecialExp", c_uint16),
+        ("DVs", 2 * c_uint8),
+        ("PP", 4 * c_uint8),
+    ]
+
+
+class PartyStruct(Structure):
+    _fields_ = BoxStruct._fields_ + [
+        ("Level", c_uint8),
+        ("MaxHP", c_uint16),
+        ("Attack", c_uint16),
+        ("Defense", c_uint16),
+        ("Speed", c_uint16),
+        ("Special", c_uint16),
+    ]
+
+
+PARTY_LENGTH_BYTES = 6 * sizeof(PartyStruct)
+
+
+class PartyMons(Union):
+    _fields_ = [("party", 6 * PartyStruct), ("asbytes", c_uint8 * PARTY_LENGTH_BYTES)]
+
+    def __init__(self, emu: PyBoy):
+        _, wPartyMons = emu.symbol_lookup("wPartyMons")
+        self.asbytes = (c_uint8 * PARTY_LENGTH_BYTES)(
+            *emu.memory[wPartyMons : wPartyMons + PARTY_LENGTH_BYTES]
+        )
+
+    def __getitem__(self, idx):
+        return self.party[idx]
diff --git a/pokemonred_puffer/environment.py b/pokemonred_puffer/environment.py
index edb2156..02738cb 100644
--- a/pokemonred_puffer/environment.py
+++ b/pokemonred_puffer/environment.py
@@ -33,6 +33,7 @@
 )
 from pokemonred_puffer.data.map import MapIds
 from pokemonred_puffer.data.missable_objects import MissableFlags
+from pokemonred_puffer.data.party import PartyMons
 from pokemonred_puffer.data.strength_puzzles import STRENGTH_SOLUTIONS
 from pokemonred_puffer.data.tilesets import Tilesets
 from pokemonred_puffer.data.tm_hm import (
@@ -564,6 +565,7 @@ def step(self, action):
         self.run_action_on_emulator(action)
         self.events = EventFlags(self.pyboy)
         self.missables = MissableFlags(self.pyboy)
+        self.party = PartyMons(self.pyboy)
         self.update_seen_coords()
         self.update_health()
         self.update_pokedex()