feat(ai): add algorithm

src/ai/__init__.py

@@ -0,0 +1,3 @@
+from .ai import run
+
+__all__ = ["run"]

src/ai/ai.py

@@ -0,0 +1,35 @@
+from game import Main
+from game.sprites import Tetromino
+from loguru import logger
+from utils import BestMove, GameMode
+
+from .move import get_best_move
+
+
+def run() -> None:
+    app = Main(GameMode.AI_TRAINING)
+    game = app.game
+
+    if not game:
+        return
+
+    tetris = game.tetris
+
+    while True:
+        app.handle_events()
+        app.run_game_loop()
+
+        phantom_tetermino = Tetromino(tetris.phantom_sprites, None, tetris.field, tetris.tetromino.figure, True)
+        best_move: BestMove = get_best_move(app, game.tetris, phantom_tetermino)
+        figure = game.tetris.tetromino.figure
+        logger.debug(f"{figure.name=}    {best_move=}")
+        for rotation in range(best_move.rotation):
+            tetris.tetromino.rotate()
+
+        for _ in range(abs(best_move.x_axis_offset)):
+            if best_move.x_axis_offset > 0:
+                tetris.move_right()
+            elif best_move.x_axis_offset < 0:
+                tetris.move_left()
+
+        tetris.drop()

src/ai/move.py

@@ -0,0 +1,55 @@
+from typing import Optional
+
+from game import Main, Tetris
+from game.sprites import Tetromino
+from loguru import logger
+from utils import CONFIG, BestMove, Direction, Figure
+
+from .score import calculate_score
+
+NUM_ROTATIONS: dict[Figure, int] = {
+    Figure.I: 2,
+    Figure.O: 1,
+    Figure.T: 4,
+    Figure.S: 2,
+    Figure.Z: 2,
+    Figure.J: 4,
+    Figure.L: 4,
+}
+
+
+def get_best_move(app: Main, game: Tetris, tetermino: Tetromino) -> BestMove:
+    best_move: Optional[BestMove] = None
+    best_score: Optional[float] = None
+
+    for rotation in range(NUM_ROTATIONS[tetermino.figure]):
+        for i in range(CONFIG.game.columns):
+            x_axis_movement: int = 0
+            for _ in range(rotation):
+                tetermino.rotate()
+
+            while tetermino.move_horizontal(Direction.LEFT):  # move maximaly to the left
+                x_axis_movement -= 1
+
+            for _ in range(i):
+                if tetermino.move_horizontal(Direction.RIGHT):  # slowly move to the right
+                    x_axis_movement += 1
+
+            tetermino.drop()
+
+            score: float = calculate_score(game)
+
+            logger.debug(
+                f"{tetermino.figure.name=:3} {score=:6.6f} {best_score=:6.6f} {rotation=:1} {x_axis_movement=:1}"
+            )
+
+            tetermino.kill()
+            tetermino = Tetromino(game.phantom_sprites, None, game.field, game.tetromino.figure, True)
+
+            if best_score is None or score > best_score:
+                best_score = score
+                best_move = BestMove(rotation, x_axis_movement)
+
+    if not best_move:
+        best_move = BestMove(0, 0)
+    return best_move

src/ai/score.py

@@ -0,0 +1,17 @@
+import numpy as np
+from game import Tetris
+
+from .heuristics import aggregate_height, complete_lines, count_holes, get_bumpiness
+
+
+def calculate_score(game: Tetris) -> float:
+    field: np.ndarray[int, np.dtype[np.uint8]] = np.where(game.field != None, 1, 0)
+    for block in game.tetromino.blocks:
+        field[int(block.pos.y), int(block.pos.x)] = 1
+
+    height = aggregate_height(field) * -0.510066
+    lines = complete_lines(field) * 0.760666
+    holes = count_holes(field) * -0.35663
+    bumpiness = get_bumpiness(field) * -0.184483
+
+    return height + lines + holes + bumpiness

src/utils/__init__.py

@@ -3,7 +3,7 @@ from .enum import Direction, GameMode, Rotation
 from .figure import Figure, FigureConfig
 from .path import BASE_PATH
 from .settings import read_settings, save_settings
-from .tuples import Size
+from .tuples import BestMove, Size
 
 __all__ = [
     "BASE_PATH",
@@ -16,4 +16,5 @@ __all__ = [
     "GameMode",
     "read_settings",
     "save_settings",
+    "BestMove",
 ]

src/utils/tuples.py

@@ -17,3 +17,16 @@ class Size(NamedTuple):
         if isinstance(other, Size):
             return Size(self.width - other.width, self.height - other.height)
         return Size(self.width - other, self.height - other)
+
+
+class BestMove(NamedTuple):
+    """
+    A best move object.
+
+    Attributes:
+        rotation: The rotation of the best move.
+        x_axis_offset: The x-axis offset of the best move.
+    """
+
+    rotation: int
+    x_axis_offset: int

tetris/__main__.py

@@ -27,6 +27,13 @@ parser.add_argument(
     help="Run app with GUI [Default]",
 )
 
+parser.add_argument(
+    "-t",
+    "--train",
+    action="store_true",
+    help="Train AI",
+)
+
 
 def setup_logger(level: str = "warning") -> None:
     from utils import BASE_PATH
@@ -64,10 +71,14 @@ def main(args) -> None:
         level = "info"
     else:
         level = "warning"
-
     setup_logger(level)
 
-    run()
+    if args.train:  # type: ignore
+        import ai
+
+        ai.train()
+    else:
+        run()
 
 
 if __name__ == "__main__":