import time

import numpy as np
import matplotlib.pyplot as plt
import math
import init
import csv
import pso_init
import UA_game
class UE:
    def __init__(self):
        self.loc = [0, 0, 0]
        self.snrth = 0
        self.uav = -1

class BS:
    def __init__(self):
        self.loc = [0, 0, 0]
        self.sort = 'UAV'
        self.power = 0

power_upper = 4
snrth = -5
power = 20 #dBm
Noise = -130
def getweight():
    # 惯性权重
    weight = 1
    return weight

def getlearningrate():
    # 分别是粒子的个体和社会的学习因子，也称为加速常数
    lr = (0.49445,1.49445)
    return lr

def getmaxgen():
    # 最大迭代次数
    maxgen = 20
    return maxgen

def getsizepop():
    # 种群规模
    sizepop = 50
    return sizepop

def getrangepop():
    # 粒子的位置的范围限制,x、y方向的限制相同
    rangepop = (0 , 100)
    #rangepop = (-2,2)
    return rangepop

def getrangespeed():
    # 粒子的速度范围限制
    rangespeed = (-2,2)
    return rangespeed

def pathlossA2G(dist2,h):
    '''
    :param dist2: 基站在水平面的投影与用户的距离
    :param h:高度
    :return:空地倾斜路径损耗
    '''
    f = 1.4e9  # MHz
    d0 = 1
    c = 3e8
    alpha = 3.5
    pi = math.pi
    A = 0.25
    C = 0.39
    E = 0.25
    G = 0
    H = 0.05
    dep_tree = 2  # 植被深度
    dist3d = math.sqrt((dist2*100)**2 + h**2)
    fspl = 20 * np.log10(4 * pi * f * d0 / c) + 10 * alpha * math.log10(dist3d / d0)
    theta = math.atan2(h, dist2*100)
    slant =  A * np.power(f / math.pow(10, 6), C) * np.power(dep_tree, E) * np.power(theta + G, H)
    return fspl + slant

# rate
# 用户 只连一个无人机
def func_fit(uav_list,ue_list):

    BS_list = []
    UE_list = []
    for i in range(len(ue_list)):
        new_ue = UE()
        new_ue.loc = [ue_list[i]]
        new_ue.uav = -1
        UE_list.append(new_ue)
    for j in range(len(uav_list)):
        new_uav = BS()
        new_uav.loc = [uav_list[j]]
        new_uav.power = power_upper
        BS_list.append(new_uav)
    # class
    snr = np.zeros([len(ue_list),len(uav_list)])
    dist = np.zeros([len(ue_list),len(uav_list)])

    for i in range(len(uav_list)):
        for j in range(len(ue_list)):
            dist[j][i] = np.sqrt((uav_list[i][0]-ue_list[j][0])**2+(uav_list[i][1]-ue_list[j][1])**2)
            pldB = pathlossA2G(dist[j][i], h=200)
            pl = dB2powerratio(pldB)
            snr[j][i] = (dBm2w(power) / pl) / dBm2w(Noise)

    utility = get_load_coalition(BS_list,UE_list,snr)
    link,linked = get_link(UE_list,BS_list)
    # 平分功率
    for i in range(len(BS_list)):
        for j in range(len(ue_list)):
            pldB = pathlossA2G(dist[j][i], h=200)
            pl = dB2powerratio(pldB)
            if(linked[i]):
                snr[j][i] = (power_upper/len(linked[i]) / pl) / dBm2w(Noise)

    cal_network_rate(BS_list, linked, snr)
    total_rate=cal_network_rate(BS_list, linked, snr)
    return total_rate

def get_load_coalition(BS_list,UE_list,snr):

    # 先清空
    for i in range(len(UE_list)):
        UE_list[i].uav = -1
    utility = np.zeros([len(BS_list),len(UE_list)])
    # 用户选择速率最高（且大于阈值）的无人机连接
    for i in range(len(UE_list)):
        idx = np.argmax(snr[i])
        if snr[i][idx] >= UE_list[i].snrth:
            UE_list[i].uav = idx
            utility[idx][i] = snr[i][idx]
    return utility


def get_link(UE_list,BS_list):
    link = np.zeros([len(BS_list),len(UE_list)])

    for j in range(len(UE_list)):
        if UE_list[j].uav > -1:
            link[UE_list[j].uav][j] = 1
    linked = []
    for i in range(len(BS_list)):
        tmp1 = []
        for j in range(len(UE_list)):
            if link[i][j] == 1:
                tmp1.append(j)
        linked.append(tmp1)
        BS_list[i].link = len(tmp1)
    # print(sum(sum(i) for i in link))
    return link,linked

