diff --git a/AI_brain/movement.py b/AI_brain/go_any_direction.py
similarity index 98%
rename from AI_brain/movement.py
rename to AI_brain/go_any_direction.py
index acd4063..74f36b9 100644
--- a/AI_brain/movement.py
+++ b/AI_brain/go_any_direction.py
@@ -14,7 +14,7 @@ class State:
         return self.x == other.x and self.y == other.y
 
 
-class StateGraphSearchBFS:
+class GoAnyDirectionBFS:
     def __init__(self, world: World, start_state: State, goal_state: State):
         self.start_state = start_state
         self.goal_state = goal_state
diff --git a/AI_brain/template_bfs.py b/AI_brain/rotate_and_go_bfs.py
similarity index 84%
rename from AI_brain/template_bfs.py
rename to AI_brain/rotate_and_go_bfs.py
index 02a258d..966650d 100644
--- a/AI_brain/template_bfs.py
+++ b/AI_brain/rotate_and_go_bfs.py
@@ -1,6 +1,5 @@
 import queue
 
-from domain.commands.vacuum_move_command import VacuumMoveCommand
 from domain.world import World
 
 
@@ -34,7 +33,7 @@ def action_sequence(node: Node):
     return actions
 
 
-class TemplateBFS:
+class RotateAndGoBFS:
     def __init__(self, world: World, start_state: State, goal_state: State):
         self.world = world
         self.start_state = start_state
@@ -49,7 +48,7 @@ class TemplateBFS:
 
         while self.fringe:
             elem = self.fringe.get()
-            if elem.state == self.goal_state:
+            if self.is_goal(elem.state):
                 self.actions = action_sequence(elem)
                 return True
             self.explored.add(elem.state)
@@ -67,12 +66,18 @@ class TemplateBFS:
 
     def successors(self, state: State):
         new_successors = [
-            # ("GO", State(state.x + state.direction[0], state.y + state.direction[1], state.direction)),
             # rotate right
             ("RR", State(state.x, state.y, (-state.direction[1], state.direction[0]))),
             # rotate left
             ("RL", State(state.x, state.y, (state.direction[1], -state.direction[0]))),
         ]
         if self.world.accepted_move(state.x + state.direction[0], state.y + state.direction[1]):
-            new_successors.append(("GO", State(state.x + state.direction[0], state.y + state.direction[1], state.direction)))
+            new_successors.append(
+                ("GO", State(state.x + state.direction[0], state.y + state.direction[1], state.direction)))
         return new_successors
+
+    def is_goal(self, state: State) -> bool:
+        return (
+                state.x == self.goal_state.x
+                and state.y == self.goal_state.y
+        )
diff --git a/main.py b/main.py
index 0b23686..e97ac54 100644
--- a/main.py
+++ b/main.py
@@ -12,7 +12,7 @@ from domain.entities.docking_station import Doc_Station
 from domain.world import World
 from view.renderer import Renderer
 # from AI_brain.movement import StateGraphSearchBFS, State
-from AI_brain.template_bfs import TemplateBFS, State
+from AI_brain.rotate_and_go_bfs import RotateAndGoBFS, State
 
 
 config = configparser.ConfigParser()
@@ -49,22 +49,26 @@ class Main:
         start_state = State(self.world.vacuum.x, self.world.vacuum.y)
         end_state = State(self.world.doc_station.x, self.world.doc_station.y)
 
-        # SGS_BFS = StateGraphSearchBFS(self.world, start_state, end_state)
-        SGS_BFS = TemplateBFS(self.world, start_state, end_state)
-        if not SGS_BFS.search():
+        # path_searcher = GoAnyDirectionBFS(self.world, start_state, end_state)
+        path_searcher = RotateAndGoBFS(self.world, start_state, end_state)
+        if not path_searcher.search():
             print("No solution")
             exit(0)
 
-        SGS_BFS.actions.reverse()
+        path_searcher.actions.reverse()
         while self.running:
-            self.renderer.render(self.world)
-            self.clock.tick(5)
-            if len(SGS_BFS.actions) > 0:
-                action_direction = SGS_BFS.actions.pop()
+            for event in pygame.event.get():
+                if event.type == pygame.QUIT:
+                    self.running = False
+
+            if len(path_searcher.actions) > 0:
+                action_direction = path_searcher.actions.pop()
                 # self.handle_action1(action_direction)
                 self.handle_action2(action_direction)
 
-                self.update()
+            self.update()
+            self.renderer.render(self.world)
+            self.clock.tick(5)
 
         pygame.quit()