@@ -7,17 +7,15 @@
 c = 0  # Marginal cost
 l = 10  # Length of binary string (chromosomes)
 T = 1000  # Max iterations
-
-# Student ID logic for probabilities
 Y = 1  # Second to last digit of ID
 Z = 3  # Last digit of ID
 
-p_cross = 0.200 + (0.01 * Y) + (0.001 * Z)
-p_mut = 0.00200 + (0.0001 * Y) + (0.00001 * Z)
+p_cross_bench = 0.200 + (0.01 * Y) + (0.001 * Z)
+p_mut_bench = 0.00200 + (0.0001 * Y) + (0.00001 * Z)
 
 # Theoretical Values
-Q_cournot_total = (n / (n + 1)) * (alpha - c)
-q_cournot_indiv = Q_cournot_total / n
+Q_nash_total = (n / (n + 1)) * (alpha - c)
+q_nash_indiv = Q_nash_total / n
 
 
 def BinaryToDecimal(binary_array):
@@ -29,14 +27,14 @@ def BinaryToDecimal(binary_array):
 
 def BinaryToOutput(binary_array):
     x_it = BinaryToDecimal(binary_array)
-    q_i = (alpha / n) * (x_it / (2**l))  ### I am not sure about this scaling
+    q_i = alpha * (
+        x_it / (2**l)
+    )  # Scale the integer [0, 1023] to output (Note: 2^10 = 1024. Max x_it is 1023.)
     return q_i
 
 
 def CalculatePrice(total_Q):
     price = alpha - total_Q
-    if price < 0:
-        price = 0
     return price
 
 
@@ -52,9 +50,14 @@ def Variance(values):
     return variance
 
 
-def run_simulation(p_cross, p_mut):
-    Population = []
+def run_simulation(p_cross, p_mut, scenario_name):
+    print(f"\n--- Running Scenario: {scenario_name} ---")
+    print(f"p_cross: {p_cross:.5f}, p_mut: {p_mut:.5f}")
 
+    random.seed(42)
+    np.random.seed(42)
+
+    Population = []
     for i in range(n):
         # Generate random binary string of length l
         genes = [random.randint(0, 1) for _ in range(l)]
@@ -65,57 +68,51 @@ def run_simulation(p_cross, p_mut):
     variance_history = []
 
     for t in range(1, T + 1):
-
         outputs_list = []
         total_Q = 0
 
+        # Calculate outputs for current population
         for firm in Population:
             q_i = BinaryToOutput(firm)
             outputs_list.append(q_i)
             total_Q += q_i
 
         current_price = CalculatePrice(total_Q)
 
-        # Check for zero-profit trap: if price is 0, restart population
-        if current_price <= 0:
-            Population = []
-            for i in range(n):
-                genes = [random.randint(0, 1) for _ in range(l)]
-                Population.append(genes)
-            continue
-
         # Store stats
         total_Q_history.append(total_Q)
         variance_history.append(Variance(outputs_list))
 
-        # A. RECOMBINATION
-        # Calculate fitness scores (profits) for all firms
+        # A. RECOMBINATION (Roulette Wheel Selection)
         fitness_scores = []
         for firm in Population:
             q_i = BinaryToOutput(firm)
             profit = CalculateProfit(q_i, current_price)
             fitness_scores.append(profit)
 
-        # Check if we can use fitness-based selection
+        # Handle negative fitness for probability calculation
         min_fitness = min(fitness_scores)
-        total_fitness = sum(fitness_scores)
+        if min_fitness < 0:
+            # Shift all fitness values to be positive
+            offset = -min_fitness
+            # Add a small epsilon to ensure non-zero probabilities if all are equal
+            fitness_scores = [f + offset + 1e-6 for f in fitness_scores]
+            # But this solution can left useless in a situation where all firms have zero fitness score
 
-        # If negative fitness or zero total fitness, restart population
-        if min_fitness < 0 or total_fitness <= 0:
-            Population = []
-            for i in range(n):
-                genes = [random.randint(0, 1) for _ in range(l)]
-                Population.append(genes)
-            continue
+        total_fitness = sum(fitness_scores)
 
-        # Normalize to get probabilities
-        probabilities = [f / total_fitness for f in fitness_scores]
+        if total_fitness == 0:
+            probabilities = [
+                1.0 / n
+            ] * n  # If all firms have zero fitness score, each firm has equal probability of being selected
+        else:
+            probabilities = [f / total_fitness for f in fitness_scores]
 
