diff --git a/bot/main.py b/bot/main.py
index 00df95b..fd76fd2 100644
--- a/bot/main.py
+++ b/bot/main.py
@@ -1,8 +1,7 @@
 from typing import Optional
 from ares import AresBot
 from ares.behaviors.combat import CombatManeuver
-from ares.cython_extensions.units_utils import cy_closest_to
-from ares.cython_extensions.geometry import cy_distance_to
+from ares.cython_extensions.units_utils import cy_closest_to, cy_distance_to
 from ares.behaviors.combat.individual import (
     AMove,
     PathUnitToTarget,
@@ -11,26 +10,16 @@
 )
 
 from sc2.ids.unit_typeid import UnitTypeId
-from sc2.units import Units
 from sc2.unit import Unit
+from sc2.units import Units
 from sc2.position import Point2
 
 import numpy as np
-from sc2.units import Units
-
 
 class MyBot(AresBot):
     combat_manager = None
 
     def __init__(self, game_step_override: Optional[int] = None):
-        """Initiate custom bot
-
-        Parameters
-        ----------
-        game_step_override :
-            If provided, set the game_step to this value regardless of how it was
-            specified elsewhere
-        """
         super().__init__(game_step_override)
 
     async def on_step(self, iteration: int):
@@ -41,82 +30,38 @@ async def on_step(self, iteration: int):
 
         # define targets and grid
         enemy_units = self.enemy_units
-        ground_grid = np.ndarray = self.mediator.get_ground_grid
+        ground_grid = self.mediator.get_ground_grid
 
-        # call the engagement for zergling and pylon attack for roaches
-        self.do_zergling_engagement(zerglings, enemy_units, ground_grid)
-        self.do_roach_pylon_attack(roaches, ground_grid) 
-        
-    
-    def do_zergling_engagement(
-        self,
-        zerglings: Units,
-        enemies: Units,
-        grid: np.ndarray,
-    ):
-        print(f"Executing do_zergling_engagement with {len(zerglings)} zerglings")
-        # map center as teh default A-move target
-        attack_target = self.game_info.map_center
+        # call the pylon attack for zergling and engagement for roaches
+        self.do_zergling_pylon_attack(zerglings, ground_grid)
+        self.do_roach_engagement(roaches, enemy_units, ground_grid)
 
-        """Engage enemy units with zerglings and use stutter step behavior.
+    def do_zergling_pylon_attack(self, zerglings: Units, grid: np.ndarray):
+        print(f"Executing do_zergling_pylon_attack with {len(zerglings)} zerglings")
+        enemy_pylon = self._close_enemy_pylon()
+        enemy_start_location = self.enemy_start_locations[0]
 
-        Parameters
-        ----------
-        zerglings : Units
-            Zerglings to engage with.
-        enemies : Units
-            Enemy units to engage.
-        grid : np.ndarray
-            The ground grid for pathing information.
-        """
-        print(enemies)
         for zergling in zerglings:
             zergling_maneuver = CombatManeuver()
-
-            # If there are enemies, find the closest one and stutter step towards it
-            if enemies:
-                closest_enemy: Unit = cy_closest_to(zergling.position, enemies)  # Use zergling.position
-                zergling_maneuver.add(StutterUnitBack(zergling, closest_enemy))
+            if enemy_pylon and cy_distance_to(zergling.position, enemy_pylon.position) < 15.0:
+                zergling_maneuver.add(AttackTarget(zergling, enemy_pylon))
             else:
-                zergling_maneuver.add(AMove(zergling, attack_target))
+                zergling_maneuver.add(PathUnitToTarget(zergling, grid, enemy_start_location, success_at_distance=5.0))
             self.register_behavior(zergling_maneuver)
 
-    def do_roach_pylon_attack(
-        self,
-        roaches: Units,
-        grid: np.ndarray,
-    ):
-        print(f"Executing do_roach_pylon_attack with {len(roaches)} roaches")
-        """Direct roaches to move towards the enemy start location and attack the pylon once in range.
-
-        Parameters
-        ----------
-        roaches : Units
-            Roaches to perform the attack.
-        grid : np.ndarray
-            The ground grid for pathing information.
-        """
-        enemy_start_location = self.enemy_start_locations[0]
-        enemy_pylon = self._close_enemy_pylon()
-
+    def do_roach_engagement(self, roaches: Units, enemies: Units, grid: np.ndarray):
+        print(f"Executing do_roach_engagement with {len(roaches)} roaches")
         for roach in roaches:
             roach_maneuver = CombatManeuver()
-            # If a pylon is found and within 15 distance units, attack it
-            if enemy_pylon and cy_distance_to(roach.position, enemy_pylon.position) < 15.0:
-                roach_maneuver.add(AttackTarget(roach, enemy_pylon))
+            if enemies:
+                closest_enemy: Unit = cy_closest_to(roach.position, enemies)
+                roach_maneuver.add(StutterUnitBack(roach, closest_enemy))
             else:
-                # Path to the enemy start location
-                roach_maneuver.add(
-                    PathUnitToTarget(
-                        roach, grid, enemy_start_location, success_at_distance=5.0
-                    )
-                )
-            
+                roach_maneuver.add(AMove(roach, self.game_info.map_center))
             self.register_behavior(roach_maneuver)
 
     def _close_enemy_pylon(self) -> Unit:
-        """Find the enemy pylon to target."""
         pylons = self.enemy_structures(UnitTypeId.PYLON)
         if pylons:
-            return pylons.closest_to(self.enemy_start_locations[0])
+            return pylons.closest_to(self.game_info.map_center)
         return None