6 files changed, 221 insertions(+), 37 deletions(-)

M board.py
M logic.py
M state.py
A tests/test_advancement_2.py
M utils/print_game.py
M utils/test_helpers.py

@@ 1,11 1,11 @@
 from copy import deepcopy
-from typing import List, Union
+from typing import Tuple, Union, Dict, List
 from state import Snake, Pos
 
 # TODO: notes
-def get_snake_bodies_2(snakes: List[Snake], you: Union[Snake, None]=None) -> List[Pos]:
+def get_snake_bodies_2(snakes: Dict[str, Snake], you: Union[Tuple[str, Snake], None]=None) -> List[Pos]:
     bodies = []
-    for snake in snakes:
+    for snake in snakes.items():
         body = deepcopy(snake[1])
         if you != None and snake == you:
             bodies += body[1:]


@@ 66,3 66,37 @@ def get_next_pos(board, pos, direction, wrapped=True):
         new_pos = get_wrapped_pos(board, new_pos)
 
     return new_pos
+
+def get_wrapped_pos_2(board: Tuple[int, int], pos: Pos):
+    new_x = pos[0]
+    new_y = pos[1]
+    width = board[0]
+    height = board[1]
+
+    if new_x < 0: new_x = width - 1
+    if new_x >= width: new_x = 0
+    if new_y < 0: new_y = height - 1
+    if new_y >= height: new_y = 0
+
+    return (new_x, new_y)
+
+# FIXME: str not Dir due to circular imports
+def get_next_pos_2(board: Tuple[int, int], pos: Pos, direction: str, wrapped=True):
+    new_pos = deepcopy(pos)
+
+    if direction == "up":
+        new_pos = (pos[0], pos[1]+1)
+
+    if direction == "down":
+        new_pos = (pos[0], pos[1]-1)
+
+    if direction == "left":
+        new_pos = (pos[0]-1, pos[1])
+
+    if direction == "right":
+        new_pos = (pos[0]+1, pos[1])
+
+    if wrapped:
+        new_pos = get_wrapped_pos_2(board, new_pos)
+
+    return new_pos

@@ 1,51 1,56 @@
 from copy import deepcopy
 import random
 import itertools
-from typing import Literal, Tuple, List
+from typing import Literal, Dict
 import bitboard
 from state import State
-from board import is_snake_in_board, get_snake_bodies, get_snake_bodies_2, get_next_pos
+from board import is_snake_in_board, get_snake_bodies, get_snake_bodies_2, get_next_pos, get_next_pos_2
 
 Dir = Literal["up", "down", "left", "right"]
-Move = Tuple[str, Dir]
 
-# TODO: advance_state that handles new state format
-def advance_state_2(prev_state: State, moves: List[Move]) -> State:
-    snakes = deepcopy(prev_state["snakes"]) # TODO: verify this copy is correct
-    food = deepcopy(prev_state["food"]) # TODO: verify this copy is correct
+# TODO: handle head-to-head interactions
+def advance_state_2(prev_state: State, moves: Dict[str, Dir]) -> State:
+    snakes = deepcopy(prev_state["snakes"])
+    food = deepcopy(prev_state["food"])
 
     # move snakes
-    for [snake_id, body] in snakes:
-        move = None
-        for [move_snake_id, snake_move] in moves:
-            if move_snake_id == snake_id:
-                move = snake_move
-                break
-        head = get_next_pos(prev_state["board"], body[0], move)
-        body = [head] + body # NOTE: might be a more efficient way to do this
+    for (snake_id, body) in snakes.items():
+        move = moves[snake_id]
+        head = get_next_pos_2(prev_state["board"], body[0], move)
+        body = [head] + body
 
     # kill off snakes that run into each other
-    # TODO: handle head-to-head interactions
-    alive_snakes = []
-    for snake in snakes:
-        if snake[1][0] not in get_snake_bodies_2(snakes, snake):
-            alive_snakes.append(snake)
-
+    alive_snakes = {}
+    for s in snakes.items():
+        if s[1][0] not in get_snake_bodies_2(snakes, s):
+            alive_snakes[s[0]] = s[1]
+
+    # remove consumed food, keep track of who's eaten
+    remaining_food = []
+    fed_snakes = []
+    for f in food:
+        consumed = False
+        for (snake_id, body) in snakes.items():
+            if body[0] == f:
+                consumed = True
+                fed_snakes.append(snake_id)
+        if not consumed:
+            remaining_food.append(f)
 
-    # handle food / growth interactions
-    # loop through food
-    # if food is at head of snake, remove snake and keep snake at length
-    # if snake hasn't eaten, remove tail
+    # if a snake hasn't eaten, remove it's tail
+    for (snake_id, body) in snakes.items():
+        if snake_id not in fed_snakes:
+            body = body[:-1]
 
     return {
         "turn": prev_state["turn"] + 1,
         "board": prev_state["board"],
         "hazards": prev_state["hazards"],
-        "food": [], # TODO
+        "food": remaining_food,
         "snakes": alive_snakes,
     }
 
-def advance_state(state, moves: List[Move]):
+def advance_state(state, moves):
     new_state = deepcopy(state) # gross
     board = new_state["board"]
 

@@ 1,7 1,7 @@
 from typing import Dict, List, Tuple, TypedDict
 
 Pos = Tuple[int, int]
-Snake = Tuple[str, List[Pos]] # TODO: health?
+Snake = List[Pos] # TODO: health?
 
 class State(TypedDict):
     """


@@ 12,7 12,7 @@ class State(TypedDict):
     board: Tuple[int, int]
     food: List[Pos]
     hazards: List[Pos]
-    snakes: List[Snake]
+    snakes: Dict[str, Snake]
 
 def dict_to_tuple(pos: Dict[str, int]) -> Pos:
     """


