import os
os.environ['MKL_DISABLE_FAST_MM'] = '1'
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from IPython.display import HTML # For displaying the animation in Jupyter
import matplotlib as mpl

# Set the animation embed limit to 50 MB (or your desired size in MB)
mpl.rcParams['animation.embed_limit'] = 200  # The value is in Megabytes

# Parameters for the simulation
NUM_BOIDS = 150
WIDTH = 600  # Width of the simulation area
HEIGHT = 400 # Height of the simulation area
MAX_SPEED = 5.0
MIN_SPEED = 2.0
PERCEPTION_RADIUS = 75.0 # How far a boid can see
AVOIDANCE_RADIUS = 15.0  # Radius to avoid other boids

# Weights for the three rules
COHESION_FACTOR = 0.05
ALIGNMENT_FACTOR = 0.5
SEPARATION_FACTOR = 0.15
TURN_FACTOR = 0.2 # How quickly boids can turn

# Boundaries for screen wrapping
MARGIN = 50

class Boid:
    def __init__(self, x, y):
        self.position = np.array([float(x), float(y)])
        angle = np.random.uniform(0, 2 * np.pi)
        self.velocity = np.array([np.cos(angle), np.sin(angle)]) * MAX_SPEED
        self.acceleration = np.zeros(2, dtype=float)

    def update(self):
        # Apply acceleration
        self.velocity += self.acceleration
        
        # Limit speed
        speed = np.linalg.norm(self.velocity)
        if speed > MAX_SPEED:
            self.velocity = (self.velocity / speed) * MAX_SPEED
        elif speed < MIN_SPEED:
             self.velocity = (self.velocity / speed) * MIN_SPEED
            
        self.position += self.velocity
        self.acceleration = np.zeros(2, dtype=float) # Reset acceleration

    def apply_force(self, force):
        self.acceleration += force

    def screen_wrap(self):
        if self.position[0] > WIDTH + MARGIN:
            self.position[0] = -MARGIN
        elif self.position[0] < -MARGIN:
            self.position[0] = WIDTH + MARGIN
        if self.position[1] > HEIGHT + MARGIN:
            self.position[1] = -MARGIN
        elif self.position[1] < -MARGIN:
            self.position[1] = HEIGHT + MARGIN

    def steer(self, target_direction):
        desired = target_direction - self.position
        desired_norm = np.linalg.norm(desired)
        
        if desired_norm > 0:
            desired = (desired / desired_norm) * MAX_SPEED
        
        steer = desired - self.velocity
        steer_norm = np.linalg.norm(steer)
        
        if steer_norm > TURN_FACTOR: # Limit the turning force
            steer = (steer / steer_norm) * TURN_FACTOR
            
        return steer

    def cohesion(self, boids):
        center_of_mass = np.zeros(2, dtype=float)
        count = 0
        for other in boids:
            if other is not self:
                distance = np.linalg.norm(self.position - other.position)
                if 0 < distance < PERCEPTION_RADIUS:
                    center_of_mass += other.position
                    count += 1
        if count > 0:
            center_of_mass /= count
            return self.steer(center_of_mass) * COHESION_FACTOR
        return np.zeros(2, dtype=float)

    def alignment(self, boids):
        avg_velocity = np.zeros(2, dtype=float)
        count = 0
        for other in boids:
            if other is not self:
                distance = np.linalg.norm(self.position - other.position)
                if 0 < distance < PERCEPTION_RADIUS:
                    avg_velocity += other.velocity
                    count += 1
        if count > 0:
            avg_velocity /= count
            # Steer towards the average velocity
            desired = avg_velocity
            desired_norm = np.linalg.norm(desired)
            if desired_norm > 0:
                desired = (desired / desired_norm) * MAX_SPEED
            
            steer = desired - self.velocity
            steer_norm = np.linalg.norm(steer)
            if steer_norm > TURN_FACTOR:
                steer = (steer / steer_norm) * TURN_FACTOR
            return steer * ALIGNMENT_FACTOR
        return np.zeros(2, dtype=float)

    def separation(self, boids):
        steering_force = np.zeros(2, dtype=float)
        count = 0
        for other in boids:
            if other is not self:
                distance = np.linalg.norm(self.position - other.position)
                if 0 < distance < AVOIDANCE_RADIUS: # Use a smaller radius for separation
                    diff = self.position - other.position
                    steering_force += diff / (distance * distance) # Force inversely proportional to square of distance
                    count += 1
        if count > 0:
            steering_force /= count # Average the force
            # Steer away
            desired = steering_force
            desired_norm = np.linalg.norm(desired)
            if desired_norm > 0:
                desired = (desired / desired_norm) * MAX_SPEED
            
            steer = desired - self.velocity
            steer_norm = np.linalg.norm(steer)
            if steer_norm > TURN_FACTOR:
                steer = (steer / steer_norm) * TURN_FACTOR
            return steer * SEPARATION_FACTOR
        return np.zeros(2, dtype=float)

    def apply_rules(self, boids):
        cohesion_force = self.cohesion(boids)
        alignment_force = self.alignment(boids)
        separation_force = self.separation(boids)
        
        self.apply_force(cohesion_force)
        self.apply_force(alignment_force)
        self.apply_force(separation_force)


