From 845aa61f7e867adc770d88d39b60e06404a1bd1b Mon Sep 17 00:00:00 2001
From: Adam Amer <136176500+adamamer20@users.noreply.github.com>
Date: Fri, 20 Sep 2024 11:44:09 +0200
Subject: [PATCH] fix: occupied cells sorted by agent_order, assertion to
 verify best moves aren't duplicated

---
 examples/sugarscape_ig/ss_polars/agents.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/examples/sugarscape_ig/ss_polars/agents.py b/examples/sugarscape_ig/ss_polars/agents.py
index 46dffe8..c0b6840 100644
--- a/examples/sugarscape_ig/ss_polars/agents.py
+++ b/examples/sugarscape_ig/ss_polars/agents.py
@@ -200,16 +200,20 @@ def get_best_moves(self, neighborhood: pl.DataFrame, agent_order):
             )
             .drop("agent_order")
         )
+assert best_moves.n_unique() == len(
+            best_moves
+        ), "Duplicates found in best_moves"
         return best_moves
 
     def _prepare_cells(self, neighborhood: pl.DataFrame):
         occupied_cells = (
-            neighborhood[["agent_id_center"]]
+            neighborhood[["agent_id_center", "agent_order"]]
             .unique()
             .join(self.pos, left_on="agent_id_center", right_on="unique_id")
             .with_columns(
                 flattened=(pl.col("dim_0") * self.space.dimensions[1] + pl.col("dim_1"))
-            )["flattened"]
+            ).sort("agent_order")
+["flattened"]
             .to_numpy()
         )
         free_cells = np.ones(