diff --git a/graphics.py b/graphics.py
index 828d3fb..6efae19 100644
--- a/graphics.py
+++ b/graphics.py
@@ -40,7 +40,6 @@ class Window:
     def __init__(self, width: int, height: int) -> None:
         self._root = Tk()
         self._root.title("Maze Solver")
-        self._root.protocol("WM_DELETE_WINDOW", self.close)
 
         # Position the window to the centre of the screen
         screen_width = self._root.winfo_screenwidth()
@@ -61,8 +60,6 @@ class Window:
         )
         self._canvas.pack()
 
-        self._running = False
-
     def redraw(self) -> None:
         """
         redraw redraws all the graphics in the window.
@@ -70,23 +67,15 @@ class Window:
         self._root.update_idletasks()
         self._root.update()
 
-    def wait_for_close(self) -> None:
+    def mainloop(self) -> None:
         """
-        wait_for_close continuously redraws the window until
-        it is set to close.
+        mainloop calls the root widget's mainloop method to
+        ensure that the window remains visible on the screen.
         """
-        self._running = True
-        while self._running:
-            self.redraw()
+        self._root.mainloop()
 
     def draw_line(self, line: Line, fill_colour: str = "black") -> None:
         """
         draw_line draws a line on the canvas.
         """
         line.draw(self._canvas, fill_colour)
-
-    def close(self) -> None:
-        """
-        close sets the window to close.
-        """
-        self._running = False
diff --git a/main.py b/main.py
index 9337fed..7a9b18e 100644
--- a/main.py
+++ b/main.py
@@ -5,9 +5,17 @@ from maze import Maze
 def main():
     window = Window(800, 800)
 
-    _ = Maze(10, 10, 30, 30, 20, 20, window)
+    _ = Maze(
+        x_position=10,
+        y_position=10,
+        num_cell_rows=30,
+        num_cells_per_row=30,
+        cell_size_x=20,
+        cell_size_y=20,
+        window=window
+    )
 
-    window.wait_for_close()
+    window.mainloop()
 
 
 if __name__ == "__main__":