# -*- coding: utf-8 -*-
"""
Created on Wed Nov 26 20:25:53 2025

@author: Panos
"""

import random
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation


# ============================================================
# 1. Objective (fitness) function
#    f(x, y) = (x - 2y + 3)^2 + (2x + y - 8)^2
# ============================================================
def fitness_function(x1, x2):
    f1 = x1 - 2 * x2 + 3
    f2 = 2 * x1 + x2 - 8
    z = f1 ** 2 + f2 ** 2
    return z


# ============================================================
# 2. Velocity and position update (standard PSO formulas)
# ============================================================
def update_velocity(particle, velocity, pbest, gbest,
                    w=0.7, c1=1.4, c2=1.4):
    """
    particle, velocity, pbest, gbest are 1D numpy arrays of size = dimension
    w  : inertia weight
    c1 : cognitive coefficient
    c2 : social coefficient
    """
    r1 = random.random()
    r2 = random.random()
    new_velocity = (
        w * velocity
        + c1 * r1 * (pbest - particle)
        + c2 * r2 * (gbest - particle)
    )
    return new_velocity


def update_position(particle, velocity):
    return particle + velocity


# ============================================================
# 3. PSO main algorithm (2D for plotting, but generic)
# ============================================================
def pso_2d(population=30,
           dimension=2,
           position_min=-10.0,
           position_max=10.0,
           max_generations=100,
           fitness_threshold=1e-3):
    """
    Returns:
        gbest_position : best position found (2D)
        gbest_fitness  : best fitness value
        swarm_history  : list of particle positions for each generation
                         (for plotting / animation)
    """

    # ---------------------
    # Step 1: Initialize swarm
    # ---------------------
    particles = np.array([
        [random.uniform(position_min, position_max)
         for _ in range(dimension)]
        for _ in range(population)
    ])

    velocities = np.zeros_like(particles)

    # Personal bests start at initial positions
    pbest_position = particles.copy()
    pbest_fitness = np.array([
        fitness_function(p[0], p[1]) for p in particles
    ])

    # Global best from personal bests
    gbest_index = np.argmin(pbest_fitness)
    gbest_position = pbest_position[gbest_index].copy()
    gbest_fitness = pbest_fitness[gbest_index]

    # For animation: store positions of all particles each generation
    swarm_history = [particles.copy()]

    # ---------------------
    # Steps 2–6: Main loop
    # ---------------------
    for gen in range(max_generations):

        for i in range(population):
            # Step 3.1: update velocity
            velocities[i] = update_velocity(
                particles[i],
                velocities[i],
                pbest_position[i],
                gbest_position
            )

            # Step 3.2: update position
            particles[i] = update_position(particles[i], velocities[i])

            # (optional) clamp positions inside search space
            particles[i] = np.clip(particles[i], position_min, position_max)

            # Step 3.3: evaluate fitness at new position
            current_fitness = fitness_function(particles[i][0],
                                               particles[i][1])

            # Step 3.4: update personal best
            if current_fitness < pbest_fitness[i]:
                pbest_fitness[i] = current_fitness
                pbest_position[i] = particles[i].copy()

                # Step 3.5: update global best (using personal bests)
                if current_fitness < gbest_fitness:
                    gbest_fitness = current_fitness
                    gbest_position = particles[i].copy()

        swarm_history.append(particles.copy())

        # Step 3.6: stopping criterion
        if gbest_fitness <= fitness_threshold:
            print(f"Stopped early at generation {gen+1}")
            break

    return gbest_position, gbest_fitness, swarm_history


# ============================================================
# 4. Plotting and animation
# ============================================================
def make_animation(swarm_history,
                   position_min=-10,
                   position_max=10,
                   filename="pso_animation.gif"):
    """
    swarm_history: list of (N, 2) arrays with particle positions
    """
    fig = plt.figure(figsize=(8, 8))
    ax = fig.add_subplot(111, projection='3d')

    # Plot the surface once
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('f(x, y)')

    x = np.linspace(position_min, position_max, 80)
    y = np.linspace(position_min, position_max, 80)
    X, Y = np.meshgrid(x, y)
    Z = fitness_function(X, Y)
    ax.plot_wireframe(X, Y, Z, linewidth=0.3)

    images = []

    # One frame per generation
    for positions in swarm_history:
        xs = positions[:, 0]
        ys = positions[:, 1]
        zs = np.array([fitness_function(xi, yi) for xi, yi in zip(xs, ys)])

        scat = ax.scatter(xs, ys, zs)
        images.append([scat])

    ani = animation.ArtistAnimation(
        fig,
        images,
        interval=150,
        blit=True,
        repeat_delay=1000
    )
    ani.save(filename, writer='pillow')
    plt.close(fig)


# ============================================================
# 5. Run PSO and create animation
# ============================================================
if __name__ == "__main__":
    gbest_pos, gbest_fit, swarm_hist = pso_2d(
        population=50,
        dimension=2,
        position_min=-10.0,
        position_max=10.0,
        max_generations=100,
        fitness_threshold=1e-3
    )

    print("Global best position:", gbest_pos)
    print("Global best fitness :", gbest_fit)

    make_animation(
        swarm_history=swarm_hist,
        position_min=-10,
        position_max=10,
        filename="pso_clean.gif"
    )