diff --git a/gameobjects.py b/gameobjects.py
index 32d641c..19bd87b 100644
--- a/gameobjects.py
+++ b/gameobjects.py
@@ -3,6 +3,7 @@ import random
 
 WIDTH, HEIGHT = 800, 600
 HUD_HEIGHT = 150
+BARRIER_SEED = 20
 class GameObject:
     x = 0
     y = 0
@@ -186,6 +187,7 @@ class Spawner:
                 self.tick_counter = 0  # Reset tick counter
 
     def spawn(self):
-        random_y = random.randint(0, HEIGHT - 100)  # Random Y position for the barrier
+
+        random_y = random.randint(0, HEIGHT - 100, )  # Random Y position for the barrier
         new_barrier = Barrier(self.barrier_image_path, WIDTH, random_y)  # Create a new barrier
         return new_barrier
diff --git a/main.py b/main.py
index 718b328..6f3a469 100644
--- a/main.py
+++ b/main.py
@@ -2,23 +2,31 @@ import pygame
 import sys
 from gameobjects import *
 from genetic_alg import GeneticAlgorithm
+import matplotlib.pyplot as plt
 
-POPULATION_SIZE = 100
+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 = 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
+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 = []
@@ -38,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)
@@ -70,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()
 
@@ -79,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:
@@ -96,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()
@@ -123,17 +137,35 @@ 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:
+            # 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()
 
     UI.draw(screen)
     UI.draw_marks(screen)
-    
-    pygame.display.flip()  # Update the display
-    pygame.time.Clock().tick(FRAME_RATE)  # Limit the frame rate to 160 FPS
+
+    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()
+    pygame.time.Clock().tick(FRAME_RATE)