import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
import matplotlib as mpl
import copy as copy
 
 

def initialization(pop,ub,lb,dim):
    ''' 种群初始化函数'''
    X = np.zeros([pop,dim]) #声明空间
    for i in range(pop):
        for j in range(dim):
            X[i,j]=(ub[j]-lb[j])*np.random.random()+lb[j] #生成[lb,ub]之间的随机数
    
    return X
     
def BorderCheck(X,ub,lb,pop,dim):
    '''边界检查函数'''
    for i in range(pop):
        for j in range(dim):
            if X[i,j]>ub[j]:
                X[i,j] = ub[j]
            elif X[i,j]<lb[j]:
                X[i,j] = lb[j]
    return X


def CaculateFitness(X,fun):
    '''计算种群的所有个体的适应度值'''
    pop = X.shape[0]
    fitness = np.zeros([pop, 1])
    for i in range(pop):
        fitness[i] = fun(X[i, :])
    return fitness


def SortFitness(Fit):
    '''适应度排序'''
    fitness = np.sort(Fit, axis=0)
    index = np.argsort(Fit, axis=0)
    return fitness,index

def SortPosition(X,index):
    '''根据适应度对位置进行排序'''
    Xnew = np.zeros(X.shape)
    for i in range(X.shape[0]):
        Xnew[i,:] = X[index[i],:]
    return Xnew


def RouletteWheelSelection(P):
    '''轮盘赌策略'''
    C = np.cumsum(P)#累加
    r = np.random.random()*C[-1]#定义选择阈值，将随机概率与总和的乘积作为阈值
    out = 0
    #若大于或等于阈值，则输出当前索引，并将其作为结果，循环结束
    for i in range(P.shape[0]):
        if r<C[i]:
            out = i
            break
    return out
        

def ABC(pop, dim, lb, ub, MaxIter, fun):
   
    L = round(0.6*dim*pop) #limit 参数
    C = np.zeros([pop,1]) #计数器
    nOnlooker=pop #引领蜂数量
    
    X= initialization(pop,ub,lb,dim)  # 初始化种群
    fitness = CaculateFitness(X, fun)  # 计算适应度值
    fitness, sortIndex = SortFitness(fitness)  # 对适应度值排序
    X = SortPosition(X, sortIndex)  # 种群排序
    GbestScore = copy.copy(fitness[0])#记录最优适应度值
    GbestPositon = np.zeros([1,dim])
    GbestPositon[0,:] = copy.copy(X[0, :])#记录最优位置
    Curve = np.zeros([MaxIter, 1])
    Xnew = np.zeros([pop,dim])
    fitnessNew = copy.copy(fitness)
    for t in range(MaxIter):
        '''引领蜂搜索'''
        for i in range(pop):
            k = np.random.randint(pop)#随机选择一个个体
            while(k==i):#当k=i时，再次随机选择，直到k不等于i
                k = np.random.randint(pop)
            phi = (2*np.random.random([1,dim]) - 1)
            Xnew[i,:] = X[i,:]+phi*(X[i,:]-X[k,:]) #公式(2.2)位置更新
        Xnew=BorderCheck(Xnew,ub,lb,pop,dim) #边界检查
        fitnessNew = CaculateFitness(Xnew, fun)  # 计算适应度值
        for i in range(pop):
            if fitnessNew[i]<fitness[i]:#如果适应度值更优，替换原始位置
                X[i,:]= copy.copy(Xnew[i,:])
                fitness[i] = copy.copy(fitnessNew[i])
            else:
                C[i] = C[i]+1 #如果位置没有更新，累加器+1
                
        #计算选择适应度权重,这里用的是e的负指数次幂，我的理解是将适应度缩小，减少计算量。
        F = np.zeros([pop,1])
        MeanCost = np.mean(fitness) #求列表均值
        for i in range(pop):
            F[i]=np.exp(-fitness[i]/MeanCost)
        P=F/sum(F) #式（2.4）
        '''侦察蜂搜索'''
        for m in range(nOnlooker):
            i=RouletteWheelSelection(P)#轮盘赌测量选择个体
            k = np.random.randint(pop)#随机选择个体
            while(k==i):
                k = np.random.randint(pop)
            phi = (2*np.random.random([1,dim]) - 1)
            Xnew[i,:] = X[i,:]+phi*(X[i,:]-X[k,:])#位置更新
        Xnew=BorderCheck(Xnew,ub,lb,pop,dim)#边界检查
        fitnessNew = CaculateFitness(Xnew, fun)  # 计算适应度值
        for i in range(pop):
            if fitnessNew[i]<fitness[i]:#如果适应度值更优，替换原始位置
                X[i,:]= copy.copy(Xnew[i,:])
                fitness[i] = copy.copy(fitnessNew[i])
            else:
                C[i] = C[i]+1 #如果位置没有更新，累加器+1
        '''判断limit条件，并进行更新'''
        for i in range(pop):
            if C[i]>=L:
                for j in range(dim):
                    X[i, j] = np.random.random() * (ub[j] - lb[j]) + lb[j]
                    C[i] = 0
                
        fitness = CaculateFitness(X, fun)  # 计算适应度值
        fitness, sortIndex = SortFitness(fitness)  # 对适应度值排序
        X = SortPosition(X, sortIndex)  # 种群排序
        if fitness[0] <= GbestScore:  # 更新全局最优
            GbestScore = copy.copy(fitness[0])
            GbestPositon[0,:] = copy.copy(X[0, :])
        Curve[t] = GbestScore
        print(t,MaxIter)

    return GbestScore, GbestPositon, Curve



'''适应度函数'''
def fun(x):
        fitness=0
        for i in range(3):
            fitness=fitness+(x[i]**2-10*np.cos(2*np.pi*x[i])+10)
        return fitness




if __name__=="__main__":
#设置参数
    pop = 50 #种群数量
    MaxIter = 500 #最大迭代次数
    dim = 3 #维度
    lb = np.array([-5,-5,-5]) #下边界
    ub = np.array([5,5,5])#上边界
    #适应度函数选择
    fobj = fun
    GbestScore,GbestPositon,Curve = ABC(pop,dim,lb,ub,MaxIter,fobj) 
    print('最优适应度值：',GbestScore)
    print('最优解：',GbestPositon)

    #绘制适应度曲线
    plt.figure(1)
    plt.plot(Curve,'r-',linewidth=2)
    plt.xlabel('Iteration',fontsize='medium')
    plt.ylabel("Fitness",fontsize='medium')
    plt.grid()
    plt.title('ABC',fontsize='large')
    plt.show()