@@ 31,11 31,8 @@ def parse_state(state) -> State:
         "board": (state["board"]["width"], state["board"]["height"]),
         "food": [dict_to_tuple(pos) for pos in state["board"]["food"]],
         "hazards": [dict_to_tuple(pos) for pos in state["board"]["hazards"]],
-        "snakes": [parse_snake(snake) for snake in state["board"]["snakes"]],
+        "snakes": {s["id"]:parse_snake(s) for s in state["board"]["snakes"]},
     }
 
 def parse_snake(snake) -> Snake:
-    return (
-        snake["id"],
-        [dict_to_tuple(pos) for pos in snake["body"]]
-    )
+    return [dict_to_tuple(pos) for pos in snake["body"]]

@@ 0,0 1,50 @@
+from typing import Dict
+import unittest
+from logic import advance_state_2, Dir
+from utils.test_helpers import generate_state_2
+from utils.print_game import print_to_screen_2
+
+class TestAdvancement2(unittest.TestCase):
+    def test_advance_bodies(self):
+        """
+        Testing that the snake head AND bodies are moved into the correct place
+        after several moves. The hypothesis is that the neck position isn't
+        moving correctly.
+        """
+        state_1 = generate_state_2([
+            ["_", "_", "_", "_"],
+            ["_", "_", "_", "_"],
+            ["a", "a", "A", "_"],
+            ["_", "_", "_", "_"],
+            ["_", "_", "_", "_"],
+        ])
+        print_to_screen_2(state_1, True)
+        snake_a = state_1["snakes"]["snake-a"]
+        self.assertEqual(snake_a[0], (2, 2))
+        self.assertIn((2, 2), snake_a)
+        self.assertIn((1, 2), snake_a)
+        self.assertIn((0, 2), snake_a)
+
+        # FIXME: it doesn't look like we're advancing
+        # state_2 = generate_state([
+        #     ["_", "_", "_", "_"],
+        #     ["_", "_", "A", "_"],
+        #     ["_", "a", "a", "_"],
+        #     ["_", "_", "_", "_"],
+        #     ["_", "_", "_", "_"],
+        # ])
+        # update the body segment order to match what the API would give (in
+        # order from head to tail
+        snake_a[0] = (2, 2)
+        snake_a[1] = (1, 2)
+        snake_a[2] = (0, 2)
+        moves_1: Dict[str, Dir] = {"snake-a": "up"}
+        state_2 = advance_state_2(state_1, moves_1)
+        snake_a = state_2["snakes"]["snake-a"]
+        self.assertEqual(snake_a[0], (2, 3))
+        self.assertIn((2, 3), snake_a)
+        self.assertIn((2, 2), snake_a)
+        self.assertIn((1, 2), snake_a)
+
+if __name__ == "__main__":
+    unittest.main()

@@ 1,4 1,60 @@
 from board import get_snake_bodies
+from state import State
+
+def print_to_screen_2(state: State, numbers=False, turn=False):
+    print(get_state_str_2(state, numbers, turn))
+
+def get_state_str_2(state: State, numbers: bool, turn: bool) -> str:
+    results = ""
+    snakes = [s[1] for s in state["snakes"].items()]
+
+    if turn:
+        results += f"===== turn: {state['turn']} =====\n"
+
+    for vert in range(state["board"][1]):
+        y = state["board"][1] - vert - 1
+        if numbers:
+            if y < 10:
+                results += f"{y} "
+            else:
+                results += f"{y}"
+        for x in range(state["board"][0]):
+            pos = (x, y)
+
+            if pos in state["hazards"]:
+                results += "x "
+            elif pos in state["food"]:
+                results += "o "
+
+            printed = False
+            for (i, s) in enumerate(snakes):
+                if i == 0:
+                    if pos == s[0]:
+                        results += "Y "
+                        printed = True
+                    elif pos in s:
+                        results += "b "
+                        printed = True
+                else:
+                    if pos == s[0]:
+                        results += "H "
+                        printed = True
+                    elif pos in s:
+                        results += "b "
+                        printed = True
+            if not printed:
+                results += "_ "
+
+        results += "\n"
+
+    if numbers:
+        results += f"  "
+        for x in range(state["board"][0]):
+            results += f"{x} "
+
+        results += "\n"
+
+    return results
 
 def print_to_screen(state, numbers=False, turn=False):
     print(get_state_str(state, numbers, turn))

@@ 1,4 1,46 @@
 from typing import List
+from state import State
+
+def generate_state_2(test_input: List[List[str]]) -> State:
+    height = len(test_input)
+    width = len(test_input[1])
+    food = []
+    hazards = []
+    snakes = { "snake-a": [], "snake-b": [] }
+    snake_a_head = None
+    snake_b_head = None
+
+    for vert, row in enumerate(test_input):
+        y = height - vert - 1
+        for x, cell in enumerate(row):
+            pos = (x, y)
+            if cell == "o":
+                food.append(pos)
+            if cell == "x":
+                hazards.append(pos)
+            if cell == "A":
+                snake_a_head = pos
+            if cell == "a":
+                snakes["snake-a"].append(pos)
+            if cell == "B":
+                snake_b_head = pos
+            if cell == "b":
+                snakes["snake-b"].append(pos)
+
+    snakes["snake-a"] = [snake_a_head] + snakes["snake-a"]
+
+    if snake_b_head != None:
+        snakes["snake-b"] = [snake_b_head] + snakes["snake-b"]
+    else:
+        del snakes["snake-b"]
+
+    return {
+        "turn": 1,
+        "board": (width, height),
+        "food": food,
+        "hazards": hazards,
+        "snakes": snakes
+    }
 
 def generate_state(test_input: List[List[str]]):
     board = generate_board(test_input)