Read one memory chunk per cycle and add method to calculate building …

…related stats

src/stats/building.py

@@ -1,5 +1,6 @@
 """Script containing code to manage building related tasks and calculations."""
 
+import numpy as np
 import pandas as pd
 
 from src import PROCESS_NAME
@@ -49,48 +50,88 @@ def list_buildings(self, player_id: int = 0) -> pd.DataFrame:
             pd.DataFrame: buildings data
         """
         num_buildings = int(read_memory(PROCESS_NAME, self.total_buildings, D_Types.INT))
-        buildings_info = [
-            read_memory_chunk(
-                PROCESS_NAME,
-                self.base + i * self.offset,
-                [0, self.owner, self.workers_needed, self.workers, self.workers_missing, self.snoozed],
-            )
-            for i in range(num_buildings)
-        ]
+        offset_list = [0, self.owner, self.workers_needed, self.workers, self.workers_missing, self.snoozed]
+        buildings_list = read_memory_chunk(
+            PROCESS_NAME,
+            self.base,
+            [i * self.offset + extra_off for i in range(num_buildings) for extra_off in offset_list],
+        )
+        buildings_array = np.array(buildings_list).reshape((num_buildings, len(offset_list)))
         if player_id != 0:
-            buildings_info = [
-                [self.building_names.get(sublist[0])] + sublist
-                for sublist in filter(lambda x: x[1] == player_id, buildings_info)
-            ]
+            mask = buildings_array[:, 2] == player_id
         else:
-            buildings_info = [[self.building_names.get(sublist[0])] + sublist for sublist in buildings_info]
+            mask = (buildings_array[:, 2] >= 0) & (buildings_array[:, 2] <= 8)
+
+        filtered_buildings = buildings_array[mask]
+
+        building_names_array = np.vectorize(self.building_names.get)(filtered_buildings[:, 0])
+        address_array = (self.base + np.arange(buildings_array.shape[0]) * self.offset).reshape(-1, 1)
+        buildings_array = np.column_stack((address_array, building_names_array, filtered_buildings))
         return pd.DataFrame(
-            buildings_info,
+            buildings_array,
             columns=[
-                "b_name",
                 "address",
+                "b_name",
                 "ID",
                 "owner",
                 "workers_needed",
                 "workers_working",
                 "workers_missing",
                 "snoozed",
             ],
+        ).astype(
+            {
+                "address": pd.Int64Dtype(),
+                "b_name": pd.StringDtype(),
+                "ID": pd.Int64Dtype(),
+                "owner": pd.Int16Dtype(),
+                "workers_needed": pd.Int16Dtype(),
+                "workers_working": pd.Int16Dtype(),
+                "workers_missing": pd.Int16Dtype(),
+                "snoozed": pd.Int16Dtype(),
+            }
         )
 
-    def calc_worker_stats(self, player_id: int = 0) -> pd.Series:
-        """Calculate the worker stats through the buildings.
-
-        Args:
-            player_id (int, optional): player id to filter. Defaults to 0 and calculating global stats.
+    def calculate_all_stats(self) -> pd.DataFrame:
+        """Calculate all building and worker related stats into a dataframe.
 
         Returns:
-            pd.Series: series containing workers_needed, workers_working and workers_missing
+            pd.DataFrame: All building stats.
         """
-        buildings_df = self.list_buildings(player_id)
-        # filter hovels, houses, quarrypiles and snoozed building to obtain correct results
-        sums = buildings_df.loc[
-            ~(buildings_df["ID"].isin([1, 2, 21])) | (buildings_df["snoozed"] == 1),
-            ["workers_needed", "workers_working", "workers_missing"],
-        ].sum()
-        return sums
+        false_worker_ids = [1, 2, 8, 9, 21, 29]
+        ground_ids = [53, 55, 56, 57, 58, 59]
+        keep_ids = [71, 72, 73]
+        siege_engines = [80, 81, 82, 83, 84, 86, 87]
+        building_info_df = pd.DataFrame(
+            columns=["num_buildings", "workers_needed", "workers_working", "workers_missing", "snoozed"]
+        )
+        building_mem_df = self.list_buildings()
+        building_mem_df = building_mem_df.loc[
+            ~(building_mem_df["ID"].isin(ground_ids + keep_ids + siege_engines)),
+            :,
+        ]
+        building_info_df["num_buildings"] = (
+            building_mem_df.loc[:, ["owner"]]
+            .groupby("owner")
+            .size()
+            .to_frame()
+            .rename(columns={"size": "num_buildings"})
+        )
+        building_info_df["snoozed"] = (
+            building_mem_df.loc[building_mem_df["snoozed"] == 1, ["owner"]]
+            .groupby("owner")
+            .size()
+            .to_frame()
+            .rename(columns={"size": "snoozed"})
+        )
+        building_info_df["snoozed"] = building_info_df["snoozed"].fillna(0).astype(pd.Int32Dtype())
+        building_mem_df = building_mem_df.loc[
+            ~(building_mem_df["ID"].isin(false_worker_ids)) & (building_mem_df["snoozed"] == 0),
+            :,
+        ]
+        building_info_df[["workers_needed", "workers_working", "workers_missing"]] = (
+            building_mem_df.loc[:, ["owner", "workers_needed", "workers_working", "workers_missing"]]
+            .groupby("owner")
+            .sum()
+        )
+        return building_info_df.reset_index(names=["p_ID"])