diff --git a/pyproject.toml b/pyproject.toml
index 6736271..43f76c6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -69,7 +69,7 @@ extend-select = [
     "TID",
     "YTT",
 ]
-ignore = ["E741"]
+ignore = ["E741", "E711"]
 show-fixes = true
 line-length = 120
 indent-width = 4
diff --git a/src/ai/ai.py b/src/ai/ai.py
index e1fc94d..8d7446d 100644
--- a/src/ai/ai.py
+++ b/src/ai/ai.py
@@ -1,5 +1,4 @@
 from game import Main
-from game.sprites import Tetromino
 from loguru import logger
 from utils import BestMove, GameMode
 
@@ -8,6 +7,7 @@ from .move import get_best_move
 
 def run() -> None:
     app = Main(GameMode.AI_TRAINING)
+    app.play()
     game = app.game
 
     if not game:
@@ -19,10 +19,10 @@ def run() -> None:
         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)
+        best_move: BestMove = get_best_move(game.tetris, tetris.tetromino.figure)
         figure = game.tetris.tetromino.figure
-        logger.debug(f"{figure.name=}    {best_move=}")
+        logger.warning(f"{figure.name=}    {best_move=}")
+
         for rotation in range(best_move.rotation):
             tetris.tetromino.rotate()
 
diff --git a/src/ai/move.py b/src/ai/move.py
index f98ce37..bd7582d 100644
--- a/src/ai/move.py
+++ b/src/ai/move.py
@@ -1,6 +1,8 @@
 from typing import Optional
 
-from game import Main, Tetris
+import numpy as np
+import pygame
+from game import Tetris
 from game.sprites import Tetromino
 from loguru import logger
 from utils import CONFIG, BestMove, Direction, Figure
@@ -18,12 +20,14 @@ NUM_ROTATIONS: dict[Figure, int] = {
 }
 
 
-def get_best_move(app: Main, game: Tetris, tetermino: Tetromino) -> BestMove:
+def get_best_move(game: Tetris, figure: Figure) -> BestMove:
     best_move: Optional[BestMove] = None
     best_score: Optional[float] = None
+    phantom_sprites = pygame.sprite.Group()  # type: ignore
 
-    for rotation in range(NUM_ROTATIONS[tetermino.figure]):
+    for rotation in range(NUM_ROTATIONS[figure]):
         for i in range(CONFIG.game.columns):
+            tetermino = Tetromino(phantom_sprites, None, game.field, game.tetromino.figure, True)
             x_axis_movement: int = 0
             for _ in range(rotation):
                 tetermino.rotate()
@@ -39,17 +43,25 @@ def get_best_move(app: Main, game: Tetris, tetermino: Tetromino) -> BestMove:
 
             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)
+            logger.debug(f"{tetermino.figure.name=:3} {score=:6.6f} {best_score=} {rotation=:1} {x_axis_movement=:1}")
 
             if best_score is None or score > best_score:
                 best_score = score
                 best_move = BestMove(rotation, x_axis_movement)
 
+            if not tetermino._are_new_positions_valid(
+                [pygame.Vector2(block.pos.x + 1, block.pos.y) for block in tetermino.blocks]
+            ):
+                break
+
+            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
+
+            logger.debug(f"{field=}")
+            tetermino.kill()
+
     if not best_move:
         best_move = BestMove(0, 0)
+    phantom_sprites.empty()
     return best_move
diff --git a/tetris/__main__.py b/tetris/__main__.py
index a42bb17..306ecbe 100755
--- a/tetris/__main__.py
+++ b/tetris/__main__.py
@@ -76,7 +76,7 @@ def main(args) -> None:
     if args.train:  # type: ignore
         import ai
 
-        ai.train()
+        ai.run()
     else:
         run()