combat!

1 files changed, 42 insertions, 0 deletions

diff --git a/robots/game.py b/robots/game.py
index 827562a..cb065ec 100644
--- a/robots/game.py
+++ b/robots/game.py
@@ -168,6 +168,47 @@ class Game:
                 for x, y, energy in robots
             ]
 
+    def resolve_combats(self):
+        to_remove = set()
+        to_append = defaultdict(list)
+
+        for (x, y), robots in self.state.robots.items():
+            if len(robots) <= 1:
+                # No combat here.
+                continue
+
+            to_remove.add((x, y))
+
+            # each player's strength is the sum of their robots' energy
+            strength_by_player = defaultdict(int)
+            total = 0
+            for player, energy in robots:
+                strength_by_player[player] += energy
+                total += energy
+
+            winner_count = 0
+            for player, strength in strength_by_player.items():
+                # subtract the sum of the other players' strength
+                strength -= total - strength
+                # remove with strength players <= 0
+                if strength > 0:
+                    energy = min(strength, Energy.MAXIMUM)
+                    to_append[player].append((x, y, energy))
+                    winner_count += 1
+
+            # we should end up with one player.
+            # (and I wrote a proof of that, so something's seriously wrong if
+            # this fails)
+            assert winner_count <= 1
+
+        for player, robots in self.state.robots_by_player.items():
+            extra_robots = to_append[player]
+            robots[:] = [
+                (x, y, energy)
+                for x, y, energy in robots
+                if (x, y) not in to_remove
+            ] + extra_robots
+
     def __iter__(self):
         while not self.finished:
             yield self.state
@@ -182,6 +223,7 @@ class Game:
         # 3. Factories produce new robots.
         self.factories_produce_robots()
         # 4. Combats are resolved.
+        self.resolve_combats()
         # 5. Each robot gains energy, up to the maximum.
         self.robots_gain_energy()