import numpy as np
import matplotlib.pyplot as plt

class Particle:
    def __init__(self, dim, bounds):
        """입자 초기화"""
        self.position = np.random.uniform(bounds[0], bounds[1], dim)
        self.velocity = np.random.uniform(-1, 1, dim)
        self.best_position = self.position.copy()
        self.best_fitness = float('inf')
    
    def update_velocity(self, global_best, w=0.7, c1=1.5, c2=1.5):
        """속도 업데이트"""
        r1, r2 = np.random.rand(2)
        cognitive = c1 * r1 * (self.best_position - self.position)
        social = c2 * r2 * (global_best - self.position)
        self.velocity = w * self.velocity + cognitive + social
    
    def update_position(self, bounds):
        """위치 업데이트"""
        self.position += self.velocity
        # 경계 처리
        self.position = np.clip(self.position, bounds[0], bounds[1])

def sphere_function(x):
    """테스트 함수: 구 함수 (최솟값은 0)"""
    return np.sum(x**2)

def pso_optimize(func, dim=2, n_particles=30, n_iterations=100, bounds=(-10, 10)):
    """PSO 최적화 실행"""
    # 입자 떼 초기화
    swarm = [Particle(dim, bounds) for _ in range(n_particles)]
    global_best_position = np.random.uniform(bounds[0], bounds[1], dim)
    global_best_fitness = float('inf')
    
    history = []
    
    for iteration in range(n_iterations):
        for particle in swarm:
            # 적합도 평가
            fitness = func(particle.position)
            
            # 개인 최고 업데이트
            if fitness < particle.best_fitness:
                particle.best_fitness = fitness
                particle.best_position = particle.position.copy()
            
            # 전역 최고 업데이트
            if fitness < global_best_fitness:
                global_best_fitness = fitness
                global_best_position = particle.position.copy()
        
        # 모든 입자 업데이트
        for particle in swarm:
            particle.update_velocity(global_best_position)
            particle.update_position(bounds)
        
        history.append(global_best_fitness)
        
        if iteration % 20 == 0:
            print(f"반복 {iteration}: 최고 적합도 = {global_best_fitness:.6f}")
    
    return global_best_position, global_best_fitness, history

# 실행
best_pos, best_fit, history = pso_optimize(sphere_function)
print(f"\n최종 결과:")
print(f"최적 위치: {best_pos}")
print(f"최적 값: {best_fit:.6f}")

# 수렴 그래프
plt.plot(history)
plt.xlabel('반복 횟수')
plt.ylabel('최고 적합도')
plt.title('PSO 수렴 과정')
plt.yscale('log')
plt.grid(True)
plt.show()