import random
import math
import matplotlib.pyplot as plt

def fitness(x):
    return sum(i**2 for i in x)

n = 10  #no. of particles
dim = 2 #no. of variables
low, high = -7.5, 7.5 #search boundary
iterations = 50

w, c1, c2 = 0.5, 1.5, 1.5
minimize = True

pos = [[random.uniform(low,high) for _ in range(dim)] for _ in range(n)]
vel = [[0]*dim for _ in range(n)]
pbest = [p.copy() for p in pos]

def better(a,b):
    return a < b if minimize else a > b

gbest = min(pbest,key=fitness) if minimize else max(pbest,key=fitness)
history = []

for _ in range(iterations):
    for i in range(n):
        for d in range(dim):
            r1,r2 = random.random(),random.random()

            vel[i][d] = (w*vel[i][d]
                         + c1*r1*(pbest[i][d]-pos[i][d])
                         + c2*r2*(gbest[d]-pos[i][d]))

            pos[i][d] += vel[i][d]
            pos[i][d] = max(low,min(high,pos[i][d]))

        if better(fitness(pos[i]),fitness(pbest[i])):
            pbest[i] = pos[i].copy()

    current = min(pbest,key=fitness) if minimize else max(pbest,key=fitness)

    if better(fitness(current),fitness(gbest)):
        gbest = current.copy()

    history.append(fitness(gbest))

print("Best Position:",gbest)
print("Best Fitness:",fitness(gbest))


distance = math.sqrt(sum(x**2 for x in gbest))
print("Distance from Origin:",distance)

plt.plot(history)
plt.xlabel("Iterations")
plt.ylabel("Best Fitness")
plt.title("PSO Convergence")
plt.grid()
plt.show()
