import numpy as np
uniRand = np.random.uniform
class Particle:
    def __init__(self, N, xRange, vRange):
        # 生成(N)維參數(shù)
        self.x = uniRand(*xRange, (N,))
        self.v = uniRand(*vRange, (N,))
        self.best = np.inf
        self.xBest = np.zeros((N,))

from copy import deepcopy
rand = np.random.rand
class Swarm:
    # pNum 粒子個(gè)數(shù)，N粒子維度，即參數(shù)個(gè)數(shù)
    # wRange 最大最小權(quán)重
    def __init__(self, pNum, N, xRange, vRange, wRange, c):
        self.ps = [Particle(N, xRange, vRange) for _ in range(pNum)]
        self.bestPs = deepcopy(self.ps)
        self.best = np.inf  #全局最優(yōu)值
        self.xBest = np.zeros((N,)) #全局最優(yōu)參數(shù)
        self.wRange = wRange
        self.c0, self.c1 = c
        self.N = N

# 下面的函數(shù)寫在class Swarm中
# 為了書寫方便，這里頂個(gè)寫，注意在Swarm中的縮進(jìn)
    # 參數(shù)更新 func為待優(yōu)化函數(shù)
def oneStep(self, func, w=1):
    for p in self.ps:
        y = func(p.x)
        # 更新歷史最優(yōu)值
        if y < p.best:
            p.best = y
            p.xBest = deepcopy(p.x)
        # 更新粒子群最優(yōu)情況
        if y < self.best:
            self.best = y
            self.xBest = deepcopy(p.x)
            self.bestPs = deepcopy(self.ps)
    for p in self.ps:
        iw = uniRand(*self.wRange, 1)[0]
        us = self.c0 * rand(self.N) * (p.xBest - p.x)
        vs = self.c1 * rand(self.N) * (self.xBest - p.x)
        p.v = iw * p.v + us + vs
        p.x = p.x + p.v
    return self.best
# nMsg為輸出提示的周期
# iter為迭代次數(shù)
def optimize(self, func, nMsg, iter):
    for i in range(iter):
        best = self.oneStep()
        if i % nMsg == 0:
            print(f"第{i}次迭代最小值為{best}")

def test(xs):
    s = 0.0
    for i in range(len(xs)):
        s += np.cos(i*xs[i]/5)*(i+1)
    return s

if __name__ == "__main__":
    xRange = (-3,3)
    vRange = (-1,1)
    wRange = (0.5,1)
    C = (1.5, 1.5)
    s = Swarm(20, 10, xRange, vRange, wRange, C)
    s.optimize(test, 20, 200)
    print("最佳位置在：\n", s.xBest)