import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.optimize import minimize, differential_evolution
import time

# Set random seed for reproducibility
np.random.seed(42)

print("=" * 70)
print("ECOSYSTEM RESOURCE ALLOCATION OPTIMIZATION")
print("Maximizing Species Survival Probability")
print("=" * 70)

# ============================================================================
# PROBLEM PARAMETERS
# ============================================================================
n_species = 4  # Number of species
R_total = 100.0  # Total available resources

# Species parameters
alpha = np.array([0.15, 0.12, 0.18, 0.10])  # Resource efficiency
beta = np.array([15.0, 20.0, 12.0, 18.0])   # Minimum resource threshold

# Interaction matrix (gamma_ij): how species j affects species i
# Negative values = competition, Positive values = mutualism
gamma = np.array([
    [ 0.00, -0.03, -0.02,  0.01],  # Species 0
    [-0.03,  0.00, -0.01,  0.02],  # Species 1
    [-0.02, -0.01,  0.00, -0.04],  # Species 2
    [ 0.01,  0.02, -0.04,  0.00]   # Species 3
])

species_names = ["Herbivore A", "Herbivore B", "Predator C", "Pollinator D"]

print(f"\nNumber of species: {n_species}")
print(f"Total resources available: {R_total}")
print(f"\nSpecies parameters:")
for i in range(n_species):
    print(f"  {species_names[i]:15s}: α={alpha[i]:.3f}, β={beta[i]:.1f}")

# ============================================================================
# SURVIVAL PROBABILITY FUNCTIONS
# ============================================================================
def survival_probability(r_i, r_others, i):
    """
    Calculate survival probability for species i
    
    P_i = 1 / (1 + exp(-alpha_i * (r_i - beta_i) + sum(gamma_ij * r_j)))
    """
    interaction_effect = sum(gamma[i, j] * r_others[j] for j in range(len(r_others)) if j != i)
    exponent = -alpha[i] * (r_i - beta[i]) + interaction_effect
    return 1.0 / (1.0 + np.exp(exponent))

def ecosystem_survival_probability(r):
    """
    Calculate overall ecosystem survival probability (product of individual probabilities)
    """
    if np.any(r < 0) or np.sum(r) > R_total:
        return 0.0
    
    prob = 1.0
    for i in range(n_species):
        p_i = survival_probability(r[i], r, i)
        prob *= p_i
    
    return prob

def negative_log_probability(r):
    """
    Objective function for optimization (minimize negative log probability)
    Using log for numerical stability
    """
    if np.any(r < 0) or np.sum(r) > R_total:
        return 1e10
    
    log_prob = 0.0
    for i in range(n_species):
        p_i = survival_probability(r[i], r, i)
        if p_i <= 0:
            return 1e10
        log_prob += np.log(p_i)
    
    return -log_prob

# ============================================================================
# OPTIMIZATION
# ============================================================================
print("\n" + "=" * 70)
print("OPTIMIZATION PROCESS")
print("=" * 70)

# Constraints
constraints = {'type': 'ineq', 'fun': lambda r: R_total - np.sum(r)}
bounds = [(0, R_total) for _ in range(n_species)]

# Initial guess: equal distribution
r_initial = np.ones(n_species) * (R_total / n_species)

print(f"\nInitial allocation (equal distribution):")
print(f"  Resources: {r_initial}")
print(f"  Survival probability: {ecosystem_survival_probability(r_initial):.6f}")

# Method 1: Sequential Least Squares Programming (SLSQP)
print(f"\n[Method 1] Running SLSQP optimization...")
start_time = time.time()
result_slsqp = minimize(
    negative_log_probability,
    r_initial,
    method='SLSQP',
    bounds=bounds,
    constraints=constraints,
    options={'maxiter': 1000}
)
time_slsqp = time.time() - start_time

r_optimal_slsqp = result_slsqp.x
prob_optimal_slsqp = ecosystem_survival_probability(r_optimal_slsqp)

print(f"  Time: {time_slsqp:.3f} seconds")
print(f"  Success: {result_slsqp.success}")
print(f"  Optimal allocation: {r_optimal_slsqp}")
print(f"  Total allocated: {np.sum(r_optimal_slsqp):.2f}")
print(f"  Optimal survival probability: {prob_optimal_slsqp:.6f}")

# Method 2: Differential Evolution (global optimizer)
print(f"\n[Method 2] Running Differential Evolution (global search)...")
start_time = time.time()
result_de = differential_evolution(
    negative_log_probability,
    bounds,
    seed=42,
    maxiter=300,
    popsize=15,
    atol=1e-6,
    workers=1
)
time_de = time.time() - start_time

r_optimal_de = result_de.x
prob_optimal_de = ecosystem_survival_probability(r_optimal_de)

print(f"  Time: {time_de:.3f} seconds")
print(f"  Success: {result_de.success}")
print(f"  Optimal allocation: {r_optimal_de}")
print(f"  Total allocated: {np.sum(r_optimal_de):.2f}")
print(f"  Optimal survival probability: {prob_optimal_de:.6f}")

# Use the better result
if prob_optimal_de > prob_optimal_slsqp:
    r_optimal = r_optimal_de
    prob_optimal = prob_optimal_de
    best_method = "Differential Evolution"
else:
    r_optimal = r_optimal_slsqp
    prob_optimal = prob_optimal_slsqp
    best_method = "SLSQP"

print(f"\n[Best Result] Using {best_method}")
print(f"  Improvement: {(prob_optimal / ecosystem_survival_probability(r_initial) - 1) * 100:.2f}%")

# ============================================================================
# DETAILED ANALYSIS
# ============================================================================
print("\n" + "=" * 70)
print("DETAILED ANALYSIS")
print("=" * 70)

print("\nOptimal Resource Allocation:")
for i in range(n_species):
    pct = (r_optimal[i] / R_total) * 100
    p_i = survival_probability(r_optimal[i], r_optimal, i)
    print(f"  {species_names[i]:15s}: {r_optimal[i]:6.2f} units ({pct:5.1f}%) | P={p_i:.4f}")

print(f"\nOverall Ecosystem Survival Probability: {prob_optimal:.6f}")

# ============================================================================
# VISUALIZATION
# ============================================================================
print("\n" + "=" * 70)
print("GENERATING VISUALIZATIONS")
print("=" * 70)

fig = plt.figure(figsize=(18, 12))

# Plot 1: Resource Allocation Comparison
ax1 = plt.subplot(2, 3, 1)
x_pos = np.arange(n_species)
width = 0.35
ax1.bar(x_pos - width/2, r_initial, width, label='Initial (Equal)', alpha=0.7, color='skyblue')
ax1.bar(x_pos + width/2, r_optimal, width, label='Optimal', alpha=0.7, color='salmon')
ax1.set_xlabel('Species', fontsize=11, fontweight='bold')
ax1.set_ylabel('Resource Allocation', fontsize=11, fontweight='bold')
ax1.set_title('Resource Allocation: Initial vs Optimal', fontsize=12, fontweight='bold')
ax1.set_xticks(x_pos)
ax1.set_xticklabels([s.split()[0] for s in species_names], rotation=45, ha='right')
ax1.legend()
ax1.grid(axis='y', alpha=0.3)

