diff --git a/main.py b/main.py
index 8308f57..6f3a469 100644
--- a/main.py
+++ b/main.py
@@ -2,29 +2,31 @@ import pygame
 import sys
 from gameobjects import *
 from genetic_alg import GeneticAlgorithm
+import matplotlib.pyplot as plt
 
 POPULATION_SIZE = 300
 MUTATION_RATE = 0.5
 POPULATION_NEW = 0.1
 POPULATION_BEST = 0.3
-FRAME_RATE = 300  # TODO: FIX THE INCORRECT FRAME RATE CORRELATION
+FRAME_RATE = 300
 BARRIER_SPEED = 10
 BARRIER_DELAY = 100
-MAX_ITERATIONS = 50
+MAX_ITERATIONS = 5
+BARRIER_SEED = 42
 
 current_iteration = 1
-
+best_fitnesses = []
 
 pygame.init()
 
-screen = pygame.display.set_mode((WIDTH, HEIGHT + HUD_HEIGHT))  # Initialize the screen
-pygame.display.set_caption("NIC_Project")  # Set window title
-font = pygame.font.SysFont("Arial", 36) 
+screen = pygame.display.set_mode((WIDTH, HEIGHT + HUD_HEIGHT))
+pygame.display.set_caption("NIC_Project")
+font = pygame.font.SysFont("Arial", 36)
 
-WHITE = (255, 255, 255)  # Define white color
+WHITE = (255, 255, 255)
 BLACK = (0, 0, 0)
 
-UI = UI("images/HUD.png")  # Initialize UI
+UI = UI("images/HUD.png")
 gameobjects = []
 horses = []
 barriers = []
@@ -36,7 +38,7 @@ last_barrier = None
 
 def init_game():
     global gameobjects, horses, barriers, upper_bound_rect, lower_bound_rect, spawner, last_barrier, BARRIER_SPEED
-    random.seed(BARRIER_SEED)
+
     grass = Background("images/Grass.jpg", 0, 0)
     grass.set_size(WIDTH, HEIGHT)
 
@@ -44,7 +46,7 @@ def init_game():
     barriers = []
 
     for horse in horses:
-        horse.set_position(50, HEIGHT/2)
+        horse.set_position(50, HEIGHT / 2)
         horse.stopped = False
         horse.set_vacceleration(0)
         horse.set_vspeed(0)
@@ -76,7 +78,7 @@ def get_features(horse: Horse):
     return features
 
 
-horses = [Horse("images/Horse_1.png", 50, HEIGHT/2 + 0 * i) for i in range(POPULATION_SIZE)]
+horses = [Horse("images/Horse_1.png", 50, HEIGHT / 2 + 0 * i) for i in range(POPULATION_SIZE)]
 
 init_game()
 
@@ -85,14 +87,20 @@ genecticAlg = GeneticAlgorithm(POPULATION_SIZE, MUTATION_RATE, POPULATION_BEST,
 
 while True:
     for event in pygame.event.get():
-        if event.type == pygame.QUIT:  # Handle window close event
+        if event.type == pygame.QUIT:
+            # Generate and display the Matplotlib graph
+            plt.plot(best_fitnesses)
+            plt.xlabel("Iteration")
+            plt.ylabel("Max Fitness")
+            plt.title("Max Fitness per Iteration")
+            plt.show()
             pygame.quit()
             sys.exit()
 
-    keys = pygame.key.get_pressed()  # Get pressed keys
+    keys = pygame.key.get_pressed()
 
     for i, horse in enumerate(horses):
-        if not horse.stopped:  # If the horse is not stopped
+        if not horse.stopped:
             data = get_features(horse)
             res = genecticAlg.predict(data, i)
             if res == 0:
@@ -102,21 +110,21 @@ while True:
             else:
                 horse.stay()
         else:
-            horse.move(-BARRIER_SPEED, 0)  # Move stopped horse to the left
+            horse.move(-BARRIER_SPEED, 0)
 
-        horse.apply_vacceleration()  # Apply acceleration to horse
-        horse.apply_vspeed()  # Apply speed to horse
-        horse.draw(screen)  # Draw the horse
+        horse.apply_vacceleration()
+        horse.apply_vspeed()
+        horse.draw(screen)
 
-    spawner.handle()  # Handle spawner logic
-    if spawner.tick_counter == 0:  # If a new barrier is spawned
-        new_barrier = spawner.spawn()  # Spawn a new barrier
-        barriers.append(new_barrier)  # Add barrier to barriers list
-        gameobjects.append(new_barrier)  # Add barrier to game objects
+    spawner.handle()
+    if spawner.tick_counter == 0:
+        new_barrier = spawner.spawn()
+        barriers.append(new_barrier)
+        gameobjects.append(new_barrier)
         last_barrier = new_barrier
 
     for barrier in barriers:
-        barrier.move(-BARRIER_SPEED, 0)  # Move barriers to the left
+        barrier.move(-BARRIER_SPEED, 0)
 
     for horse in horses:
         horse.update_animation()
@@ -129,16 +137,23 @@ while True:
                 horse.stop()
             horse.count_fitness()
 
-
     for object in gameobjects:
-        object.draw(screen)  # Draw all game objects  
-
+        object.draw(screen)
 
     if all(horse.stopped for horse in horses):
         UI.iteration_num += 1
         genecticAlg.learn([x.fitness for x in horses])
+        
+        # store the best fitness
+        best_fitnesses.append(max(genecticAlg.fitnessBest) if len(genecticAlg.fitnessBest) != 0 else 0)
         current_iteration += 1
-        if current_iteration == MAX_ITERATIONS:
+        if current_iteration > MAX_ITERATIONS:
+            # Generate and display the Matplotlib graph
+            plt.plot(best_fitnesses)
+            plt.xlabel("Iteration")
+            plt.ylabel("Max Fitness")
+            plt.title("Max Fitness per Iteration")
+            plt.show()
             pygame.quit()
             sys.exit()
         init_game()
@@ -146,11 +161,11 @@ while True:
     UI.draw(screen)
     UI.draw_marks(screen)
 
-    
     iteration_num_txt = font.render(f"iteration: {current_iteration}", True, BLACK)
-    best_fitness_txt = font.render(f"Best: {max(genecticAlg.fitnessBest) if len(genecticAlg.fitnessBest) != 0 else 0}",True, BLACK)
-    screen.blit(iteration_num_txt, (WIDTH/2 - 90, HEIGHT + HUD_HEIGHT - 120)) 
-    screen.blit(best_fitness_txt, (WIDTH/2 - 90, HEIGHT + HUD_HEIGHT - 80)) 
-    
-    pygame.display.flip()  # Update the display
-    pygame.time.Clock().tick(FRAME_RATE)  # Limit the frame rate to 160 FPS
+    best_fitness_txt = font.render(
+        f"Best: {max(genecticAlg.fitnessBest) if len(genecticAlg.fitnessBest) != 0 else 0}", True, BLACK)
+    screen.blit(iteration_num_txt, (WIDTH / 2 - 90, HEIGHT + HUD_HEIGHT - 120))
+    screen.blit(best_fitness_txt, (WIDTH / 2 - 90, HEIGHT + HUD_HEIGHT - 80))
+
+    pygame.display.flip()
+    pygame.time.Clock().tick(FRAME_RATE)