diff --git a/agent.py b/agent.py
index 63d3450..c87c891 100644
--- a/agent.py
+++ b/agent.py
@@ -5,7 +5,6 @@ class trashmaster(pygame.sprite.Sprite):
     def __init__(self):
         super().__init__()
         
-        
         self.images = []
         self.image = pygame.image.load("resources/textures/trashmaster_blu.png")
         self.image = pygame.transform.scale(self.image, (16, 16))
diff --git a/main.py b/main.py
index 31c2b4d..e707f7a 100644
--- a/main.py
+++ b/main.py
@@ -3,47 +3,56 @@ from map import preparedMap
 from agent import trashmaster
 
 
-class WalleGame():
-    def __init__(self):
-        pass
-    
-    
-    
-pygame.init()
-
 #config
 SCREEN_SIZE = [512, 512]
 BACKGROUND_COLOR = '#ffffff'
 
-if __name__ == '__main__':
+class WalleGame():
+    
+    def __init__(self):
+        self.SCREEN_SIZE = [512, 512]
+        self.BACKGROUND_COLOR = '#ffffff'
+        
+        pygame.init()
+        pygame.display.set_caption('Wall-e')
+        
+        
+        self.screen = pygame.display.set_mode(SCREEN_SIZE)
+        self.screen.fill(pygame.Color(self.BACKGROUND_COLOR))
+        
+        # krata
+        self.map = preparedMap(self.SCREEN_SIZE)
+        self.screen.blit(self.map, (0,0))
+        
+    def update_window(self):
+        pygame.display.update()
+    
+    def draw_trashmaster(self, smieciara: trashmaster):
+        smieciara_list = pygame.sprite.Group()
+        smieciara_list.add(smieciara)
+        smieciara_list.draw(self.screen)
+        
+        
 
-    pygame.init()
+def main():
+    game = WalleGame()
+    game.update_window()
+    
+    smieciara_object = trashmaster()
+    game.draw_trashmaster(smieciara_object)
 
-    # tytul okna
-    pygame.display.set_caption('Wall-e')
+    game.update_window()
 
-    screen = pygame.display.set_mode(SCREEN_SIZE)
-    screen.fill(pygame.Color(BACKGROUND_COLOR))
-
-    # krata
-    map = preparedMap(SCREEN_SIZE)
-    screen.blit(map, (0,0))
-
-    # update okna
-    pygame.display.update()
-
-smieciara1 = trashmaster()
-smieciara_list = pygame.sprite.Group()
-smieciara_list.add(smieciara1)
-smieciara_list.draw(screen)
-
-pygame.display.update()
-
-running = True
-while running:
-    for event in pygame.event.get():
-        if event.type == pygame.QUIT:
-            running = False
+    running = True
+    
+    while running:
+        for event in pygame.event.get():
+            if event.type == pygame.QUIT:
+                running = False
+                pygame.quit()
+    pygame.quit()
    
 
-pygame.quit()
+
+if __name__ == '__main__':
+    main()
\ No newline at end of file