# Initialize boids
boids = [Boid(np.random.uniform(0, WIDTH), np.random.uniform(0, HEIGHT)) for _ in range(NUM_BOIDS)]

# Matplotlib setup
fig, ax = plt.subplots(figsize=(10, 7.5)) # Adjusted for better aspect ratio
ax.set_xlim(0, WIDTH)
ax.set_ylim(0, HEIGHT)
ax.set_xticks([]) # Remove x-axis ticks
ax.set_yticks([]) # Remove y-axis ticks
ax.set_facecolor('lightgrey') # Background color

# Using quivers to represent boids (arrows showing direction)
# We need to store the quiver objects to update them
quivers = []
for boid in boids:
    # Initial quiver plot for each boid
    # Scale the velocity for better arrow representation
    # The head_width, head_length, etc., might need tweaking based on MAX_SPEED
    q = ax.quiver(boid.position[0], boid.position[1], 
                  boid.velocity[0], boid.velocity[1], 
                  color='blue', scale=50, headwidth=4, headlength=6) # Adjust scale and head
    quivers.append(q)

def update_frame(frame_number):
    # Update positions and velocities of boids
    for boid in boids:
        boid.apply_rules(boids)
        boid.update()
        boid.screen_wrap()

    # Update the quiver plots
    for i, boid in enumerate(boids):
        quivers[i].set_offsets(boid.position)
        # Normalize velocity for direction if MAX_SPEED is large, or scale it down
        # This ensures arrows don't become excessively long
        vel_norm = np.linalg.norm(boid.velocity)
        if vel_norm > 0:
            uv = boid.velocity / vel_norm 
        else:
            uv = np.zeros(2) # In case velocity is zero
        quivers[i].set_UVC(uv[0], uv[1])
        
    return quivers

# Create animation
# interval is the delay between frames in milliseconds.
# Adjust blit=True for potentially smoother animation if it works well on your system.
# blit=False is often more robust.
ani = FuncAnimation(fig, update_frame, frames=400, interval=50, blit=False)

# Display the animation in Jupyter Notebook
plt.close(fig) # Prevent static plot from displaying


# ani.save('boid_simulation.mp4', writer='ffmpeg', fps=30)
HTML(ani.to_jshtml()) # or use ani.to_html5_video()

# Import necessary libraries
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from IPython.display import HTML  # For displaying animation in Jupyter
import random
import math
import time # To track performance if needed
import matplotlib as mpl

# Set the animation embed limit to 50 MB (or your desired size in MB)
mpl.rcParams['animation.embed_limit'] = 600  # The value is in Megabytes

# --- Constants (mostly same as pygame version) ---
WIDTH, HEIGHT = 400, 300 # Logical world dimensions
PHEROMONE_RESOLUTION_FACTOR = 5 # Pheromones calculated on a coarser grid (HIGHER value = COARSER grid = FASTER)
P_WIDTH, P_HEIGHT = WIDTH // PHEROMONE_RESOLUTION_FACTOR, HEIGHT // PHEROMONE_RESOLUTION_FACTOR

N_ANTS = 150
ANT_SPEED = 1.0
ANT_ROTATION_SPEED = 0.1  # Radians per step
SENSOR_DISTANCE = 15     # How far ahead ants "smell"
SENSOR_ANGLE = math.pi / 4 # Angle offset for side sensors (radians)

# Pheromones
EVAPORATION_RATE = 0.99 # Percentage remaining after 1 step
PHEROMONE_DROP_RATE = 2.0
MAX_PHEROMONE = 100.0 # Cap pheromone intensity

# States (using integers for easy indexing/coloring)
SEARCHING = 0
RETURNING = 1

# Colors for matplotlib plotting (Ants)
ANT_COLORS = ['red', 'blue'] # red for searching, blue for returning

# --- Helper Functions (mostly same) ---
def normalize_angle(angle):
    while angle > math.pi: angle -= 2 * math.pi
    while angle <= -math.pi: angle += 2 * math.pi
    return angle

def world_to_pheromone_grid(pos):
    x, y = pos
    px = int(x / PHEROMONE_RESOLUTION_FACTOR)
    py = int(y / PHEROMONE_RESOLUTION_FACTOR)
    # Clamp values to be within grid bounds
    px = max(0, min(P_WIDTH - 1, px))
    py = max(0, min(P_HEIGHT - 1, py))
    return px, py

def is_within_bounds(pos, w=WIDTH, h=HEIGHT):
    x, y = pos
    return 0 <= x < w and 0 <= y < h

