import init
from virtual_force import *
from UA_game import *
from ratio_pso import optimize_common_power_eta
from rsma import *
from icsi_rician import get_total_rate_icsi
numfbs = 0
nummbs = 0
sumgbs = nummbs + numfbs
R_mbs = 25
R_uav = 25
snrth = -5
itermax = 100
init_upper = 20
num1 = 1.01
num2 = 0
num3 = 5
num4 = 0
font_set = matplotlib.font_manager.FontProperties(fname=r"HYQiHei-60S.ttf")
obstacle = get_obstacle(ob_params)

class UE:
    def __init__(self):
        self.loc = [0, 0, 0]
        self.snrth = 0
        self.uav = -1
        self.ratec = 0
        self.ratep = 0
        self.sinr_c = 0
        self.sinr_p = 0

class BS:
    def __init__(self):
        self.loc = [0, 0, 0]
        self.sort = 'UAV'
        self.power_upper = 0
        self.power_allocation = []  # 长度为用户数
        self.power_a2a = 0
        self.power = 0
        self.rate = 0
        self.rate_upper = 0


def read_UElist(ue_addr):
    ue_list = []
    with open(ue_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(float, ue_list[i]))
    UE_list = []
    for i in range(len(ue_list)):
        new_ue = UE()
        new_ue.loc = [ue_list[i][0], ue_list[i][1], 0]
        new_ue.sort = 'ue'
        new_ue.uav = -1
        new_ue.rate = 0
        new_ue.sinrc = 0
        new_ue.sinrp = 0
        new_ue.snrth = dB2powerratio(snrth)
        UE_list.append(new_ue)
    return ue_list,UE_list


# BS_list,UE_list ->matrix
def get_km_link_matrix(BS_list,UE_list,snr):
    matrix = []

    for i in range(len(UE_list)):
        for j in range(len(UE_list)):
            if j < len(BS_list):
                if snr[i][j] >= UE_list[i].snrth:
                    matrix.append((i,j,snr[i][j]))
                else:
                    matrix.append((i,j,0))
            else:
                matrix.append((i,j,0))

    return matrix