# Plot 2: Individual Survival Probabilities
ax2 = plt.subplot(2, 3, 2)
probs_initial = [survival_probability(r_initial[i], r_initial, i) for i in range(n_species)]
probs_optimal = [survival_probability(r_optimal[i], r_optimal, i) for i in range(n_species)]
ax2.bar(x_pos - width/2, probs_initial, width, label='Initial', alpha=0.7, color='skyblue')
ax2.bar(x_pos + width/2, probs_optimal, width, label='Optimal', alpha=0.7, color='salmon')
ax2.set_xlabel('Species', fontsize=11, fontweight='bold')
ax2.set_ylabel('Survival Probability', fontsize=11, fontweight='bold')
ax2.set_title('Individual Species Survival Probabilities', fontsize=12, fontweight='bold')
ax2.set_xticks(x_pos)
ax2.set_xticklabels([s.split()[0] for s in species_names], rotation=45, ha='right')
ax2.legend()
ax2.grid(axis='y', alpha=0.3)
ax2.set_ylim([0, 1])

# Plot 3: Survival Probability vs Resource (for each species)
ax3 = plt.subplot(2, 3, 3)
r_range = np.linspace(0, R_total * 0.5, 100)
for i in range(n_species):
    probs = []
    for r in r_range:
        r_test = r_optimal.copy()
        r_test[i] = r
        if np.sum(r_test) <= R_total:
            p = survival_probability(r, r_test, i)
            probs.append(p)
        else:
            probs.append(np.nan)
    ax3.plot(r_range, probs, label=species_names[i].split()[0], linewidth=2)
    ax3.axvline(r_optimal[i], color='gray', linestyle='--', alpha=0.5)

ax3.set_xlabel('Resource Allocation', fontsize=11, fontweight='bold')
ax3.set_ylabel('Survival Probability', fontsize=11, fontweight='bold')
ax3.set_title('Survival Probability vs Resource Allocation', fontsize=12, fontweight='bold')
ax3.legend()
ax3.grid(alpha=0.3)

# Plot 4: 3D Surface - Two-species interaction
ax4 = fig.add_subplot(2, 3, 4, projection='3d')
species_pair = (0, 2)  # Herbivore A vs Predator C
r1_range = np.linspace(0, 50, 30)
r2_range = np.linspace(0, 50, 30)
R1, R2 = np.meshgrid(r1_range, r2_range)
Z = np.zeros_like(R1)

r_base = r_optimal.copy()
for i in range(R1.shape[0]):
    for j in range(R1.shape[1]):
        r_base[species_pair[0]] = R1[i, j]
        r_base[species_pair[1]] = R2[i, j]
        if np.sum(r_base) <= R_total:
            Z[i, j] = ecosystem_survival_probability(r_base)
        else:
            Z[i, j] = np.nan

surf = ax4.plot_surface(R1, R2, Z, cmap='viridis', alpha=0.8, edgecolor='none')
ax4.scatter([r_optimal[species_pair[0]]], [r_optimal[species_pair[1]]], 
           [prob_optimal], color='red', s=100, marker='o', label='Optimal')
ax4.set_xlabel(species_names[species_pair[0]], fontsize=10, fontweight='bold')
ax4.set_ylabel(species_names[species_pair[1]], fontsize=10, fontweight='bold')
ax4.set_zlabel('Ecosystem Survival Prob', fontsize=10, fontweight='bold')
ax4.set_title('3D Surface: Ecosystem Survival Probability', fontsize=12, fontweight='bold')
fig.colorbar(surf, ax=ax4, shrink=0.5)

# Plot 5: Resource Allocation Pie Chart
ax5 = plt.subplot(2, 3, 5)
colors_pie = plt.cm.Set3(np.linspace(0, 1, n_species))
wedges, texts, autotexts = ax5.pie(r_optimal, labels=[s.split()[0] for s in species_names], 
                                     autopct='%1.1f%%', colors=colors_pie, startangle=90)
for autotext in autotexts:
    autotext.set_color('white')
    autotext.set_fontweight('bold')
ax5.set_title('Optimal Resource Distribution', fontsize=12, fontweight='bold')

# Plot 6: Interaction Matrix Heatmap
ax6 = plt.subplot(2, 3, 6)
im = ax6.imshow(gamma, cmap='RdYlGn', aspect='auto', vmin=-0.05, vmax=0.05)
ax6.set_xticks(range(n_species))
ax6.set_yticks(range(n_species))
ax6.set_xticklabels([s.split()[0] for s in species_names], rotation=45, ha='right')
ax6.set_yticklabels([s.split()[0] for s in species_names])
ax6.set_title('Species Interaction Matrix (γ)', fontsize=12, fontweight='bold')

for i in range(n_species):
    for j in range(n_species):
        text = ax6.text(j, i, f'{gamma[i, j]:.3f}', ha='center', va='center', 
                       color='black' if abs(gamma[i, j]) < 0.03 else 'white', fontsize=9)

plt.colorbar(im, ax=ax6, label='Interaction strength')

plt.tight_layout()
plt.savefig('ecosystem_optimization_results.png', dpi=150, bbox_inches='tight')
print("\nVisualization saved as 'ecosystem_optimization_results.png'")
plt.show()

# ============================================================================
# SENSITIVITY ANALYSIS
# ============================================================================
print("\n" + "=" * 70)
print("SENSITIVITY ANALYSIS")
print("=" * 70)

fig2, axes = plt.subplots(1, 2, figsize=(14, 5))

# Sensitivity to total resources
ax_sens1 = axes[0]
R_range = np.linspace(50, 150, 20)
probs_vs_R = []

print("\nAnalyzing sensitivity to total resource availability...")
for R in R_range:
    R_total_temp = R
    constraints_temp = {'type': 'ineq', 'fun': lambda r: R_total_temp - np.sum(r)}
    result_temp = minimize(
        negative_log_probability,
        np.ones(n_species) * (R_total_temp / n_species),
        method='SLSQP',
        bounds=[(0, R_total_temp) for _ in range(n_species)],
        constraints=constraints_temp,
        options={'maxiter': 500}
    )
    probs_vs_R.append(ecosystem_survival_probability(result_temp.x))

ax_sens1.plot(R_range, probs_vs_R, 'b-o', linewidth=2, markersize=6)
ax_sens1.axvline(R_total, color='red', linestyle='--', label=f'Current R={R_total}')
ax_sens1.set_xlabel('Total Resources Available', fontsize=11, fontweight='bold')
ax_sens1.set_ylabel('Optimal Ecosystem Survival Prob', fontsize=11, fontweight='bold')
ax_sens1.set_title('Sensitivity to Total Resource Availability', fontsize=12, fontweight='bold')
ax_sens1.legend()
ax_sens1.grid(alpha=0.3)