# --- Ant Class (logic mostly same, removed drawing) ---
class Ant:
    def __init__(self, nest_pos):
        self.pos = np.array(nest_pos, dtype=float)
        self.angle = random.uniform(0, 2 * math.pi)
        self.state = SEARCHING
        self.speed = ANT_SPEED
        self.nest_pos = np.array(nest_pos, dtype=float)
        self.color = ANT_COLORS[self.state]

    def sense_pheromones(self, pheromone_grid):
        sensor_main_pos = self.pos + SENSOR_DISTANCE * np.array([math.cos(self.angle), math.sin(self.angle)])
        sensor_left_pos = self.pos + SENSOR_DISTANCE * np.array([math.cos(self.angle - SENSOR_ANGLE), math.sin(self.angle - SENSOR_ANGLE)])
        sensor_right_pos = self.pos + SENSOR_DISTANCE * np.array([math.cos(self.angle + SENSOR_ANGLE), math.sin(self.angle + SENSOR_ANGLE)])

        # Read pheromone levels only if sensors are within bounds
        p_main, p_left, p_right = 0, 0, 0
        if is_within_bounds(sensor_main_pos):
            px, py = world_to_pheromone_grid(sensor_main_pos)
            p_main = pheromone_grid[px, py]
        if is_within_bounds(sensor_left_pos):
            px, py = world_to_pheromone_grid(sensor_left_pos)
            p_left = pheromone_grid[px, py]
        if is_within_bounds(sensor_right_pos):
            px, py = world_to_pheromone_grid(sensor_right_pos)
            p_right = pheromone_grid[px, py]

        return p_main, p_left, p_right

    def steer(self, p_main, p_left, p_right):
        if p_main > p_left and p_main > p_right:
            pass # Go straight
        elif p_left > p_right:
            self.angle -= ANT_ROTATION_SPEED
        elif p_right > p_left:
            self.angle += ANT_ROTATION_SPEED
        else:
             self.angle += random.uniform(-ANT_ROTATION_SPEED * 0.5, ANT_ROTATION_SPEED * 0.5)
        self.angle = normalize_angle(self.angle)

    def move(self):
        self.pos += self.speed * np.array([math.cos(self.angle), math.sin(self.angle)])
        # Boundary bouncing
        if not (0 < self.pos[0] < WIDTH):
            self.angle = math.pi - self.angle
            self.pos[0] = np.clip(self.pos[0], 1, WIDTH - 1)
        if not (0 < self.pos[1] < HEIGHT):
            self.angle = -self.angle
            self.pos[1] = np.clip(self.pos[1], 1, HEIGHT - 1)
        self.angle = normalize_angle(self.angle)

    def update(self, pheromone_food, pheromone_home, food_sources):
        original_state = self.state # Remember state before potential change

        if self.state == SEARCHING:
            p_main, p_left, p_right = self.sense_pheromones(pheromone_food) # Sense food trail
            self.steer(p_main, p_left, p_right)
            # Check for food
            for food_pos, food_radius in food_sources:
                if np.linalg.norm(self.pos - food_pos) < food_radius:
                    self.state = RETURNING
                    self.angle = normalize_angle(self.angle + math.pi) # Turn around
                    break # Found food, stop checking others

        elif self.state == RETURNING:
             # Drop 'food' pheromones BEFORE sensing 'home' trail
             px, py = world_to_pheromone_grid(self.pos)
             pheromone_food[px, py] = min(MAX_PHEROMONE, pheromone_food[px, py] + PHEROMONE_DROP_RATE)

             p_main, p_left, p_right = self.sense_pheromones(pheromone_home) # Sense home trail
             self.steer(p_main, p_left, p_right)
             # Check for nest
             if np.linalg.norm(self.pos - self.nest_pos) < 15: # Nest radius
                 self.state = SEARCHING
                 self.angle = normalize_angle(self.angle + math.pi) # Turn around

        # Only move if the state didn't *just* change (prevents moving away immediately after finding food/nest)
        if self.state == original_state:
             self.move()
        else:
             # Update color immediately if state changed
             self.color = ANT_COLORS[self.state]

        # After moving (or state change), ants returning home deposit home pheromone at new location
        if self.state == SEARCHING and original_state == RETURNING: # Just arrived at nest
             px, py = world_to_pheromone_grid(self.pos)
             pheromone_home[px, py] = min(MAX_PHEROMONE, pheromone_home[px, py] + PHEROMONE_DROP_RATE * 2) # Stronger drop at nest?

