diff --git a/src/machine_learning/helpers.py b/src/machine_learning/helpers.py
new file mode 100644
index 0000000..b7ab0df
--- /dev/null
+++ b/src/machine_learning/helpers.py
@@ -0,0 +1,12 @@
+from tile import Tile
+
+
+def get_dataset_from_tile(tile: Tile):
+    return {
+        'visibility': tile.visibility,
+        'stability': tile.stability,
+        'armed': tile.mine.armed,
+        'ground': tile.ground,
+        'mine_type': tile.mine.mine_type,
+        'pressure_gt_two': True if tile.mine.pressure > 2 else False
+    }
diff --git a/src/main.py b/src/main.py
index 82a8d59..1c74944 100644
--- a/src/main.py
+++ b/src/main.py
@@ -7,6 +7,7 @@ from environment import Environment
 from tilesFactory import TilesFactory
 from src.search_algoritms.a_star import a_star
 from search_algoritms.BFS import breadth_first_search
+from machine_learning.helpers import get_dataset_from_tile
 from machine_learning.decision_tree import get_decision, tree_root
 
 
@@ -24,15 +25,27 @@ def handle_keys(env: Environment, agent: Agent, game_ui: GameUi, factory: TilesF
                          or (agent.direction == 3 and agent.x * 80 > 5)
                          or (agent.direction == 0 and agent.y * 80 > 0)
                          or (agent.direction == 2 and agent.y * 80 < 700)):
-                print('Moving forward')
                 game_ui.move()
+                print('Moving forward')
             elif event.key == pg.K_SPACE:
                 if env.field[agent.y][agent.x].mine:
                     print('Defusing')
+                    tile = env.field[agent.y][agent.x]
+
+                    dataset = get_dataset_from_tile(tile)
+                    decision = get_decision(dataset, tree_root)
+
+                    print(f'Data: {dataset}')
+                    print(f'Decision: {decision}')
+
+                    if decision == 'detonation':
+                        game_ui.kaboom()
+
                     env.mine_count -= 1
                     env.field[agent.y][agent.x] = factory.create_tile(
                         IMAGES[env.field[agent.y][agent.x].number].parent
                     )
+
                     game_ui.update()
 
 
@@ -68,27 +81,19 @@ def main():
                     # path, actions = breadth_first_search(env.field, agent.x, agent.y, agent.direction)
                     goal = breadth_first_search(env.field, agent.x, agent.y, agent.direction, True)
                     path, actions = a_star(env.field, agent.x, agent.y, agent.direction, goal)
+
                     if not path and not env.field[agent.y][agent.x].mine:
                         print('Unable to find path, rocks are in the way')
                         break
                     print(f'Path{path}')
+
                     print(f'Actions:{actions}')
                     for action in actions:
                         pg.fastevent.post(pg.event.Event(pg.KEYDOWN, {'key': action}))
-                    goal_tile = env.field[goal[1]][goal[0]]
-                    data = {
-                        'visibility': goal_tile.visibility,
-                        'stability': goal_tile.stability,
-                        'armed': goal_tile.mine.armed,
-                        'ground': goal_tile.ground,
-                        'mine_type': goal_tile.mine.mine_type,
-                        'pressure_gt_two': True if goal_tile.mine.pressure > 2 else False
-                    }
-                    decision = get_decision(data, tree_root)
-                    pg.fastevent.post(pg.event.Event(pg.KEYDOWN, {'key': pg.K_SPACE}))
+
                     handle_keys(env, agent, game_ui, factory)
-                    print(data)
-                    print(f'Decision: {decision}')
+                    pg.fastevent.post(pg.event.Event(pg.KEYDOWN, {'key': pg.K_SPACE}))
+
                     print('Sector clear')
                 else:
                     pg.fastevent.post(event)