def main(addr,K=1,SIGMA_E=0):
    # 导入用户位置
    ue_list,UE_list = read_UElist(addr)
    # print(len(UE_list), UE_list[0].loc)
    # 确定用户通信需求 各用户所需的信噪比
    numuav = 4
    R = 23
    BS_list = init.main_rsma(UE_list.copy(), sumgbs, R_mbs=R_mbs, sumuav=numuav, R_uav=R_uav)
    # BS_list = KMeans.main(ue_list,R)
    snr, dist = cal_rsma_sinr(UE_list, BS_list)

    # link, linked = comm_connection(UE_list, BS_list, snr)  # 未建立连接的确保发射功率为0
    link,linked = game_process(BS_list, UE_list, snr,dist)
    power_rsma_equal_allocation(BS_list, UE_list, link)
    coverate = cal_coverate(link)

    times = 0
    init_times = 0
    total_rate = []
    total_rate1 = []

    start = time.time()
    # 求最小数量
    while (coverate != 1):

        # 位置调整
        update_bs_power(BS_list)
        Fa, Fa_list = virtual_rsma_fa(UE_list, BS_list)
        Fr, Fr_list = virtual_fr(BS_list)
        Fa2, Fa_list2 = virtual_fa(UE_list, BS_list)
        F = total_force(Fa, Fr)
        # print(len(Fa),len(Fa[0]),len(Fa2),len(Fa2[0]),Fa[0][0],Fa2[0][0])
        v = force2v(F, Max_Speed=1)
        move_length = update_state_position(BS_list, v)
        snr, dist = cal_rsma_sinr(UE_list, BS_list)
        # link, linked = comm_connection(UE_list, BS_list, snr)  # 确立连接关系 未建立连接的确保功率为0
        link,linked = game_process(BS_list, UE_list, snr,dist)
        coverate = cal_coverate(link)
        if (coverate == 1):
            print(times,coverate,linked, '实现全覆盖，退出迭代')
            break
        # 重新初始化
        if (coverate != 1) and init_times <= init_upper:

            print('此次无法全覆盖，重新初始化', numuav, coverate)
            BS_list = init.main_rsma(UE_list.copy(), sumgbs, R_mbs=R_mbs, sumuav=numuav, R_uav=R_uav)

            snr, dist = cal_rsma_sinr(UE_list, BS_list)
            link, linked = game_process(BS_list, UE_list, snr, dist)
            power_rsma_equal_allocation(BS_list, UE_list, link)
        elif (coverate != 1):
            print('无法全覆盖，增加无人机数量并重新初始化', numuav, coverate)
            init_times = 0
            numuav += 1

            BS_list = init.main_rsma(UE_list.copy(), sumgbs, R_mbs=R_mbs, sumuav=numuav, R_uav=R_uav)
            snr, dist = cal_rsma_sinr(UE_list, BS_list)
            # link, linked = comm_connection(UE_list, BS_list, snr)  # 确立连接关系 未建立连接的确保功率为0

            link,linked = game_process(BS_list, UE_list, snr,dist)
            power_rsma_equal_allocation(BS_list, UE_list, link)
            coverate = cal_coverate(link)

            draw_uavus_link_circle2(UE_list.copy(),BS_list.copy(),linked,basecolor='red',size=40,)
        init_times += 1  # """

    # plt.figure()
    # for i in range(len(UE_list)):
    #     plt.scatter(UE_list[i].loc[0], UE_list[i].loc[1], color='blue', marker='x', s=40)
    # plt.scatter(UE_list[i].loc[0], UE_list[i].loc[1], color='blue', marker='x', s=40, label='UE')
    # plt.xlim(-10, 110)
    # plt.ylim(-10, 110)
    # draw_obstacle(obstacle)
    # draw_uavus_link(UE_list, BS_list, linked, basecolor='red', size=40, )
    uav_posi = []
    uav_power = []
    for i in range(len(BS_list)):
        uav_posi.append(BS_list[i].loc)
        uav_power.append(BS_list[i].power_allocation)

    # plt.legend(loc='upper right', borderaxespad=num4, prop=font_set)
    # draw_uavus_link_circle2(UE_list.copy(), BS_list.copy(), linked, basecolor='red', size=40, )


    power_rsma_equal_allocation(BS_list, UE_list, link)
    BS_list1 = copy.deepcopy(BS_list)
    rate_tmp = cal_network_rate(BS_list, UE_list, linked)
    rate_1 = get_total_rate_icsi(BS_list,UE_list,dist,K=K,SIGMA_E=SIGMA_E)#cal_network_rate_rsma(BS_list, UE_list, linked)
    rate_2 = get_total_rate_icsi(BS_list,UE_list,dist,K=K,SIGMA_E=SIGMA_E)#cal_network_rate_rsma(BS_list, UE_list, linked)
    total_rate.append(rate_1)
    total_rate1.append(rate_2)
    matrix = get_km_link_matrix(BS_list, UE_list, snr)
    print('开始联调')
    # return len(BS_list), total_rate[-1]

    while (times < 15):
    # while (len(total_rate) < itermax):
        optimize_common_power_eta(UE_list,BS_list)
        # print(BS_list[0].power_allocation)
        # print(get_total_rate_icsi(BS_list,UE_list,dist))
        power_ = []
        # 位置调整
        Fa, Fa_list = virtual_fa(UE_list, BS_list)
        Fr, Fr_list = virtual_fr(BS_list)

        F = total_force(Fa, Fr)
        v = force2v(F, 5)
        #print(v)
        move_length = update_state_position(BS_list, v)
        snr, dist = cal_rsma_sinr(UE_list, BS_list, ob=True)
        # link, linked = comm_connection(UE_list, BS_list, snr)  # 确立连接关系 未建立连接的确保功率为0
        link, linked = game_process(BS_list, UE_list, snr,dist)
        update_bs_power(BS_list)

        # 功率调整
        Fa, Fa_list = virtual_fa(UE_list, BS_list)
        Fr, Fr_list = virtual_fr(BS_list)

        va = force2v(Fa, Max_Speed=power_step)
        va_list = forcelist2vlist(Fa_list, 1)
        vr_list = power2v(total_fr_power(Fr_list), 1)
        update_state_power(BS_list, va_list, vr_list, link)
        # update_state_power2(BS_list, va, vr_list, link)
        update_bs_power(BS_list)
        for i in range(len(BS_list)):
            power_.append(BS_list[i].power)
        snr, dist = cal_rsma_sinr(UE_list, BS_list,ob=True)
        link, linked = game_process(BS_list, UE_list, snr,dist)
        # link, linked = comm_connection(UE_list, BS_list, snr)  # 确立连接关系 未建立连接的确保功率为0
        rate_tmp = get_total_rate_icsi(BS_list,UE_list,dist,K=K,SIGMA_E=SIGMA_E)#cal_network_rate_rsma(BS_list, UE_list, linked)
        # print(times,rate_tmp)

        # if rate_tmp > total_rate[-1]:
        total_rate.append(rate_tmp)
        # total_rate.append(cal_network_rate(BS_list, UE_list, linked))
        # print(times,total_rate[-1])
        # print(times, coverate)

        # if move_length < 1:
        #    break
        times += 1
    time1 = time.time() - start
    linked_ = copy.deepcopy(linked)
    uav_posi = []
    uav_power = []
    for i in range(len(BS_list)):
        uav_posi.append(BS_list[i].loc)
        uav_power.append(BS_list[i].power_allocation)
    print(total_rate)
    # print(uav_posi,cal_network_rate(BS_list, UE_list, linked))
    # """
    # 只进行位置调整
    """print('位置调整')
    times = 0
    # 只进行位置调整
    vforce2.power_equal_allocation(BS_list1, UE_list, linked, link)
    while (times < itermax):
        # 位置调整
        Fa, Fa_list = vforce2.virtual_fa(UE_list, BS_list1)
        Fr, Fr_list = vforce2.virtual_fr(BS_list1)

        F = total_force(Fa, Fr)
        v = force2v(F, 0.2)
        move_length = update_state_position(BS_list1, v)
        snr, dist = cal_snr(UE_list, BS_list1)
        link, linked = comm_connection(UE_list, BS_list1, snr)  # 确立连接关系 未建立连接的确保功率为0
        update_bs_power(BS_list1)
        total_rate1.append(cal_network_rate(BS_list1, UE_list, linked))
        # print(times,total_rate1[-1])
        times += 1
    time2 = time.time() - start - time1"""
    """
    while (times < itermax):

        # 位置调整
        Fa, Fa_list = vforce2.virtual_fa(UE_list, BS_list1)
        Fr, Fr_list = vforce2.virtual_fr(BS_list1)

        F = vforce2.total_force(Fa, Fr)
        v = vforce2.force2v(F,0.2)
        move_length = vforce2.update_state_position(BS_list1, v)
        snr, dist = vforce2.cal_snr(UE_list, BS_list1)
        link, linked = vforce2.comm_connection(UE_list, BS_list1, snr)  # 确立连接关系 未建立连接的确保功率为0
        vforce2.power_equal_allocation(BS_list1, UE_list, linked, link)
        vforce2.update_bs_power(BS_list1)

        total_rate1.append(vforce2.cal_network_rate(BS_list1, UE_list,linked))

        #print(times, coverate)

        # if move_length < 1:
        #    break
        times += 1
    # jointly optimize#"""

    """while (times < itermax):
        power_ = []
        # 位置调整
        Fa, Fa_list = virtual_fa(UE_list, BS_list1)
        Fr, Fr_list = virtual_fr(BS_list1)

        F = total_force(Fa, Fr)
        v = force2v(F)
        move_length = update_state_position(BS_list1, v)
        snr, dist = cal_snr(UE_list, BS_list1)
        link, linked = comm_connection(UE_list, BS_list1, snr)  # 确立连接关系 未建立连接的确保功率为0
        update_bs_power(BS_list1)

        total_rate1.append(cal_network_rate(BS_list1, UE_list,linked))

        print(times, coverate)

        # if move_length < 1:
        #    break
        times += 1  # """

    """rate = []
    for i in range(len(total_rate)):
        tmp = [i, total_rate[i],total_rate1[i]]
        rate.append(tmp)
    csv_name = 'va_joint_solution.csv'
    csv_file = open(csv_name, 'w', newline='')
    writer = csv.writer(csv_file)
    writer.writerows(rate)
    csv_file.close()"""

    # """
    # print('用时', total_time,total_rate,BS_list[0].power_allocation,BS_list[0].power)
    # 加KM

    # plt.figure()
    # for i in range(len(UE_list)):
    #     plt.scatter(UE_list[i].loc[0], UE_list[i].loc[1], color='blue', marker='x', s=40)
    # plt.scatter(UE_list[i].loc[0], UE_list[i].loc[1], color='blue', marker='x', s=40, label='UE')
    # plt.xlim(-10, 110)
    # plt.ylim(-10, 110)
    #
    # draw_obstacle(obstacle)
    # draw_uavus_link(UE_list, BS_list, linked, basecolor='red', size=40, )
    # print(len(BS_list), coverate, cal_network_rate(BS_list, UE_list, linked))
    # # pic_name = 'deployprocess'
    # # title = pic_name + str(numuav) + '.jpg'
    # plt.legend(loc='upper right', borderaxespad=num4, prop=font_set)
    # draw_uavus_link_circle2(UE_list.copy(), BS_list.copy(), linked, basecolor='red', size=40, name='resul')
    #
    # plt.figure()
    # for i in range(len(UE_list)):
    #     plt.scatter(UE_list[i].loc[0], UE_list[i].loc[1], color='blue', marker='x', s=40)
    # plt.xlim(-10, 110)
    # plt.ylim(-10, 110)
    #
    # draw_uav_update(BS_list, basecolor='k', LoSscale=1, size=20)
    # plt.figure()
    # for i in range(len(UE_list)):
    #     plt.scatter(UE_list[i].loc[0], UE_list[i].loc[1], color='blue', marker='x', s=40)
    # plt.scatter(UE_list[i].loc[0], UE_list[i].loc[1], color='blue', marker='x', s=40, label='UE')
    # plt.xlim(-10, 110)
    # plt.ylim(-10, 110)
    #
    # draw_obstacle(obstacle)
    # draw_uavus_link(UE_list, BS_list, linked, basecolor='red', size=40, )
    # plt.figure()
    #
    # x = np.arange(len(total_rate))
    # #plt.plot(x, total_rate1, 'b-', label='optimize position only')
    # plt.plot(x, total_rate, 'g-', label='jointly optimize position ans power')
    # plt.xlabel('iterations', fontproperties=font_set)
    # plt.ylabel('system throughput（bit/s·Hz）', fontproperties=font_set)
    # plt.legend(loc='center right', borderaxespad=num4, prop=font_set)
    # draw_uavus_link_circle2(UE_list.copy(), BS_list.copy(), linked_, basecolor='red', size=40,name='result_' )
    # plt.show()  # """
    # # print(total_rate[-1], total_rate1[-1])
    return len(BS_list), max(total_rate), time1,
    # return len(BS_list), total_rate[-1], total_rate1[-1], time1,


if __name__ == '__main__':
    addr = r"D:\pycharm\deploy_algorithms\data\list_ue100_6.csv"
    result = main(addr)
    print(result)