# --- Simulation Setup ---
# Environment elements
nest_pos = (WIDTH // 2, HEIGHT // 2)
food_sources = [
    (np.array([WIDTH * 0.8, HEIGHT * 0.2]), 20),
    (np.array([WIDTH * 0.2, HEIGHT * 0.8]), 20),
]

# Initialize ants
ants = [Ant(nest_pos) for _ in range(N_ANTS)]

# Pheromone grids
pheromone_food = np.zeros((P_WIDTH, P_HEIGHT), dtype=float)
pheromone_home = np.zeros((P_WIDTH, P_HEIGHT), dtype=float)

# --- Matplotlib Visualization Setup ---
fig, ax = plt.subplots(figsize=(10, 7.5)) # Adjust figsize as needed
ax.set_xlim(0, WIDTH)
ax.set_ylim(0, HEIGHT)
ax.set_aspect('equal')
ax.set_facecolor('black') # Set background color
plt.xticks([])
plt.yticks([])
plt.title("Ant Foraging Simulation (Matplotlib)")

# Create plot elements to update
# Pheromones: Use imshow with separate colormaps and alpha blending
# Normalize pheromone values for color mapping (0 to 1)
norm_ph = plt.Normalize(0, MAX_PHEROMONE * 0.5) # Adjust max value for better visibility if needed
im_food = ax.imshow(pheromone_food.T, cmap='Greens', alpha=0.6, origin='lower',
                     extent=[0, WIDTH, 0, HEIGHT], interpolation='bilinear', norm=norm_ph)
im_home = ax.imshow(pheromone_home.T, cmap='Blues', alpha=0.6, origin='lower',
                     extent=[0, WIDTH, 0, HEIGHT], interpolation='bilinear', norm=norm_ph)


# Ants: Use scatter plot. Store the scatter object to update its offsets and colors.
ant_positions = np.array([ant.pos for ant in ants])
ant_colors_array = [ant.color for ant in ants]
scatter_ants = ax.scatter(ant_positions[:, 0], ant_positions[:, 1], s=10, c=ant_colors_array) # s is marker size

# Static elements (Nest and Food)
ax.scatter(nest_pos[0], nest_pos[1], s=150, c='saddlebrown', marker='o', label='Nest') # Use s for size
for food_pos, food_radius in food_sources:
    ax.scatter(food_pos[0], food_pos[1], s=food_radius*15, c='lime', marker='o', label='Food') # Scale radius for size

# --- Animation Update Function ---
def update(frame):
    global pheromone_food, pheromone_home, ants

    # 1. Update Pheromones (Evaporation)
    pheromone_food *= EVAPORATION_RATE
    pheromone_home *= EVAPORATION_RATE
    # Clip minimum value (optional, prevents floating point issues)
    pheromone_food[pheromone_food < 0.01] = 0
    pheromone_home[pheromone_home < 0.01] = 0

    # 2. Update Ants
    new_positions = np.zeros((N_ANTS, 2))
    new_colors = [None] * N_ANTS
    for i, ant in enumerate(ants):
        ant.update(pheromone_food, pheromone_home, food_sources)
        new_positions[i] = ant.pos
        new_colors[i] = ant.color # Update color based on state

    # 3. Update Plot Elements
    # Pheromones (set data for imshow objects) - Transpose needed as imshow expects (row, col) -> (y, x)
    im_food.set_data(pheromone_food.T)
    im_home.set_data(pheromone_home.T)

    # Ants (update scatter plot offsets and colors)
    scatter_ants.set_offsets(new_positions)
    scatter_ants.set_facecolors(new_colors) # Use set_facecolors for individual colors

    # Update title with frame number (optional)
    ax.set_title(f"Ant Foraging Simulation (Frame: {frame})")

    # Return list of modified plot elements
    return [im_food, im_home, scatter_ants]

# --- Create and Display Animation ---
# Note: interval is the delay between frames in milliseconds.
# Setting blit=True generally improves performance but can sometimes cause issues. Try False if it glitches.
# frames = number of frames to generate. Can be set high or left None for continuous (until stopped)
ani = animation.FuncAnimation(fig, update, frames=1500, interval=30, blit=True)

# Close the initial static plot display
plt.close(fig)

# Display the animation in Jupyter using HTML
# This generates Javascript/HTML to embed the animation
print("Generating animation... (this may take a moment)")
start_time = time.time()
html_output = HTML(ani.to_jshtml())
end_time = time.time()
print(f"Animation generated in {end_time - start_time:.2f} seconds.")

# ani.save('ant_simulation.mp4', writer='ffmpeg', fps=30)
html_output # Display the animation

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from IPython.display import HTML
import warnings

# Suppress Matplotlib animation warning if it occurs
warnings.filterwarnings("ignore", category=UserWarning, module="matplotlib.animation")

# --- Fitness Functions (for minimization) ---

def rastrigin_2d(x, y):
    """Rastrigin function in 2D."""
    return 20 + (x**2 - 10 * np.cos(2 * np.pi * x)) + \
           (y**2 - 10 * np.cos(2 * np.pi * y))

def schwefel_2d(x, y):
    """Schwefel function in 2D."""
    # Global minimum is at x=y=420.9687, f(x,y)=0
    # To make it a minimization problem with min near 0 for typical FA brightness:
    # We can use the standard form: 418.9829 * 2 - (x * sin(sqrt|x|) + y * sin(sqrt|y|))
    # Or a shifted/scaled one if needed for brightness calculation,
    # but for now, let's use the standard.
    term1 = x * np.sin(np.sqrt(np.abs(x)))
    term2 = y * np.sin(np.sqrt(np.abs(y)))
    return 418.9829 * 2 - (term1 + term2)

# --- Function Configuration ---
# Choose which function to use: 'rastrigin' or 'schwefel'
SELECTED_FUNCTION = 'schwefel' #@param ['rastrigin', 'schwefel']

if SELECTED_FUNCTION == 'rastrigin':
    fitness_function = rastrigin_2d
    DOMAIN_MIN = -5.12
    DOMAIN_MAX = 5.12
    GLOBAL_MIN_POS = np.array([0.0, 0.0])
    GLOBAL_MIN_VAL = 0.0
elif SELECTED_FUNCTION == 'schwefel':
    fitness_function = schwefel_2d
    DOMAIN_MIN = -500.0
    DOMAIN_MAX = 500.0
    GLOBAL_MIN_POS = np.array([420.9687, 420.9687])
    GLOBAL_MIN_VAL = 0.0 # Schwefel is 0 at its global min
else:
    raise ValueError("Invalid function selected")

# For scaling the random walk
DOMAIN_WIDTH = DOMAIN_MAX - DOMAIN_MIN

# Firefly Algorithm Parameters
N_FIREFLIES = 50       # Number of fireflies
MAX_GENERATIONS = 100  # Maximum number of generations
ALPHA = 0.01            # Randomness factor (for random walk part)
BETA0 = 0.01            # Attractiveness at r=0
GAMMA = 0.5            # Light absorption coefficient
                       # For Rastrigin with domain [-5.12, 5.12], r^2 can be up to (10.24*sqrt(2))^2 approx 200.
                       # If gamma=0.5, exp(-100) is tiny.
                       # If gamma=0.01, exp(-2) approx 0.13.
                       # Let's adjust gamma based on domain, or make it smaller.
                       # For a domain width of ~10, r_max_sq is ~200.
                       # For a domain width of ~1000 (Schwefel), r_max_sq is ~2,000,000
                       # We need gamma * r^2 to be in a reasonable range.
                       # Let's make gamma inversely proportional to characteristic squared length scale
CHARACTERISTIC_LENGTH_SQUARED = (DOMAIN_WIDTH)**2 / 10 # Heuristic
GAMMA = 1.0 / CHARACTERISTIC_LENGTH_SQUARED if CHARACTERISTIC_LENGTH_SQUARED > 0 else 0.1
ALPHA_DELTA = 0.97 # Optional: factor to reduce alpha over generations

# Initialize fireflies
# Each firefly is [x, y]
firefly_positions = np.random.uniform(DOMAIN_MIN, DOMAIN_MAX, (N_FIREFLIES, 2))
firefly_brightness = np.zeros(N_FIREFLIES)
firefly_fitness = np.zeros(N_FIREFLIES)

# Best solution found so far
global_best_pos = None
global_best_fitness = float('inf')

def calculate_brightness(fitness_values):
    # Assuming minimization, fitness >= 0. Higher brightness = better (lower fitness)
    # Add a small epsilon to avoid division by zero if fitness is 0
    return 1.0 / (1.0 + fitness_values + 1e-6)

# Initial fitness and brightness calculation
for i in range(N_FIREFLIES):
    x, y = firefly_positions[i]
    f_val = fitness_function(x, y)
    firefly_fitness[i] = f_val
    if f_val < global_best_fitness:
        global_best_fitness = f_val
        global_best_pos = firefly_positions[i].copy()

firefly_brightness = calculate_brightness(firefly_fitness)

print(f"Using function: {SELECTED_FUNCTION}")
print(f"Domain: [{DOMAIN_MIN}, {DOMAIN_MAX}]")
print(f"Gamma: {GAMMA:.4e}")
print(f"Initial best fitness: {global_best_fitness:.4f} at {global_best_pos}")

def update_fireflies_and_get_best(positions, brightness, current_fitness,
                                  current_global_best_pos, current_global_best_fitness):
    n = positions.shape[0]
    num_dims = positions.shape[1]
    new_positions = np.copy(positions) # Work on a copy for this generation's new positions

    # Store current brightness for comparisons within this generation
    # This is I(t) based on positions(t)
    current_brightness_for_comparison = np.copy(brightness)

    for i in range(n):
        # This firefly's position will be updated based on attractions
        # It starts from its position at time t
        x_i_candidate = np.copy(positions[i])
        
        # Flag to check if firefly 'i' was attracted and moved by any 'j'
        # This helps decide if it needs a standalone random walk (if it's a global best type)
        attracted_and_moved = False

        for j in range(n):
            if i == j:
                continue

            # Compare based on brightness at the start of the generation
            if current_brightness_for_comparison[j] > current_brightness_for_comparison[i]:
                attracted_and_moved = True
                
                # Calculate distance: r_ij
                # Distance is between x_i_candidate (which evolves within this i-loop)
                # and positions[j] (which is fixed at t for this generation)
                r_sq = np.sum((x_i_candidate - positions[j])**2)
                
                beta = BETA0 * np.exp(-GAMMA * r_sq)
                
                # Random vector for this specific interaction
                # Scale factor ensures random step is proportional to domain width
                random_vec = ALPHA * (np.random.rand(num_dims) - 0.5) * DOMAIN_WIDTH
                
                # Update x_i_candidate based on attraction to this j AND its own random jitter for THIS interaction
                x_i_candidate += beta * (positions[j] - x_i_candidate) + random_vec
                
                # Apply boundary conditions immediately after each micro-move
                x_i_candidate = np.clip(x_i_candidate, DOMAIN_MIN, DOMAIN_MAX)
        
        # If firefly 'i' was not attracted to any brighter firefly (e.g., it's the current global best or in a "dark" area alone)
        # it performs a simple random walk.
        if not attracted_and_moved:
            random_vec_standalone = ALPHA * (np.random.rand(num_dims) - 0.5) * DOMAIN_WIDTH
            x_i_candidate += random_vec_standalone
            x_i_candidate = np.clip(x_i_candidate, DOMAIN_MIN, DOMAIN_MAX)
            
        new_positions[i] = x_i_candidate

    # Now, evaluate fitness for all new positions and update brightness
    new_fitness = np.zeros(n)
    updated_global_best_pos = current_global_best_pos.copy() if current_global_best_pos is not None else None
    updated_global_best_fitness = current_global_best_fitness

    for i in range(n):
        f_val = fitness_function(new_positions[i, 0], new_positions[i, 1])
        new_fitness[i] = f_val
        if f_val < updated_global_best_fitness:
            updated_global_best_fitness = f_val
            updated_global_best_pos = new_positions[i].copy()
            
    new_brightness = calculate_brightness(new_fitness)
    
    return new_positions, new_brightness, new_fitness, updated_global_best_pos, updated_global_best_fitness

fig, ax = plt.subplots(figsize=(8, 6))
ax.set_xlim(DOMAIN_MIN, DOMAIN_MAX)
ax.set_ylim(DOMAIN_MIN, DOMAIN_MAX)

# Plot contour of the fitness function
x_grid = np.linspace(DOMAIN_MIN, DOMAIN_MAX, 100)
y_grid = np.linspace(DOMAIN_MIN, DOMAIN_MAX, 100)
X_mesh, Y_mesh = np.meshgrid(x_grid, y_grid)
Z_mesh = fitness_function(X_mesh, Y_mesh)
contour = ax.contourf(X_mesh, Y_mesh, Z_mesh, levels=50, cmap='viridis') # Or 'plasma', 'magma'
fig.colorbar(contour, ax=ax, label='Fitness Value')

# Plot global minimum if known
if GLOBAL_MIN_POS is not None:
    ax.plot(GLOBAL_MIN_POS[0], GLOBAL_MIN_POS[1], 'X', color='red', markersize=10, label=f'Global Min ({GLOBAL_MIN_VAL:.2f})')
ax.legend(loc='upper right')

# Initial scatter plot for fireflies
# We can use brightness to modulate color or size if desired, but simple points for now
# firefly_colors = firefly_brightness / np.max(firefly_brightness + 1e-6) # Normalize for color
scatter = ax.scatter(firefly_positions[:, 0], firefly_positions[:, 1], 
                     color='orange', edgecolor='black', s=50, zorder=3, label='Fireflies') # s=size

title = ax.set_title(f"Generation: 0, Best Fitness: {global_best_fitness:.4f}")

# Animation function
def animate(generation):
    global firefly_positions, firefly_brightness, firefly_fitness, global_best_pos, global_best_fitness
    
    firefly_positions, firefly_brightness, firefly_fitness, \
    global_best_pos, global_best_fitness = update_fireflies_and_get_best(
                                                firefly_positions, 
                                                firefly_brightness, # Pass current brightness for comparison
                                                firefly_fitness,    # Pass current fitness (for info, not direct use in update logic)
                                                global_best_pos, 
                                                global_best_fitness
                                            )
    
    scatter.set_offsets(firefly_positions)
    # Optionally update firefly colors/sizes based on new brightness
    # norm_brightness = firefly_brightness / (np.max(firefly_brightness) + 1e-6)
    # scatter.set_array(norm_brightness) # If using color to represent brightness
    
    title.set_text(f"Generation: {generation + 1}, Best Fitness: {global_best_fitness:.4f}")
    
    if (generation + 1) % 10 == 0 or generation == MAX_GENERATIONS -1 : # Print progress
        print(f"Gen: {generation + 1}/{MAX_GENERATIONS}, Best Fit: {global_best_fitness:.4f}, Pos: {global_best_pos}")
        
    return scatter, title,

# Create and display animation
# blit=False because title is also an artist we are returning and it might cause issues with blit=True sometimes
ani = FuncAnimation(fig, animate, frames=MAX_GENERATIONS, interval=150, blit=False, repeat=False)

plt.close(fig) # Prevent static plot from showing up before the animation
HTML(ani.to_jshtml())
# To save as mp4 (ensure ffmpeg is installed):
# ani.save('firefly_algorithm.mp4', writer='ffmpeg', fps=15)
# print("Animation saved as firefly_algorithm.mp4")

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from IPython.display import HTML
import warnings

# Suppress Matplotlib animation warning if it occurs
warnings.filterwarnings("ignore", category=UserWarning, module="matplotlib.animation")

# --- Fitness Functions (for minimization) ---

def rastrigin_2d(x, y):
    """Rastrigin function in 2D."""
    return 20 + (x**2 - 10 * np.cos(2 * np.pi * x)) + \
           (y**2 - 10 * np.cos(2 * np.pi * y))

def schwefel_2d(x, y):
    """Schwefel function in 2D."""
    term1 = x * np.sin(np.sqrt(np.abs(x)))
    term2 = y * np.sin(np.sqrt(np.abs(y)))
    return 418.9829 * 2 - (term1 + term2)

# --- Function Configuration ---
# Choose which function to use: 'rastrigin' or 'schwefel'
SELECTED_FUNCTION = 'rastrigin' #@param ['rastrigin', 'schwefel']

if SELECTED_FUNCTION == 'rastrigin':
    fitness_function = rastrigin_2d
    DOMAIN_MIN = -5.12
    DOMAIN_MAX = 5.12
    GLOBAL_MIN_POS_THEORETICAL = np.array([0.0, 0.0])
    GLOBAL_MIN_VAL_THEORETICAL = 0.0
elif SELECTED_FUNCTION == 'schwefel':
    fitness_function = schwefel_2d
    DOMAIN_MIN = -500.0
    DOMAIN_MAX = 500.0
    GLOBAL_MIN_POS_THEORETICAL = np.array([420.9687, 420.9687])
    GLOBAL_MIN_VAL_THEORETICAL = 0.0
else:
    raise ValueError("Invalid function selected")

DOMAIN_WIDTH = DOMAIN_MAX - DOMAIN_MIN
NUM_DIMS = 2 # Working in 2D

# PSO Parameters
N_PARTICLES = 30
MAX_ITERATIONS = 100

# Inertia weight: linearly decreasing from W_MAX to W_MIN
W_MAX = 0.5
W_MIN = 0.1

C1 = 0.5  # Cognitive coefficient (pbest influence)
C2 = 0.5  # Social coefficient (gbest influence)

# Velocity clamping (max velocity as a fraction of domain width per dimension)
V_MAX_FACTOR = 0.01 # Max velocity is 5% of domain width
V_MAX = V_MAX_FACTOR * DOMAIN_WIDTH
V_MIN = -V_MAX

# Initial velocity range (fraction of domain width)
V_INIT_RANGE_FACTOR = 0.05
V_INIT_MAX = V_INIT_RANGE_FACTOR * DOMAIN_WIDTH
V_INIT_MIN = -V_INIT_MAX

# Initialize particles
particle_positions = np.random.uniform(DOMAIN_MIN, DOMAIN_MAX, (N_PARTICLES, NUM_DIMS))
particle_velocities = np.random.uniform(V_INIT_MIN, V_INIT_MAX, (N_PARTICLES, NUM_DIMS))

# Personal best positions and fitness
particle_pbest_positions = np.copy(particle_positions)
particle_pbest_fitness = np.full(N_PARTICLES, float('inf'))

# Global best position and fitness
gbest_position = None
gbest_fitness = float('inf')

# Initial fitness calculation and pbest/gbest update
for i in range(N_PARTICLES):
    current_fitness = fitness_function(particle_positions[i, 0], particle_positions[i, 1])
    particle_pbest_fitness[i] = current_fitness
    # pbest_positions are already set to initial positions

    if current_fitness < gbest_fitness:
        gbest_fitness = current_fitness
        gbest_position = particle_positions[i].copy()

print(f"Using function: {SELECTED_FUNCTION}")
print(f"Domain: [{DOMAIN_MIN}, {DOMAIN_MAX}]")
print(f"Initial global best fitness: {gbest_fitness:.4f} at {gbest_position}")

def update_pso_and_get_best(iteration, positions, velocities,
                             pbest_pos, pbest_fit,
                             gbest_pos, gbest_fit):
    
    n_particles = positions.shape[0]
    num_dims = positions.shape[1]

    # Linearly decreasing inertia weight
    w = W_MAX - (W_MAX - W_MIN) * (iteration / MAX_ITERATIONS)

    # Update pbest and gbest based on current positions
    current_gbest_pos_iter = gbest_pos.copy() if gbest_pos is not None else None
    current_gbest_fit_iter = gbest_fit
    
    new_pbest_pos = np.copy(pbest_pos)
    new_pbest_fit = np.copy(pbest_fit)

    for i in range(n_particles):
        current_fitness = fitness_function(positions[i, 0], positions[i, 1])
        
        if current_fitness < new_pbest_fit[i]:
            new_pbest_fit[i] = current_fitness
            new_pbest_pos[i] = positions[i].copy()
            
            if current_fitness < current_gbest_fit_iter:
                current_gbest_fit_iter = current_fitness
                current_gbest_pos_iter = positions[i].copy()
                
    # Update velocities and positions
    new_velocities = np.zeros_like(velocities)
    new_positions = np.zeros_like(positions)

    for i in range(n_particles):
        r1 = np.random.rand(num_dims) # Random numbers for cognitive component
        r2 = np.random.rand(num_dims) # Random numbers for social component
        
        cognitive_velocity = C1 * r1 * (new_pbest_pos[i] - positions[i])
        social_velocity = C2 * r2 * (current_gbest_pos_iter - positions[i])
        
        new_velocities[i] = w * velocities[i] + cognitive_velocity + social_velocity
        
        # Apply velocity clamping
        new_velocities[i] = np.clip(new_velocities[i], V_MIN, V_MAX)
        
        # Update position
        new_positions[i] = positions[i] + new_velocities[i]
        
        # Apply boundary conditions (clamping)
        new_positions[i] = np.clip(new_positions[i], DOMAIN_MIN, DOMAIN_MAX)
        
    return new_positions, new_velocities, new_pbest_pos, new_pbest_fit, current_gbest_pos_iter, current_gbest_fit_iter


# --- Ensure a clean slate for Matplotlib ---
plt.close('all') 

# --- Setup the new figure and axes ---
fig, ax = plt.subplots(figsize=(8, 6))
ax.set_xlim(DOMAIN_MIN, DOMAIN_MAX)
ax.set_ylim(DOMAIN_MIN, DOMAIN_MAX)

# Plot contour of the fitness function
x_grid = np.linspace(DOMAIN_MIN, DOMAIN_MAX, 100)
y_grid = np.linspace(DOMAIN_MIN, DOMAIN_MAX, 100)
X_mesh, Y_mesh = np.meshgrid(x_grid, y_grid)
Z_mesh = fitness_function(X_mesh, Y_mesh)

contour = ax.contourf(X_mesh, Y_mesh, Z_mesh, levels=50, cmap='viridis')
fig.colorbar(contour, ax=ax, label='Fitness Value')

# Plot theoretical global minimum if known
if GLOBAL_MIN_POS_THEORETICAL is not None:
    ax.plot(GLOBAL_MIN_POS_THEORETICAL[0], GLOBAL_MIN_POS_THEORETICAL[1], 'X', color='red', 
            markersize=10, label=f'Global Min ({GLOBAL_MIN_VAL_THEORETICAL:.2f})')

# Initial scatter plot for particles
scatter = ax.scatter(particle_positions[:, 0], particle_positions[:, 1], 
                     color='orange', edgecolor='black', s=30, zorder=3, label='Particles')

ax.legend(loc='upper right')
title_artist = ax.set_title(f"Iteration: 0, Best Fitness: {gbest_fitness:.4f}")

# Animation function
def animate(iteration):
    global particle_positions, particle_velocities, \
           particle_pbest_positions, particle_pbest_fitness, \
           gbest_position, gbest_fitness
    
    # Update particle states using PSO logic
    particle_positions, particle_velocities, \
    particle_pbest_positions, particle_pbest_fitness, \
    gbest_position, gbest_fitness = update_pso_and_get_best(
                                        iteration,
                                        particle_positions, 
                                        particle_velocities,
                                        particle_pbest_positions,
                                        particle_pbest_fitness,
                                        gbest_position, 
                                        gbest_fitness
                                    )
    
    scatter.set_offsets(particle_positions) # Update particle positions on the plot
    title_artist.set_text(f"Iteration: {iteration + 1}, Best Fitness: {gbest_fitness:.4f}")
    
    if (iteration + 1) % 10 == 0 or iteration == MAX_ITERATIONS - 1:
        print(f"Iter: {iteration + 1}/{MAX_ITERATIONS}, Best Fit: {gbest_fitness:.4f}, Pos: {gbest_position}")
        
    return scatter, title_artist,

# Ensure initial state for title before animation starts
if gbest_position is not None:
    title_artist.set_text(f"Iteration: 0, Best Fitness: {gbest_fitness:.4f} at ({gbest_position[0]:.2f}, {gbest_position[1]:.2f})")
else:
    title_artist.set_text(f"Iteration: 0, Best Fitness: {gbest_fitness:.4f}")

# Create and display animation
ani = FuncAnimation(fig, animate, frames=MAX_ITERATIONS, interval=150, blit=False, repeat=False)

animation_html = HTML(ani.to_jshtml())
plt.close(fig) 
animation_html
# To save as mp4 (ensure ffmpeg is installed):
# ani.save('pso_algorithm.mp4', writer='ffmpeg', fps=15)
# print("Animation saved as pso_algorithm.mp4")