diff --git a/square.py b/square.py
index 7e11650..3455d8c 100644
--- a/square.py
+++ b/square.py
@@ -4,28 +4,35 @@ import robot
 DRIVE_A = 5.3 # 10 degrees turn at power 1
 TURN_A = 16 # 10 centimeter drive at power 1
 
+def square(arlo, power, turn_t, drive_t):
+    while True:
+        # Driving forward
+        arlo.go_diff(power, power, 1, 1)
+        sleep(drive_t * 10)
+        arlo.stop()
+
+        # Turning 90 degrees
+        arlo.go_diff(power, power, 1, 0)
+        sleep(turn_t * 9)
+        arlo.stop()
+
 def main():
+    # Gets the power for the robot
     power = input("Power (64):")
     if power == "":
         power = 64
     else:
         power = int(power)
 
-    turn_t = (DRIVE_A / power) * 9 # 10 degrees turn at given power
-    drive_t = (TURN_A / power) * 10 # 10 centimeter drive at given power
+    # Calculates sleep times for the robot
+    turn_t = DRIVE_A / power # 10 degrees turn at given power
+    drive_t = TURN_A / power # 10 centimeter drive at given power
+
+    # Initializes the robot and runs the
     arlo = robot.Robot()
 
     try:
-        while True:
-            # Driving forward
-            arlo.go_diff(power, power, 1, 1)
-            sleep(drive_t)
-            arlo.stop()
-
-            # Turning 90 degrees
-            arlo.go_diff(power, power, 1, 0)
-            sleep(turn_t)
-            arlo.stop()
+        square(arlo, power, turn_t, drive_t)
 
     except KeyboardInterrupt:
         arlo.stop()