import math
import Users
import numpy as np
import random
from pprint import pprint


class PSO:
    def __init__(self,users,cal_tool,N=60):
        #之后解析JSON，从JSON中读参数
        self.N=N
        self.cal_tool=cal_tool
        self.d = 3  # 可行解维数
        self.limit = [0, 1000]  # 边界位置限制
        self.vlimit = [-100, 100]  # 边界速度限制
        self.height = 100
        self.wmax = 1.5  # 惯性权重
        self.wmin = 0.3
        self.c1 = 2  # 自我学习因子
        self.c2 = 2  # 群体学习因子
        # self.c1 = 0.5  # 自我学习因子
        # self.c2 = 0.5  # 群体学习因子
        self.maxgen = 200
        # self.users = Users.getUsers("./data.csv")
        self.users = users


    # #todo：防止无人机位置碰撞
    # def punish_function(self,uavs):
    #     # user_arr = [1, 2, 5]
    #     punish=0
    #     for i in range(len(uavs)):
    #         for j in range(i+1,len(uavs)):
    #             dis=self.cal_tool.getDistance(uavs[i],uavs[j])
    #             if dis<200:
    #                 punish+=100
    #
    #     return punish


    # 初始化粒子群，k是无人机数量
    def initParticle(self, k, users):
        X = []
        # users=Users.generate(100)#用户位置
        V = np.random.uniform(self.vlimit[0], self.vlimit[1], (self.N, k, 3))
        for i in range(self.N):
            uavs=[]
            for index in range(k):
                x=np.random.randint(self.limit[0],self.limit[1])
                y=np.random.randint(self.limit[0],self.limit[1])
                z=self.height
                uavs.append([x,y,z])
            # uavs = self.cal_tool.getUavInitialPos(k, users,self.height[0])
            X.append(uavs)
        X = np.array(X)
        return X, V

    # 初始化个体和全局最佳历史位置和最佳适应度
    def initBest(self, X, k):
        p_pos = X.copy()  # 每个个体的历史最佳位置
        global_pos = np.zeros((k, 3))  # 种群的历史最佳位置5*3
        p_best = np.full((self.N, 1), float('-inf'))  # 每个个体的历史最佳适应度
        global_best = float('-inf')  # 种群历史最佳适应度
        return p_pos, global_pos, p_best, global_best


    def pso(self, X, V, k,get_fitness):
        # X, V, users = initParticlePos()
        p_pos, global_pos, p_best, global_best = self.initBest(X, k)
        result = np.zeros(self.maxgen)

        for iter in range(self.maxgen):
            # fitness = self.getfitness(X, users, k,P)  # 计算所有粒子适应度
            fitness = np.zeros(self.N)
            for i in range(len(X)):
                _,_,fitness[i]= get_fitness(self.users,X[i])
            for i in range(self.N):
                if fitness[i] > p_best[i]:
                    p_best[i] = fitness[i]  # 更新个体历史最佳适应度
                    p_pos[i] = X[i].copy()  # 更新个体历史最佳位置(取值)
            if max(p_best) > global_best:
                global_best = max(p_best)  # 更新群体历史最佳适应度
                max_index = p_best.argmax()
                global_pos = X[max_index].copy()  # 更新群体历史最佳位置


            # 权重更新
            w = self.wmax - iter * (self.wmax - self.wmin) / self.maxgen
            for i in range(self.N):
                V[i] = V[i] * w + self.c1 * random.random() * (p_pos[i] - X[i]) + self.c2 * random.random() * (
                        global_pos - X[i])
                # print(V[i])
                for j in range(len(V[i])):
                    for p in range(len(V[i][j])):
                        if V[i][j][p] > self.vlimit[1]:
                            V[i][j][p] = self.vlimit[1]
                            # V[i][j][p] = self.vlimit[0]+np.random.rand(1)*(self.vlimit[1]-self.vlimit[0])
                        elif V[i][j][p] < self.vlimit[0]:
                            V[i][j][p] = self.vlimit[0]
                # print(V[i])

            # 位置更新
            X = X + V
            # 边界位置处理
            for i in range(self.N):
                for j in range(len(X[i])):
                    for p in range(2):
                        if X[i][j][p] > self.limit[1]:
                            X[i][j][p] = self.limit[1]
                        elif X[i][j][p] < self.limit[0]:
                            X[i][j][p] = self.limit[0]
                    if X[i][j][2] > self.height:
                        X[i][j][2] = self.height
                    elif X[i][j][2] < self.height:
                        X[i][j][2] = self.height

            result[iter] = global_best
            # print(iter,global_pos,global_best)

        # print(global_best)
        # print(global_pos)

        return global_pos,global_best,result






if __name__ == '__main__':
    print()