# Sensitivity to individual species resource allocation
ax_sens2 = axes[1]
species_idx = 0  # Test for first species
r_test_range = np.linspace(0, 50, 50)
probs_vs_r = []

print(f"Analyzing sensitivity to {species_names[species_idx]} resource allocation...")
for r_test in r_test_range:
    r_temp = r_optimal.copy()
    r_temp[species_idx] = r_test
    if np.sum(r_temp) <= R_total:
        probs_vs_r.append(ecosystem_survival_probability(r_temp))
    else:
        probs_vs_r.append(np.nan)

ax_sens2.plot(r_test_range, probs_vs_r, 'g-', linewidth=2)
ax_sens2.axvline(r_optimal[species_idx], color='red', linestyle='--', 
                label=f'Optimal r={r_optimal[species_idx]:.1f}')
ax_sens2.set_xlabel(f'{species_names[species_idx]} Resource Allocation', fontsize=11, fontweight='bold')
ax_sens2.set_ylabel('Ecosystem Survival Probability', fontsize=11, fontweight='bold')
ax_sens2.set_title(f'Sensitivity to {species_names[species_idx]} Allocation', fontsize=12, fontweight='bold')
ax_sens2.legend()
ax_sens2.grid(alpha=0.3)

plt.tight_layout()
plt.savefig('sensitivity_analysis.png', dpi=150, bbox_inches='tight')
print("Sensitivity analysis saved as 'sensitivity_analysis.png'")
plt.show()

print("\n" + "=" * 70)
print("ANALYSIS COMPLETE")
print("=" * 70)
print("\nKey Findings:")
print(f"1. Optimal allocation improves survival probability by "
      f"{(prob_optimal / ecosystem_survival_probability(r_initial) - 1) * 100:.1f}%")
print(f"2. Species with higher efficiency (α) and lower interaction costs receive more resources")
print(f"3. The optimization balances individual needs with ecosystem-level interactions")
print(f"4. Predator-prey and competitive relationships significantly affect optimal allocation")
print("\n" + "=" * 70)

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import torch
import torch.nn as nn
import torch.optim as optim
from time import time

# Set random seeds for reproducibility
np.random.seed(42)
torch.manual_seed(42)

print("="*60)
print("Neural Network Loss Function Minimization")
print("Scientific Data Analysis - Deep Model Optimization")
print("="*60)

# ============================================================
# 1. Generate Scientific Data
# ============================================================
print("\n[1] Generating scientific data...")

def scientific_function(x, y):
    """
    Complex scientific function: f(x,y) = sin(2πx)·cos(2πy)·exp(-(x²+y²))
    Common in wave physics and signal processing
    """
    return np.sin(2*np.pi*x) * np.cos(2*np.pi*y) * np.exp(-(x**2 + y**2))

# Create dense grid for training
n_train = 2000
x_train = np.random.uniform(-1.5, 1.5, n_train)
y_train = np.random.uniform(-1.5, 1.5, n_train)
z_train = scientific_function(x_train, y_train)

# Create grid for visualization
n_grid = 50
x_grid = np.linspace(-1.5, 1.5, n_grid)
y_grid = np.linspace(-1.5, 1.5, n_grid)
X_grid, Y_grid = np.meshgrid(x_grid, y_grid)
Z_true = scientific_function(X_grid, Y_grid)

print(f"Training samples: {n_train}")
print(f"Input range: x,y ∈ [-1.5, 1.5]")
print(f"Output range: z ∈ [{z_train.min():.3f}, {z_train.max():.3f}]")

# ============================================================
# 2. Define Deep Neural Network Architecture
# ============================================================
print("\n[2] Building deep neural network...")

class ScientificNN(nn.Module):
    """
    Deep neural network for scientific function approximation
    Architecture: Input(2) -> Hidden(64) -> Hidden(64) -> Hidden(32) -> Output(1)
    Uses Tanh activation for smooth gradients
    """
    def __init__(self):
        super(ScientificNN, self).__init__()
        self.layer1 = nn.Linear(2, 64)
        self.layer2 = nn.Linear(64, 64)
        self.layer3 = nn.Linear(64, 32)
        self.layer4 = nn.Linear(32, 1)
        self.activation = nn.Tanh()
    
    def forward(self, x):
        x = self.activation(self.layer1(x))
        x = self.activation(self.layer2(x))
        x = self.activation(self.layer3(x))
        x = self.layer4(x)
        return x

# Initialize model
model = ScientificNN()
criterion = nn.MSELoss()  # Mean Squared Error loss
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Count parameters
total_params = sum(p.numel() for p in model.parameters())
print(f"Model architecture: 2 -> 64 -> 64 -> 32 -> 1")
print(f"Total parameters: {total_params}")
print(f"Activation function: Tanh")
print(f"Optimizer: Adam (lr=0.001)")

# ============================================================
# 3. Training Loop with Loss Minimization
# ============================================================
print("\n[3] Training neural network (loss minimization)...")

# Prepare training data
X_train = torch.FloatTensor(np.column_stack([x_train, y_train]))
y_train_tensor = torch.FloatTensor(z_train).reshape(-1, 1)

# Training parameters
n_epochs = 1000
batch_size = 128
n_batches = len(X_train) // batch_size

# Storage for tracking
loss_history = []
epoch_times = []

start_time = time()

for epoch in range(n_epochs):
    epoch_start = time()
    epoch_loss = 0.0
    
    # Mini-batch gradient descent
    indices = torch.randperm(len(X_train))
    for i in range(n_batches):
        batch_indices = indices[i*batch_size:(i+1)*batch_size]
        X_batch = X_train[batch_indices]
        y_batch = y_train_tensor[batch_indices]
        
        # Forward pass
        predictions = model(X_batch)
        loss = criterion(predictions, y_batch)
        
        # Backward pass and optimization
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        epoch_loss += loss.item()
    
    avg_loss = epoch_loss / n_batches
    loss_history.append(avg_loss)
    epoch_times.append(time() - epoch_start)
    
    if (epoch + 1) % 100 == 0:
        print(f"Epoch [{epoch+1}/{n_epochs}], Loss: {avg_loss:.6f}, Time: {epoch_times[-1]:.3f}s")

total_time = time() - start_time
print(f"\nTraining completed in {total_time:.2f}s")
print(f"Final loss: {loss_history[-1]:.6f}")
print(f"Average epoch time: {np.mean(epoch_times):.3f}s")

# ============================================================
# 4. Model Evaluation
# ============================================================
print("\n[4] Evaluating trained model...")

