import random
import numpy as np

# -----------------------------
# Objective function (K-D)
# -----------------------------
def f(x):
    return np.sum(x**2)

# -----------------------------
# PSO parameters
# -----------------------------
w = 0.5      # inertia
c1 = 1.0     # cognitive
c2 = 1.0     # social

n_particles = 10
iterations = 30
k = 3  # dimensionality (change this!)

# -----------------------------
# Initialize swarm
# -----------------------------
particles = [np.random.uniform(-10, 10, k) for _ in range(n_particles)]
velocities = [np.zeros(k) for _ in range(n_particles)]

# Personal bests
pbest = particles.copy()
pbest_val = [f(x) for x in particles]

# Global best
gbest_index = np.argmin(pbest_val)
gbest = pbest[gbest_index].copy()

print("Initial gbest:", gbest, "f =", f(gbest))

# -----------------------------
# PSO main loop
# -----------------------------
for it in range(iterations):

    for i in range(n_particles):

        r1 = np.random.random(k)
        r2 = np.random.random(k)

        # Velocity update
        velocities[i] = (
            w * velocities[i]
            + c1 * r1 * (pbest[i] - particles[i])
            + c2 * r2 * (gbest - particles[i])
        )

        # Position update
        particles[i] = particles[i] + velocities[i]

        # Evaluate
        fitness = f(particles[i])

        # Update personal best
        if fitness < pbest_val[i]:
            pbest[i] = particles[i].copy()
            pbest_val[i] = fitness

    # Update global best
    gbest_index = np.argmin(pbest_val)
    gbest = pbest[gbest_index].copy()

    print(f"Iteration {it+1:2d} | gbest = {gbest} | f = {f(gbest):.6f}")

# -----------------------------
# Final result
# -----------------------------
print("\nFinal solution:")
print("gbest =", gbest)
print("f(gbest) =", f(gbest))