def cal_network_rate(BS_list,linked,snr):

    total_rate = 0
    for i in range(len(BS_list)):
        for j in range(len(linked[i])):
            tmp = linked[i][j]
            snr_ij = snr[tmp][i]
            total_rate += np.log10(1 + 1 * snr_ij)

    return total_rate

# 初始化
# 初始位置：贪婪算法得到无人机位置
def initpopvfit(uav_list,ue_list,sizepop,rangespeed):
    # sizepop * len(uav_list)
    pop = np.zeros((sizepop,len(uav_list),2))
    v = np.zeros((sizepop,len(uav_list),2))
    fitness = np.zeros(sizepop)

    for i in range(sizepop):
        for j in range(len(uav_list)):
            pop[i][j] = [uav_list[j][0] + (np.random.rand()-0.5)*rangespeed[0]*2,uav_list[j][1] + (np.random.rand()-0.5)*rangespeed[1]*2]
            v[i][j] = [(np.random.rand()-0.5)*rangespeed[0]*2,(np.random.rand()-0.5)*rangespeed[1]*2]
        fitness[i] = func_fit(pop[i],ue_list)
    return pop,v,fitness

def getinitbest(fitness,pop):
    # 群体最优的粒子位置及其适应度值
    gbestpop,gbestfitness = pop[fitness.argmax()].copy(),fitness.max()
    #个体最优的粒子位置及其适应度值,使用copy()使得对pop的改变不影响pbestpop，pbestfitness类似
    pbestpop,pbestfitness = pop.copy(),fitness.copy()

    return gbestpop,gbestfitness,pbestpop,pbestfitness

# pl(dB) = 10*log(pr/pt)
def dB2powerratio(dB):
    return np.power(10,dB/10)
def dBm2w(dBm):
    return 0.001*np.power(10,dBm/10)

def read_csv(addr):
    ue_list = []
    with open(addr, 'r') as f:
        f = csv.reader(f, delimiter=',')
        next(f)
        for row in f:
            ue_list.append(row)
    for i in range(len(ue_list)):
        ue_list[i] = list(map(int, ue_list[i]))
    return ue_list

def main(addr):
    ue_list = read_csv(addr)
    numuav = 5
    R_uav = 20
    uav_list = pso_init.main(ue_list,R_uav)

    w = getweight()
    lr = getlearningrate()  # 学习因子
    maxgen = getmaxgen()    # 迭代次数
    sizepop = getsizepop()  # 种群规模
    rangepop = getrangepop()    # 运动范围限制？
    rangespeed = getrangespeed()    # 速度范围限制？
    start = time.time()
    pop,v,fitness = initpopvfit(uav_list,ue_list,sizepop,rangespeed)   # 位置 速度 适应度
    print(time.time()-start,'完成初始化')
    gbestpop,gbestfitness,pbestpop,pbestfitness = getinitbest(fitness,pop)  # 群体最优 和 个体最优
    print('开始迭代')
    result = np.zeros(maxgen)
    # 迭代
    for i in range(maxgen):

        t=0.5
        #速度更新
        for j in range(sizepop):
            v[j] += lr[0]*np.random.rand()*(pbestpop[j]-pop[j])+lr[1]*np.random.rand()*(gbestpop-pop[j])
        # 速度限制
        v[v<rangespeed[0]] = rangespeed[0]
        v[v>rangespeed[1]] = rangespeed[1]

        #粒子位置更新
        for j in range(sizepop):
            #pop[j] += 0.5*v[j]
            # pop[j] = t*(0.5*v[j])+(1-t)*pop[j]
            pop[j] += t * (0.5 * v[j])
        # 位置限制

        pop[pop<rangepop[0]] = rangepop[0]
        pop[pop>rangepop[1]] = rangepop[1]

        #适应度更新
        for j in range(sizepop):
            fitness[j] = func_fit(pop[j],ue_list)
        # 个体最优
        for j in range(sizepop):
            if fitness[j] > pbestfitness[j]:
                pbestfitness[j] = fitness[j]
                pbestpop[j] = pop[j].copy()
        # 群体最优
        if pbestfitness.max() > gbestfitness :
            gbestfitness = pbestfitness.max()
            gbestpop = pop[pbestfitness.argmax()].copy()
        # 第i次迭代的最大适应度
        result[i] = gbestfitness
        # print(i,time.time()-start,result[i])
    # rate = []
    # for i in range(len(result)):
    #     tmp = [i, result[i]]
    #     rate.append(tmp)
    # csv_name = 'pso_solution_Mar.csv'
    # csv_file = open(csv_name, 'w', newline='')
    # writer = csv.writer(csv_file)
    # writer.writerows(rate)
    # csv_file.close()
    # plt.plot(result)
    # plt.show()
    return len(uav_list),result[-1],time.time()-start
if __name__ == '__main__':
    addr = r"D:\pycharm\deploy_algorithms\data\list_ue140_1.csv"
    print(main(addr))
