import numpy as np
import matplotlib.pyplot as plt
from IPython import embed
 
 
target = np.random.uniform(low=0.0, high=64, size=(1, 2))
 
 
def eval_func(particle):
    dist = ((target - particle) ** 2).sum() ** 0.5
    return (64 - dist) / 64
 
 
class PSO:
    def __init__(self, n_particles, n_dims, v_max, v_min, func):
        super().__init__()
        # 입자들의 위치와 속도 초기화
        self.particles = np.random.uniform(high=v_max, low=v_min, size=(n_particles, n_dims))
        diff = v_max - v_min
        self.velocity = np.random.uniform(high=diff, low=-diff, size=(n_particles, n_dims))
 
        # 개별 입자가 찾은 최적해 및 적합도
        self.particle_best = np.random.uniform(high=v_max, low=v_min, size=(n_particles, n_dims))
        self.particle_best_fitness = -1e10 * np.ones(n_particles)
 
        # 군집이 찾은 최고 적합도
        self.swarm_best_fitness = -1e10
 
        # 초기 평가
        self.func = func
        self.eval()
 
    def eval(self):
        # 평가
        fitness = [f for f in map(self.func, self.particles)]
        self.fitness = np.array(fitness, dtype=np.float32)
 
        # 개별 입자가 찾은 최적해 및 적합도 갱신
        cond = self.fitness > self.particle_best_fitness
        self.particle_best = np.where(cond.reshape(-1, 1), self.particles, self.particle_best)
        self.particle_best_fitness = np.where(cond, self.fitness, self.particle_best_fitness)
 
        # 군집이 찾은 최적해 및 적합도 갱신
        swarm_best_idx = self.fitness.argmax(axis=0)
        if self.fitness[swarm_best_idx] > self.swarm_best_fitness:
            self.swarm_best_fitness = self.fitness[swarm_best_idx]
            self.swarm_best = self.particles[swarm_best_idx]
        return fitness
 
    def step(self, omega=0.33, theta_p=0.33, theta_g=0.33):
        r_p, r_g = np.random.random(), np.random.random()
        # 현재 속력
        intertia = omega * self.velocity 
        # 자신이 찾은 최적해에 끌리는 힘
        to_particle_best = theta_p * r_p * (self.particle_best - self.particles)
        # 군집이 찾은 최적해에 끌리는 힘
        to_swarm_best = theta_g * r_g * (self.swarm_best - self.particles)
        # 속도 갱신
        self.velocity = intertia + to_particle_best + to_swarm_best
        # 위치 갱신
        self.particles = self.particles + self.velocity
        # 새로운 위치 평가
        self.eval()
        return self.swarm_best, self.swarm_best_fitness
 
    def plot(self):
        m = np.zeros((64, 64))
        y = self.particles[:, 0].astype(np.int32).clip(max=63, min=0)
        x = self.particles[:, 1].astype(np.int32).clip(max=63, min=0)
        m[y, x] = self.fitness
        plt.plot([target[0, 0]], [target[0, 1]], 'o')
        plt.plot([self.swarm_best[0]], [self.swarm_best[1]], 'o')
        plt.imshow(m.T, cmap='gray', vmax=1.0, vmin=0.0)
        plt.show()
        plt.clf()
 
 
if __name__ == '__main__':
 
    pso = PSO(n_particles=20, n_dims=2, v_max=64, v_min=0, func=eval_func)
    for i in range(20):
        best, fitness = pso.step()
        print(i, target, best, fitness)
    pso.plot()
 
    # embed(); exit()