# 功率分配，以及重构算法
# 公有功率 + 私有功率
# 虚拟引力计算,输出合力以及各个分力
import numpy as np

from ray import *

T = 10
def get_weight_p(ue_idx,uav):
    up = np.exp(-1*uav.power_allocation[ue_idx]/T)
    sum = 0
    for i in range(len(uav.power_allocation)):
        sum += np.exp(-1*uav.power_allocation[i]/T)
    return up/sum

def get_gain_based_dist(dist2):
    pldB = pathlossA2G(dist2, 200)
    dist3 = np.linalg.norm([dist2,200])
    g = racian(dist3,200)
    pl = dB2powerratio(pldB)
    return 1/pl*g

def power_rsma_equal_allocation(BS_list,UE_list,link):
    # 链路均分
    for i in range(len(BS_list)):
        if BS_list[i].link:
            link_power = BS_list[i].power_upper*(1-BS_list[i].eta) / BS_list[i].link
            if link_power > BS_list[i].power_aver_upper*0.8:
                link_power = BS_list[i].power_aver_upper*0.8
            for j in range(len(UE_list)):
                if link[i][j] == 1:
                    BS_list[i].power_allocation[j] = link_power
                else:
                    BS_list[i].power_allocation[j] = 0
        else:
            for j in range(len(UE_list)):
                BS_list[i].power_allocation[j] = 0


def cal_rsma_sinr(UE_list,BS_list,ob=False):
    '''
    :return: 信噪比矩阵
    '''
    sinr = np.zeros([len(UE_list),len(BS_list)])*(-100)
    dist = np.zeros([len(UE_list),len(BS_list)])
    for i in range(len(BS_list)):
        for j in range(len(UE_list)):
            if (ob == False) or ( ob == True and not isNLoS(BS_list[i].loc,UE_list[j].loc,ob_params)):
                dist2 = np.sqrt((BS_list[i].loc[0]-UE_list[j].loc[0])**2+(BS_list[i].loc[1]-UE_list[j].loc[1])**2)
                dist[j][i] = dist2
                # pldB = pathlossA2G(dist2, BS_list[i].loc[2])
                # pl = dB2powerratio(pldB)
                # ray_gain = np.mean(get_rayleigh_gain())
                sinr_c = np.mean(calculate_common_sinr(j, BS_list[i],dist2))
                sinr_p = np.mean(calculate_private_sinr(j, BS_list[i],dist2))
                UE_list[j].ratec = np.log2(1+sinr_c)
                UE_list[j].ratec = np.log2(1+sinr_p)
                sinr[j][i] = sinr_p + sinr_c
                # sinr[j][i] = (BS_list[i].common_power+BS_list[i].power_allocation[j]) * (get_gain_based_dist(dist2)) / dBm2w(Noise)

    return sinr,dist

def cal_network_rsma_rate(BS_list,UE_list,linked):
    total_rate = 0

    for i in range(len(BS_list)):
        for j in range(len(linked[i])):
            tmp = linked[i][j]
            dist2 = np.sqrt((BS_list[i].loc[0] - UE_list[tmp].loc[0]) ** 2 + (BS_list[i].loc[1] - UE_list[tmp].loc[1]) ** 2)

            snr = (BS_list[i].power_allocation[tmp] / get_gain_based_dist(dist2)) / dBm2w(Noise)
            total_rate += np.log10(1 + 1 * snr) # 私有
            # print(i,tmp,snr,total_rate)
    return total_rate

def virtual_rsma_fa(UE_list,BS_list):

    snr,dist2 = cal_snr(UE_list, BS_list)

    # 产生引力的前提是建立连接，二维的力
    Fa = [] #长度：基站数
    Fa_list = [] #长度：基站*用户 Fa为Fa_list求和
    for i in range(len(BS_list)):
        Fa_tmp = []
        Fa_i = [0,0]
        for j in range(len(UE_list)):
            Fa_tmp_j = [0,0]
            # 建立连接了 则snr一定大于snr upper
            if dist2[j][i] < 30 and BS_list[i].power_allocation[j] < BS_list[
                i].power_aver_upper and BS_list[i].power < BS_list[i].power_upper:
                w = get_weight_p(j,BS_list[i])
                sinr_c = np.mean(calculate_common_sinr(j,BS_list[i],dist2[j][i]))
                sinr_p = np.mean(calculate_private_sinr(j,BS_list[i],dist2[j][i]))
                if sinr_c * sinr_p == 0:
                    continue
            # if snr[j][i] > dB2powerratio(snr_upper1) and BS_list[i].power_allocation[j] < BS_list[i].power_aver_upper and BS_list[i].power < BS_list[i].power_upper:
                Fa_i_abs = Ka * (w*sinr_p/sinr_c) /(dist2[j][i]**2 + (BS_list[i].loc[2]/100)**2) # 力的大小的计算
                Fa_tmp_j[0] = Fa_i_abs * (UE_list[j].loc[0]-BS_list[i].loc[0])/np.sqrt(dist2[j][i]**2+BS_list[i].loc[2]**2) # x轴上的大小
                Fa_tmp_j[1] = Fa_i_abs * (UE_list[j].loc[1]-BS_list[i].loc[1])/np.sqrt(dist2[j][i]**2+BS_list[i].loc[2]**2)
                Fa_i[0] += copy.deepcopy(Fa_tmp_j[0])
                Fa_i[1] += copy.deepcopy(Fa_tmp_j[1])

            Fa_tmp.append(Fa_tmp_j)
        Fa_list.append(Fa_tmp)
        Fa.append(Fa_i)
    return Fa,Fa_list