model.eval()
with torch.no_grad():
    # Predict on grid
    X_grid_flat = np.column_stack([X_grid.ravel(), Y_grid.ravel()])
    X_grid_tensor = torch.FloatTensor(X_grid_flat)
    Z_pred_flat = model(X_grid_tensor).numpy()
    Z_pred = Z_pred_flat.reshape(X_grid.shape)
    
    # Calculate errors
    absolute_error = np.abs(Z_true - Z_pred)
    relative_error = absolute_error / (np.abs(Z_true) + 1e-8)
    
    print(f"Mean Absolute Error: {np.mean(absolute_error):.6f}")
    print(f"Max Absolute Error: {np.max(absolute_error):.6f}")
    print(f"Mean Relative Error: {np.mean(relative_error)*100:.2f}%")
    print(f"R² Score: {1 - np.sum(absolute_error**2)/np.sum((Z_true - np.mean(Z_true))**2):.6f}")

# ============================================================
# 5. Visualization
# ============================================================
print("\n[5] Creating visualizations...")

fig = plt.figure(figsize=(18, 12))

# Plot 1: Loss Curve
ax1 = plt.subplot(2, 3, 1)
ax1.plot(loss_history, linewidth=2, color='blue')
ax1.set_xlabel('Epoch', fontsize=12)
ax1.set_ylabel('Loss (MSE)', fontsize=12)
ax1.set_title('Loss Function Minimization', fontsize=14, fontweight='bold')
ax1.grid(True, alpha=0.3)
ax1.set_yscale('log')

# Plot 2: Loss Gradient (derivative)
ax2 = plt.subplot(2, 3, 2)
loss_gradient = np.gradient(loss_history)
ax2.plot(loss_gradient, linewidth=2, color='red')
ax2.set_xlabel('Epoch', fontsize=12)
ax2.set_ylabel('Loss Gradient', fontsize=12)
ax2.set_title('Loss Function Gradient (∂L/∂epoch)', fontsize=14, fontweight='bold')
ax2.grid(True, alpha=0.3)
ax2.axhline(y=0, color='k', linestyle='--', alpha=0.5)

# Plot 3: Training Time per Epoch
ax3 = plt.subplot(2, 3, 3)
ax3.plot(epoch_times, linewidth=2, color='green')
ax3.set_xlabel('Epoch', fontsize=12)
ax3.set_ylabel('Time (seconds)', fontsize=12)
ax3.set_title('Training Time per Epoch', fontsize=14, fontweight='bold')
ax3.grid(True, alpha=0.3)

# Plot 4: 3D True Function
ax4 = fig.add_subplot(2, 3, 4, projection='3d')
surf1 = ax4.plot_surface(X_grid, Y_grid, Z_true, cmap='viridis', alpha=0.9)
ax4.set_xlabel('x', fontsize=10)
ax4.set_ylabel('y', fontsize=10)
ax4.set_zlabel('z', fontsize=10)
ax4.set_title('True Scientific Function', fontsize=14, fontweight='bold')
ax4.view_init(elev=25, azim=45)

# Plot 5: 3D Predicted Function
ax5 = fig.add_subplot(2, 3, 5, projection='3d')
surf2 = ax5.plot_surface(X_grid, Y_grid, Z_pred, cmap='viridis', alpha=0.9)
ax5.set_xlabel('x', fontsize=10)
ax5.set_ylabel('y', fontsize=10)
ax5.set_zlabel('z', fontsize=10)
ax5.set_title('Neural Network Prediction', fontsize=14, fontweight='bold')
ax5.view_init(elev=25, azim=45)

# Plot 6: 3D Error Distribution
ax6 = fig.add_subplot(2, 3, 6, projection='3d')
surf3 = ax6.plot_surface(X_grid, Y_grid, absolute_error, cmap='hot', alpha=0.9)
ax6.set_xlabel('x', fontsize=10)
ax6.set_ylabel('y', fontsize=10)
ax6.set_zlabel('|Error|', fontsize=10)
ax6.set_title('Absolute Error Distribution', fontsize=14, fontweight='bold')
ax6.view_init(elev=25, azim=45)
plt.colorbar(surf3, ax=ax6, shrink=0.5)

plt.tight_layout()
plt.savefig('neural_network_loss_minimization.png', dpi=300, bbox_inches='tight')
print("\nVisualization saved as 'neural_network_loss_minimization.png'")
plt.show()

# ============================================================
# 6. Additional Analysis: Loss Landscape
# ============================================================
print("\n[6] Analyzing loss landscape...")

fig2 = plt.figure(figsize=(16, 5))

# Plot 1: 2D Contour - True Function
ax1 = plt.subplot(1, 3, 1)
contour1 = ax1.contourf(X_grid, Y_grid, Z_true, levels=20, cmap='viridis')
ax1.set_xlabel('x', fontsize=12)
ax1.set_ylabel('y', fontsize=12)
ax1.set_title('True Function Contour', fontsize=14, fontweight='bold')
plt.colorbar(contour1, ax=ax1)

# Plot 2: 2D Contour - Predicted Function
ax2 = plt.subplot(1, 3, 2)
contour2 = ax2.contourf(X_grid, Y_grid, Z_pred, levels=20, cmap='viridis')
ax2.set_xlabel('x', fontsize=12)
ax2.set_ylabel('y', fontsize=12)
ax2.set_title('Predicted Function Contour', fontsize=14, fontweight='bold')
plt.colorbar(contour2, ax=ax2)

# Plot 3: 2D Contour - Error
ax3 = plt.subplot(1, 3, 3)
contour3 = ax3.contourf(X_grid, Y_grid, absolute_error, levels=20, cmap='hot')
ax3.set_xlabel('x', fontsize=12)
ax3.set_ylabel('y', fontsize=12)
ax3.set_title('Absolute Error Contour', fontsize=14, fontweight='bold')
plt.colorbar(contour3, ax=ax3)

plt.tight_layout()
plt.savefig('loss_landscape_analysis.png', dpi=300, bbox_inches='tight')
print("Loss landscape saved as 'loss_landscape_analysis.png'")
plt.show()

print("\n" + "="*60)
print("Analysis Complete!")
print("="*60)

def scientific_function(x, y):
    return np.sin(2*np.pi*x) * np.cos(2*np.pi*y) * np.exp(-(x**2 + y**2))

class ScientificNN(nn.Module):
    def __init__(self):
        self.layer1 = nn.Linear(2, 64)
        self.layer2 = nn.Linear(64, 64)
        self.layer3 = nn.Linear(64, 32)
        self.layer4 = nn.Linear(32, 1)
        self.activation = nn.Tanh()

criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(n_epochs):
    indices = torch.randperm(len(X_train))
    for i in range(n_batches):
        predictions = model(X_batch)
        loss = criterion(predictions, y_batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.optimize import differential_evolution
import warnings
warnings.filterwarnings('ignore')

# Set random seed for reproducibility
np.random.seed(42)

print("=" * 70)
print("BATTERY MATERIAL CHARGING CHARACTERISTICS OPTIMIZATION")
print("Trade-off Analysis: Charging Speed vs Lifespan vs Capacity")
print("=" * 70)
print()

# ============================================================================
# OBJECTIVE FUNCTIONS DEFINITION
# ============================================================================

def charging_time(x1, x2):
    """
    Charging time (minutes) - MINIMIZE
    x1: C-rate (charging rate)
    x2: Active material loading (mg/cm^2)
    """
    return (100 / x1) * (1 + 0.02 * x2)

def battery_lifespan(x1, x2):
    """
    Battery lifespan (cycles) - MAXIMIZE (return negative for minimization)
    x1: C-rate (charging rate)
    x2: Active material loading (mg/cm^2)
    """
    return -(1000 - 50 * (x1 ** 1.5) - 10 * ((x2 - 15) ** 2))

def energy_capacity(x1, x2):
    """
    Energy capacity (mAh/cm^2) - MAXIMIZE (return negative for minimization)
    x1: C-rate (charging rate)
    x2: Active material loading (mg/cm^2)
    """
    return -(150 * x2 * np.exp(-0.1 * x1) - 0.5 * (x2 ** 2))

# ============================================================================
# MULTI-OBJECTIVE OPTIMIZATION USING NSGA-II APPROACH
# ============================================================================

class BatteryOptimizer:
    """
    Multi-objective optimizer for battery charging characteristics
    Uses weighted sum approach with multiple weight combinations
    """
    
    def __init__(self):
        self.bounds = [(0.5, 5.0), (5, 25)]  # [C-rate, Loading]
        self.pareto_solutions = []
        
    def weighted_objective(self, x, weights):
        """
        Weighted sum of normalized objectives
        """
        x1, x2 = x
        
        # Calculate objectives
        f1 = charging_time(x1, x2)
        f2 = battery_lifespan(x1, x2)
        f3 = energy_capacity(x1, x2)
        
        # Normalize (approximate ranges)
        f1_norm = f1 / 200.0  # Charging time up to ~200 min
        f2_norm = f2 / 1000.0  # Lifespan range ~1000
        f3_norm = f3 / 2000.0  # Capacity range ~2000
        
        # Weighted sum
        return weights[0] * f1_norm + weights[1] * f2_norm + weights[2] * f3_norm
    
    def optimize_with_weights(self, weights):
        """
        Optimize with specific weight combination
        """
        result = differential_evolution(
            lambda x: self.weighted_objective(x, weights),
            self.bounds,
            maxiter=300,
            popsize=20,
            seed=42,
            atol=1e-6,
            tol=1e-6
        )
        return result.x
    
    def generate_pareto_front(self, n_points=50):
        """
        Generate Pareto front by varying weights
        """
        print("Generating Pareto-optimal solutions...")
        print(f"Running {n_points} optimization scenarios...")
        print()
        
        solutions = []
        
        # Generate diverse weight combinations
        for i in range(n_points):
            # Random weights that sum to 1
            w = np.random.dirichlet([1, 1, 1])
            
            try:
                x_opt = self.optimize_with_weights(w)
                x1, x2 = x_opt
                
                # Calculate actual objective values
                f1 = charging_time(x1, x2)
                f2_actual = -battery_lifespan(x1, x2)  # Convert back to positive
                f3_actual = -energy_capacity(x1, x2)   # Convert back to positive
                
                solutions.append({
                    'x1': x1,
                    'x2': x2,
                    'charging_time': f1,
                    'lifespan': f2_actual,
                    'capacity': f3_actual,
                    'weights': w
                })
                
            except Exception as e:
                continue
        
        self.pareto_solutions = solutions
        print(f"✓ Generated {len(solutions)} Pareto-optimal solutions")
        print()
        
        return solutions

# ============================================================================
# OPTIMIZATION EXECUTION
# ============================================================================

optimizer = BatteryOptimizer()
pareto_solutions = optimizer.generate_pareto_front(n_points=60)

# Extract solution data
x1_vals = [s['x1'] for s in pareto_solutions]
x2_vals = [s['x2'] for s in pareto_solutions]
time_vals = [s['charging_time'] for s in pareto_solutions]
life_vals = [s['lifespan'] for s in pareto_solutions]
capacity_vals = [s['capacity'] for s in pareto_solutions]

# ============================================================================
# IDENTIFY KEY SOLUTIONS
# ============================================================================

print("=" * 70)
print("KEY PARETO-OPTIMAL SOLUTIONS")
print("=" * 70)
print()

# Find extreme solutions
idx_fast = np.argmin(time_vals)
idx_long_life = np.argmax(life_vals)
idx_high_cap = np.argmax(capacity_vals)

# Find balanced solution (closest to center of normalized space)
def normalize(arr):
    return (arr - np.min(arr)) / (np.max(arr) - np.min(arr))

time_norm = normalize(np.array(time_vals))
life_norm = 1 - normalize(np.array(life_vals))  # Invert (want high)
cap_norm = 1 - normalize(np.array(capacity_vals))  # Invert (want high)

distances = np.sqrt(time_norm**2 + life_norm**2 + cap_norm**2)
idx_balanced = np.argmin(distances)

solutions_of_interest = [
    ('Fastest Charging', idx_fast),
    ('Longest Lifespan', idx_long_life),
    ('Highest Capacity', idx_high_cap),
    ('Balanced Solution', idx_balanced)
]

for name, idx in solutions_of_interest:
    sol = pareto_solutions[idx]
    print(f"📊 {name}:")
    print(f"   C-rate (x₁):        {sol['x1']:.3f}")
    print(f"   Loading (x₂):       {sol['x2']:.3f} mg/cm²")
    print(f"   Charging Time:      {sol['charging_time']:.2f} minutes")
    print(f"   Battery Lifespan:   {sol['lifespan']:.0f} cycles")
    print(f"   Energy Capacity:    {sol['capacity']:.2f} mAh/cm²")
    print()

# ============================================================================
# VISUALIZATION
# ============================================================================

print("=" * 70)
print("GENERATING VISUALIZATIONS")
print("=" * 70)
print()

fig = plt.figure(figsize=(20, 12))

# ========== 3D Pareto Front ==========
ax1 = fig.add_subplot(2, 3, 1, projection='3d')
scatter = ax1.scatter(time_vals, life_vals, capacity_vals, 
                     c=capacity_vals, cmap='viridis', s=50, alpha=0.6)
ax1.set_xlabel('Charging Time (min)', fontsize=10, labelpad=10)
ax1.set_ylabel('Battery Lifespan (cycles)', fontsize=10, labelpad=10)
ax1.set_zlabel('Energy Capacity (mAh/cm²)', fontsize=10, labelpad=10)
ax1.set_title('3D Pareto Front: Trade-off Surface', fontsize=12, fontweight='bold', pad=20)
plt.colorbar(scatter, ax=ax1, label='Capacity (mAh/cm²)', shrink=0.5)
ax1.view_init(elev=20, azim=45)

# Mark key solutions
for name, idx in solutions_of_interest:
    sol = pareto_solutions[idx]
    ax1.scatter([sol['charging_time']], [sol['lifespan']], [sol['capacity']], 
               c='red', s=200, marker='*', edgecolors='black', linewidths=2)

# ========== 2D Projections ==========
# Time vs Lifespan
ax2 = fig.add_subplot(2, 3, 2)
scatter2 = ax2.scatter(time_vals, life_vals, c=capacity_vals, 
                       cmap='plasma', s=50, alpha=0.6)
ax2.set_xlabel('Charging Time (min)', fontsize=10)
ax2.set_ylabel('Battery Lifespan (cycles)', fontsize=10)
ax2.set_title('Time vs Lifespan Trade-off', fontsize=11, fontweight='bold')
ax2.grid(True, alpha=0.3)
plt.colorbar(scatter2, ax=ax2, label='Capacity')

for name, idx in solutions_of_interest[:2]:
    sol = pareto_solutions[idx]
    ax2.scatter([sol['charging_time']], [sol['lifespan']], 
               c='red', s=150, marker='*', edgecolors='black', linewidths=1.5)

# Time vs Capacity
ax3 = fig.add_subplot(2, 3, 3)
scatter3 = ax3.scatter(time_vals, capacity_vals, c=life_vals, 
                       cmap='coolwarm', s=50, alpha=0.6)
ax3.set_xlabel('Charging Time (min)', fontsize=10)
ax3.set_ylabel('Energy Capacity (mAh/cm²)', fontsize=10)
ax3.set_title('Time vs Capacity Trade-off', fontsize=11, fontweight='bold')
ax3.grid(True, alpha=0.3)
plt.colorbar(scatter3, ax=ax3, label='Lifespan')

# ========== Design Variable Space ==========
ax4 = fig.add_subplot(2, 3, 4)
scatter4 = ax4.scatter(x1_vals, x2_vals, c=time_vals, 
                       cmap='RdYlGn_r', s=60, alpha=0.7)
ax4.set_xlabel('C-rate (x₁)', fontsize=10)
ax4.set_ylabel('Active Material Loading (x₂) [mg/cm²]', fontsize=10)
ax4.set_title('Design Variable Space (colored by Charging Time)', fontsize=11, fontweight='bold')
ax4.grid(True, alpha=0.3)
plt.colorbar(scatter4, ax=ax4, label='Time (min)')

# ========== 3D Design Space ==========
ax5 = fig.add_subplot(2, 3, 5, projection='3d')
scatter5 = ax5.scatter(x1_vals, x2_vals, time_vals, 
                      c=life_vals, cmap='viridis', s=50, alpha=0.6)
ax5.set_xlabel('C-rate (x₁)', fontsize=10, labelpad=8)
ax5.set_ylabel('Loading (x₂) [mg/cm²]', fontsize=10, labelpad=8)
ax5.set_zlabel('Charging Time (min)', fontsize=10, labelpad=8)
ax5.set_title('Design Space: 3D View', fontsize=11, fontweight='bold', pad=15)
plt.colorbar(scatter5, ax=ax5, label='Lifespan', shrink=0.5)
ax5.view_init(elev=25, azim=135)

# ========== Objective Correlations ==========
ax6 = fig.add_subplot(2, 3, 6)
objectives = np.array([time_vals, life_vals, capacity_vals])
correlation = np.corrcoef(objectives)

im = ax6.imshow(correlation, cmap='RdBu_r', vmin=-1, vmax=1, aspect='auto')
ax6.set_xticks([0, 1, 2])
ax6.set_yticks([0, 1, 2])
ax6.set_xticklabels(['Time', 'Lifespan', 'Capacity'])
ax6.set_yticklabels(['Time', 'Lifespan', 'Capacity'])
ax6.set_title('Objective Function Correlations', fontsize=11, fontweight='bold')

for i in range(3):
    for j in range(3):
        text = ax6.text(j, i, f'{correlation[i, j]:.2f}',
                       ha="center", va="center", color="black", fontsize=11)

plt.colorbar(im, ax=ax6, label='Correlation')

plt.tight_layout()
plt.savefig('battery_optimization_results.png', dpi=150, bbox_inches='tight')
print("✓ Visualization complete")
print()

# ============================================================================
# ADDITIONAL ANALYSIS
# ============================================================================

print("=" * 70)
print("STATISTICAL ANALYSIS OF PARETO SOLUTIONS")
print("=" * 70)
print()

print(f"Charging Time:      {np.min(time_vals):.2f} - {np.max(time_vals):.2f} min")
print(f"Battery Lifespan:   {np.min(life_vals):.0f} - {np.max(life_vals):.0f} cycles")
print(f"Energy Capacity:    {np.min(capacity_vals):.2f} - {np.max(capacity_vals):.2f} mAh/cm²")
print()
print(f"C-rate Range:       {np.min(x1_vals):.3f} - {np.max(x1_vals):.3f}")
print(f"Loading Range:      {np.min(x2_vals):.2f} - {np.max(x2_vals):.2f} mg/cm²")
print()

print("=" * 70)
print("OPTIMIZATION COMPLETE")
print("=" * 70)

plt.show()

w = np.random.dirichlet([1, 1, 1])

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import time

# Set random seed for reproducibility
np.random.seed(42)

# ==================== Rastrigin Function ====================
def rastrigin(x):
    """
    Rastrigin function: f(x) = 10n + sum(x_i^2 - 10*cos(2*pi*x_i))
    Global minimum at x = (0, 0, ..., 0) with f(x) = 0
    """
    n = len(x)
    return 10 * n + np.sum(x**2 - 10 * np.cos(2 * np.pi * x))

# ==================== Genetic Algorithm Implementation ====================
class GeneticAlgorithm:
    def __init__(self, pop_size, dimensions, bounds, mutation_rate, 
                 tournament_size, generations):
        self.pop_size = pop_size
        self.dimensions = dimensions
        self.bounds = bounds
        self.mutation_rate = mutation_rate
        self.tournament_size = tournament_size
        self.generations = generations
        
        # Initialize population randomly within bounds
        self.population = np.random.uniform(
            bounds[0], bounds[1], (pop_size, dimensions)
        )
        
        # History tracking
        self.best_fitness_history = []
        self.avg_fitness_history = []
        self.diversity_history = []
        self.best_solution = None
        self.best_fitness = float('inf')
    
    def evaluate_population(self):
        """Evaluate fitness for entire population"""
        return np.array([rastrigin(ind) for ind in self.population])
    
    def calculate_diversity(self):
        """
        Calculate population diversity as average pairwise distance
        Diversity = (1/n(n-1)) * sum(||x_i - x_j||)
        """
        n = len(self.population)
        total_distance = 0
        count = 0
        
        for i in range(n):
            for j in range(i + 1, n):
                distance = np.linalg.norm(self.population[i] - self.population[j])
                total_distance += distance
                count += 1
        
        return total_distance / count if count > 0 else 0
    
    def tournament_selection(self, fitness):
        """
        Tournament selection: randomly select tournament_size individuals
        and return the best one
        """
        tournament_indices = np.random.choice(
            self.pop_size, self.tournament_size, replace=False
        )
        tournament_fitness = fitness[tournament_indices]
        winner_idx = tournament_indices[np.argmin(tournament_fitness)]
        return self.population[winner_idx].copy()
    
    def crossover(self, parent1, parent2):
        """
        Simulated Binary Crossover (SBX)
        Creates two offspring from two parents
        """
        eta = 20  # Distribution index
        offspring1 = np.zeros(self.dimensions)
        offspring2 = np.zeros(self.dimensions)
        
        for i in range(self.dimensions):
            if np.random.rand() < 0.5:
                if abs(parent1[i] - parent2[i]) > 1e-14:
                    if parent1[i] < parent2[i]:
                        y1, y2 = parent1[i], parent2[i]
                    else:
                        y1, y2 = parent2[i], parent1[i]
                    
                    beta = 1.0 + (2.0 * (y1 - self.bounds[0]) / (y2 - y1))
                    alpha = 2.0 - beta ** -(eta + 1.0)
                    
                    rand = np.random.rand()
                    if rand <= 1.0 / alpha:
                        betaq = (rand * alpha) ** (1.0 / (eta + 1.0))
                    else:
                        betaq = (1.0 / (2.0 - rand * alpha)) ** (1.0 / (eta + 1.0))
                    
                    offspring1[i] = 0.5 * ((y1 + y2) - betaq * (y2 - y1))
                    offspring2[i] = 0.5 * ((y1 + y2) + betaq * (y2 - y1))
                else:
                    offspring1[i] = parent1[i]
                    offspring2[i] = parent2[i]
            else:
                offspring1[i] = parent1[i]
                offspring2[i] = parent2[i]
        
        return offspring1, offspring2
    
    def mutate(self, individual):
        """
        Polynomial mutation
        """
        eta_m = 20  # Distribution index for mutation
        
        for i in range(self.dimensions):
            if np.random.rand() < self.mutation_rate:
                y = individual[i]
                delta1 = (y - self.bounds[0]) / (self.bounds[1] - self.bounds[0])
                delta2 = (self.bounds[1] - y) / (self.bounds[1] - self.bounds[0])
                
                rand = np.random.rand()
                mut_pow = 1.0 / (eta_m + 1.0)
                
                if rand < 0.5:
                    xy = 1.0 - delta1
                    val = 2.0 * rand + (1.0 - 2.0 * rand) * (xy ** (eta_m + 1.0))
                    deltaq = val ** mut_pow - 1.0
                else:
                    xy = 1.0 - delta2
                    val = 2.0 * (1.0 - rand) + 2.0 * (rand - 0.5) * (xy ** (eta_m + 1.0))
                    deltaq = 1.0 - val ** mut_pow
                
                y = y + deltaq * (self.bounds[1] - self.bounds[0])
                individual[i] = np.clip(y, self.bounds[0], self.bounds[1])
        
        return individual
    
    def evolve(self):
        """Main evolution loop"""
        for generation in range(self.generations):
            # Evaluate fitness
            fitness = self.evaluate_population()
            
            # Track statistics
            best_gen_fitness = np.min(fitness)
            avg_gen_fitness = np.mean(fitness)
            diversity = self.calculate_diversity()
            
            self.best_fitness_history.append(best_gen_fitness)
            self.avg_fitness_history.append(avg_gen_fitness)
            self.diversity_history.append(diversity)
            
            # Update best solution
            if best_gen_fitness < self.best_fitness:
                self.best_fitness = best_gen_fitness
                self.best_solution = self.population[np.argmin(fitness)].copy()
            
            # Create new population
            new_population = []
            
            # Elitism: keep the best individual
            best_idx = np.argmin(fitness)
            new_population.append(self.population[best_idx].copy())
            
            # Generate offspring
            while len(new_population) < self.pop_size:
                # Selection
                parent1 = self.tournament_selection(fitness)
                parent2 = self.tournament_selection(fitness)
                
                # Crossover
                offspring1, offspring2 = self.crossover(parent1, parent2)
                
                # Mutation
                offspring1 = self.mutate(offspring1)
                offspring2 = self.mutate(offspring2)
                
                # Add to new population
                new_population.append(offspring1)
                if len(new_population) < self.pop_size:
                    new_population.append(offspring2)
            
            self.population = np.array(new_population[:self.pop_size])
        
        return self.best_solution, self.best_fitness

# ==================== Experiment: Parameter Comparison ====================
def run_experiments():
    """Run GA with different parameter settings"""
    
    dimensions = 2
    bounds = (-5.12, 5.12)
    generations = 100
    
    # Experiment configurations
    experiments = {
        'Small Pop + Low Mutation': {
            'pop_size': 30, 'mutation_rate': 0.01, 'tournament_size': 2
        },
        'Small Pop + High Mutation': {
            'pop_size': 30, 'mutation_rate': 0.2, 'tournament_size': 2
        },
        'Large Pop + Low Mutation': {
            'pop_size': 100, 'mutation_rate': 0.01, 'tournament_size': 2
        },
        'Large Pop + High Mutation': {
            'pop_size': 100, 'mutation_rate': 0.2, 'tournament_size': 2
        },
        'Balanced (Recommended)': {
            'pop_size': 60, 'mutation_rate': 0.1, 'tournament_size': 3
        },
    }
    
    results = {}
    
    print("=" * 70)
    print("GENETIC ALGORITHM PARAMETER OPTIMIZATION EXPERIMENT")
    print("=" * 70)
    print(f"\nObjective: Minimize Rastrigin function in {dimensions}D")
    print(f"Search space: [{bounds[0]}, {bounds[1]}]^{dimensions}")
    print(f"Global optimum: f(0, 0) = 0")
    print(f"\nGenerations: {generations}")
    print("=" * 70)
    
    for name, params in experiments.items():
        print(f"\n[{name}]")
        print(f"  Population: {params['pop_size']}, "
              f"Mutation Rate: {params['mutation_rate']}, "
              f"Tournament Size: {params['tournament_size']}")
        
        start_time = time.time()
        
        ga = GeneticAlgorithm(
            pop_size=params['pop_size'],
            dimensions=dimensions,
            bounds=bounds,
            mutation_rate=params['mutation_rate'],
            tournament_size=params['tournament_size'],
            generations=generations
        )
        
        best_sol, best_fit = ga.evolve()
        elapsed_time = time.time() - start_time
        
        results[name] = {
            'ga': ga,
            'best_solution': best_sol,
            'best_fitness': best_fit,
            'time': elapsed_time
        }
        
        print(f"  Best fitness: {best_fit:.6f}")
        print(f"  Best solution: ({best_sol[0]:.4f}, {best_sol[1]:.4f})")
        print(f"  Time: {elapsed_time:.2f}s")
    
    print("\n" + "=" * 70)
    
    return results

# ==================== Visualization ====================
def create_visualizations(results):
    """Create comprehensive visualizations"""
    
    fig = plt.figure(figsize=(20, 12))
    
    # 1. Fitness convergence comparison
    ax1 = plt.subplot(2, 3, 1)
    for name, data in results.items():
        ga = data['ga']
        ax1.plot(ga.best_fitness_history, label=name, linewidth=2)
    ax1.set_xlabel('Generation', fontsize=11)
    ax1.set_ylabel('Best Fitness', fontsize=11)
    ax1.set_title('Convergence Speed Comparison', fontsize=12, fontweight='bold')
    ax1.legend(fontsize=9)
    ax1.grid(True, alpha=0.3)
    ax1.set_yscale('log')
    
    # 2. Diversity evolution
    ax2 = plt.subplot(2, 3, 2)
    for name, data in results.items():
        ga = data['ga']
        ax2.plot(ga.diversity_history, label=name, linewidth=2)
    ax2.set_xlabel('Generation', fontsize=11)
    ax2.set_ylabel('Population Diversity', fontsize=11)
    ax2.set_title('Diversity Maintenance', fontsize=12, fontweight='bold')
    ax2.legend(fontsize=9)
    ax2.grid(True, alpha=0.3)
    
    # 3. Average fitness
    ax3 = plt.subplot(2, 3, 3)
    for name, data in results.items():
        ga = data['ga']
        ax3.plot(ga.avg_fitness_history, label=name, linewidth=2)
    ax3.set_xlabel('Generation', fontsize=11)
    ax3.set_ylabel('Average Fitness', fontsize=11)
    ax3.set_title('Population Average Fitness', fontsize=12, fontweight='bold')
    ax3.legend(fontsize=9)
    ax3.grid(True, alpha=0.3)
    ax3.set_yscale('log')
    
    # 4. Final results bar chart
    ax4 = plt.subplot(2, 3, 4)
    names = list(results.keys())
    final_fitness = [results[name]['best_fitness'] for name in names]
    colors = plt.cm.viridis(np.linspace(0, 1, len(names)))
    bars = ax4.bar(range(len(names)), final_fitness, color=colors)
    ax4.set_xticks(range(len(names)))
    ax4.set_xticklabels(names, rotation=45, ha='right', fontsize=9)
    ax4.set_ylabel('Final Best Fitness', fontsize=11)
    ax4.set_title('Final Performance Comparison', fontsize=12, fontweight='bold')
    ax4.grid(True, alpha=0.3, axis='y')
    
    # Add value labels on bars
    for i, (bar, val) in enumerate(zip(bars, final_fitness)):
        ax4.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.1, 
                f'{val:.3f}', ha='center', va='bottom', fontsize=9)
    
    # 5. 3D Rastrigin surface with solutions
    ax5 = plt.subplot(2, 3, 5, projection='3d')
    
    # Create mesh for Rastrigin function
    x = np.linspace(-5.12, 5.12, 100)
    y = np.linspace(-5.12, 5.12, 100)
    X, Y = np.meshgrid(x, y)
    Z = np.zeros_like(X)
    
    for i in range(X.shape[0]):
        for j in range(X.shape[1]):
            Z[i, j] = rastrigin(np.array([X[i, j], Y[i, j]]))
    
    # Plot surface
    surf = ax5.plot_surface(X, Y, Z, cmap=cm.viridis, alpha=0.6, 
                           linewidth=0, antialiased=True)
    
    # Plot best solutions
    colors_3d = plt.cm.Set1(np.linspace(0, 1, len(results)))
    for idx, (name, data) in enumerate(results.items()):
        sol = data['best_solution']
        fit = data['best_fitness']
        ax5.scatter([sol[0]], [sol[1]], [fit], color=colors_3d[idx], 
                   s=100, marker='o', edgecolors='black', linewidths=2,
                   label=name)
    
    ax5.set_xlabel('X₁', fontsize=10)
    ax5.set_ylabel('X₂', fontsize=10)
    ax5.set_zlabel('f(X)', fontsize=10)
    ax5.set_title('Rastrigin Function with Best Solutions', 
                  fontsize=12, fontweight='bold')
    ax5.legend(fontsize=8, loc='upper left')
    
    # 6. Convergence-Diversity trade-off
    ax6 = plt.subplot(2, 3, 6)
    for name, data in results.items():
        ga = data['ga']
        # Plot final 20% of evolution
        start_idx = int(0.8 * len(ga.best_fitness_history))
        ax6.scatter(ga.diversity_history[start_idx:], 
                   ga.best_fitness_history[start_idx:],
                   label=name, alpha=0.6, s=30)
    ax6.set_xlabel('Population Diversity', fontsize=11)
    ax6.set_ylabel('Best Fitness', fontsize=11)
    ax6.set_title('Convergence-Diversity Trade-off (Final 20%)', 
                  fontsize=12, fontweight='bold')
    ax6.legend(fontsize=9)
    ax6.grid(True, alpha=0.3)
    ax6.set_yscale('log')
    
    plt.tight_layout()
    plt.savefig('ga_parameter_optimization.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    # Additional 3D visualization: Population evolution
    fig2 = plt.figure(figsize=(15, 5))
    
    for idx, (name, data) in enumerate(list(results.items())[:3]):
        ax = fig2.add_subplot(1, 3, idx+1, projection='3d')
        ga = data['ga']
        
        # Plot surface
        x = np.linspace(-5.12, 5.12, 50)
        y = np.linspace(-5.12, 5.12, 50)
        X, Y = np.meshgrid(x, y)
        Z = np.zeros_like(X)
        
        for i in range(X.shape[0]):
            for j in range(X.shape[1]):
                Z[i, j] = rastrigin(np.array([X[i, j], Y[i, j]]))
        
        ax.plot_surface(X, Y, Z, cmap=cm.viridis, alpha=0.3, 
                       linewidth=0, antialiased=True)
        
        # Plot final population
        pop = ga.population
        fitness = np.array([rastrigin(ind) for ind in pop])
        ax.scatter(pop[:, 0], pop[:, 1], fitness, c='red', 
                  marker='o', s=20, alpha=0.6)
        
        # Plot best solution
        sol = data['best_solution']
        fit = data['best_fitness']
        ax.scatter([sol[0]], [sol[1]], [fit], color='gold', 
                  s=200, marker='*', edgecolors='black', linewidths=2)
        
        ax.set_xlabel('X₁', fontsize=10)
        ax.set_ylabel('X₂', fontsize=10)
        ax.set_zlabel('f(X)', fontsize=10)
        ax.set_title(f'{name}\nFinal Population', fontsize=11, fontweight='bold')
    
    plt.tight_layout()
    plt.savefig('ga_final_populations.png', dpi=150, bbox_inches='tight')
    plt.show()

# ==================== Main Execution ====================
if __name__ == "__main__":
    # Run experiments
    results = run_experiments()
    
    # Create visualizations
    print("\n[Generating visualizations...]")
    create_visualizations(results)
    
    print("\n[Analysis Complete]")
    print("\nKey Insights:")
    print("1. Population size affects exploration capability")
    print("2. Mutation rate controls diversity maintenance")
    print("3. Tournament size determines selection pressure")
    print("4. Balance is crucial for optimal performance")