-        # Select firms based on their fitness probabilities
+        # Create mating pool
         mating_pool = []
-        for _ in range(n):
-            selected_firm_idx = np.random.choice(len(Population), p=probabilities)
-            mating_pool.append(Population[selected_firm_idx].copy())
+        indices = np.random.choice(len(Population), size=n, p=probabilities)
+        for idx in indices:
+            mating_pool.append(Population[idx].copy())
 
         # B. CROSSOVER
         offspring_pool = []
@@ -150,14 +147,14 @@ def run_simulation(p_cross, p_mut):
             q_parent = BinaryToOutput(Parent)
             q_child = BinaryToOutput(Child)
 
-            # Calculate Q_others
+            # Calculate Q_others (aggregate of all others in current generation)
             Q_others = total_Q - q_parent
 
             # Parent Profit (Actual)
-            # Uses the actual market price that occurred
             profit_parent = CalculateProfit(q_parent, current_price)
 
-            # Recalculate price if Child were in the market
+            # Child Profit (Hypothetical)
+            # Evaluate Child against the *current* behavior of others
             new_total_Q = Q_others + q_child
             new_price = CalculatePrice(new_total_Q)
             profit_child = CalculateProfit(q_child, new_price)
@@ -169,43 +166,46 @@ def run_simulation(p_cross, p_mut):
 
         Population = Next_Generation
 
-    # ---------------------------------------------------------
-    # RESULTS
-    # ---------------------------------------------------------
-    print(f"\nSimulation completed!")
-    print(f"Final Total Q: {total_Q_history[-1]}")
-    print(f"Target Cournot Total Q: {Q_cournot_total:.2f}")
-    print(f"Final Variance: {variance_history[-1]}")
+    # RESULTS & PLOTTING
+    print(f"Simulation completed!")
+    print(f"Final Total Q: {total_Q_history[-1]:.2f}")
+    print(f"Target Nash Total Q: {Q_nash_total:.2f}")
+    print(f"Final Variance: {variance_history[-1]:.2f}")
 
+    # Plotting
     plt.figure(figsize=(12, 5))
 
+    # Plot 1: Total Quantity
     plt.subplot(1, 2, 1)
-    plt.plot(total_Q_history, label="Actual Total Q")
-    plt.axhline(y=Q_cournot_total, color="r", linestyle="--", label="Cournot Eq")
+    plt.plot(total_Q_history, label="Actual Total Q", color="blue")
+    plt.axhline(y=Q_nash_total, color="r", linestyle="--", label="Nash Equilibrium")
     plt.xlabel("Iteration")
     plt.ylabel("Total Quantity (Q)")
-    plt.title("Total Quantity Convergence")
+    plt.title(f"Industry Output Convergence\n({scenario_name})")
     plt.legend()
-    plt.grid(True)
+    plt.grid(True, alpha=0.3)
 
+    # Plot 2: Variance
     plt.subplot(1, 2, 2)
-    plt.plot(variance_history)
+    plt.plot(variance_history, color="green")
     plt.xlabel("Iteration")
-    plt.ylabel("Variance of Outputs")
-    plt.title("Variance Over Time")
-    plt.grid(True)
+    plt.ylabel("Variance")
+    plt.title(f"Population Variance\n({scenario_name})")
+    plt.grid(True, alpha=0.3)
 
     plt.tight_layout()
-    plt.savefig("genetic_algorithm_results.png")
-    plt.show()
+    filename = f"results_{scenario_name.replace(' ', '_')}.png"
+    plt.savefig(f"results/{filename}")
+    print(f"Saved plot to results/{filename}")
+    plt.close()
 
 
 if __name__ == "__main__":
-    # Scenario 1: Benchmark
-    run_simulation(p_cross, p_mut)
+    # (ii) Benchmark Scenario
+    run_simulation(p_cross_bench, p_mut_bench, "Benchmark")
 
-    # Scenario 2: High Crossover (+0.250)
-    run_simulation(p_cross + 0.250, p_mut)
+    # (iii) High Crossover (+0.250)
+    run_simulation(p_cross_bench + 0.250, p_mut_bench, "High_Crossover")
 
-    # Scenario 3: High Mutation (+0.00250)
-    run_simulation(p_cross, p_mut + 0.00250)
+    # (iv) High Mutation (+0.00250)
+    run_simulation(p_cross_bench, p_mut_bench + 0.00250, "High_Mutation")