Refactoring lab

games/snake_v1/__init__.py

@@ -0,0 +1,92 @@
+from random import randint, choice
+from games.snake.apple import Apple
+from games.snake.snake import SnakeBodySegment
+from retro.game import Game
+
+
+class SnakeHead:
+    RIGHT = (1, 0)
+    UP = (0, -1)
+    LEFT = (-1, 0)
+    DOWN = (0, 1)
+    name = "Snake head"
+    position = (0, 0)
+    direction = DOWN
+    character = 'v'
+    next_segment = None
+    growing = False
+
+    def play_turn(self, game):
+        x, y = self.position
+        dx, dy = self.direction
+        next_pos = (x + dx, y + dy)
+        if self.can_move(next_pos, game):
+            apple = game.get_agent_by_name("Apple")
+            ax, ay = apple.position
+            old_dist = abs(x - ax) + abs(y - ay)
+            new_dist = abs(next_pos[0] - ax) + abs(next_pos[1] - ay)
+            game.state['reward'] += old_dist - new_dist
+            game.state['energy'] -= 1
+            self.position = next_pos
+            if self.is_on_apple(self.position, game):
+                apple.relocate(game)
+                self.growing = True
+                game.state['score'] += 50
+                game.state['reward'] += 50
+                game.state['energy'] = 150
+            if self.next_segment:
+                self.next_segment.move((x, y), game, growing=self.growing)
+            elif self.growing:
+                self.next_segment = SnakeBodySegment(1, (x, y))
+                game.add_agent(self.next_segment)
+            self.growing = False
+            if game.state['energy'] <= 0:
+                game.state['reward'] -= 10
+                game.end()
+        else:
+            game.state['reward'] -= 10
+            game.end()
+
+    def handle_keystroke(self, keystroke, game):
+        if keystroke.name == "KEY_RIGHT":
+            self.direction = self.RIGHT
+            self.character = '>'
+        elif keystroke.name == "KEY_UP":
+            self.direction = self.UP
+            self.character = '^'
+        elif keystroke.name == "KEY_LEFT":
+            self.direction = self.LEFT
+            self.character = '<'
+        elif keystroke.name == "KEY_DOWN":
+            self.direction = self.DOWN
+            self.character = 'v'
+
+    def can_move(self, position, game):
+        return game.on_board(position) and (game.is_empty(position) or self.is_on_apple(position, game))
+
+    def is_on_apple(self, position, game):
+        return game.get_agent_by_name("Apple").position == position
+
+
+def create_game():
+    head = SnakeHead()
+    apple = Apple()
+    game = Game(
+        [head, apple],
+        {'score': 0, 'reward': 0, 'energy': 150},
+        board_size=(32, 16),
+        framerate=12,
+        show_state=['score'],
+    )
+    bw, bh = game.board_size
+    head.position = (randint(1, bw - 2), randint(1, bh - 2))
+    direction, character = choice([
+        (SnakeHead.RIGHT, '>'),
+        (SnakeHead.UP,    '^'),
+        (SnakeHead.LEFT,  '<'),
+        (SnakeHead.DOWN,  'v'),
+    ])
+    head.direction = direction
+    head.character = character
+    apple.relocate(game)
+    return game