import math
import random
import warnings

import tool
import Users
import cvxpy as cp
import numpy as np
import constants as C


def cal_b_u(users,uavs,uav_index,uav_to_user,r_c,p_alloc,B_max,R_min,u=1.05):
    users_index = uav_to_user[uav_index]
    n=len(users_index) # 变量数量
    # 定义变量
    x = cp.Variable(n, pos=True)  # 非负变量
    c=[]# 目标函数系数
    for i in users_index:
        p=p_alloc[i]
        snr = tool.cal_snr_f(uavs[uav_index], users[i], p)
        factor=math.log2(1 + snr)
        c.append(factor)
    c = np.array(c)
    # print(c)
    # 定义目标函数
    objective = cp.Maximize(c @ x)
    constant_limit = B_max  # 约束条件中的常数上限
    # 定义约束条件
    constraints = [cp.sum(x) <= constant_limit]
    min_cons=[]
    # u=1.05# 1.05
    for i in range(len(users_index)):
        index=users_index[i]
        p = p_alloc[index]
        snr = tool.cal_snr_f(uavs[uav_index], users[index], p)
        lower_limit = R_min*(r_c[index])*u/math.log2(1 + snr)
        lower_limit=round(lower_limit, 4)
        # print(lower_limit)
        min_cons.append(lower_limit)
        constraints.append(x[i] >= lower_limit)
    # 构建问题并求解
    problem = cp.Problem(objective, constraints)
    problem.solve()

    # 输出求解结果
    # print("带宽分配后速率最大值：", problem.value)
    # print("带宽最优解：", x.value)
    # rate=[]
    # for i in range(len(users_index)):
    #     rate.append(c[i]*x.value[i])
    # print("速率：",rate)
    b_alloc = {}
    if problem.status == cp.OPTIMAL:
        b_alloc_list=x.value
        for i in range(len(users_index)):
            b_alloc[users_index[i]]=b_alloc_list[i]
        return b_alloc, problem.value
    else:
        # todo:没有最优解的情况
        sum_weight = 0
        sum_rate=0
        for i in users_index:
            sum_weight += r_c[i]
        for i in users_index:
            # b_alloc[i] = C.B_max * r_c[i] / sum_weight
            b_alloc[i]=C.B_max/len(users_index)
            cur_rate = tool.cal_rate(b_alloc[i], p_alloc[i], uavs[uav_index], users[i])
            sum_rate += cur_rate
        return b_alloc,sum_rate
        # lower_cons={}
        # for i in range(len(users_index)):
        #     lower_cons[min_cons[i]]=users_index[i]
        # sorted_list=sorted(min_cons)
        # sum_b=0
        # cur_index=len(sorted_list)
        # for i in range(len(sorted_list)):
        #     if sorted_list[i]+sum_b<=B_max:
        #         sum_b+=sorted_list[i]
        #         index=lower_cons[sorted_list[i]]
        #         b_alloc[index]=sorted_list[i]
        #     else:
        #         cur_index=i
        #         break
        # num_b=0
        # rest_b=B_max-sum_b
        # for i in range(cur_index,len(sorted_list)):
        #     num_b+=sorted_list[i]
        # for i in range(cur_index, len(sorted_list)):
        #     index=lower_cons[sorted_list[i]]
        #     b_alloc[index]=rest_b*sorted_list[i]/num_b

def cal_b_maxmin(users,uavs,uav_index,uav_to_user,r_c,p_alloc,B_max,R_min):
    users_index = uav_to_user[uav_index]
    n=len(users_index) # 变量数量
    # 定义变量
    x = cp.Variable(n, nonneg=True)  # 非负变量
    min_rate=cp.Variable()
    c=[]# 目标函数系数
    constant_limit = B_max  # 约束条件中的常数上限
    # 定义约束条件
    constraints = [cp.sum(x) <= constant_limit]
    for i in range(len(users_index)):
        index=users_index[i]
        p=p_alloc[index]
        snr = tool.cal_snr_f(uavs[uav_index], users[index], p)
        factor=math.log2(1 + snr)
        constraints.append(min_rate<=cp.multiply(x[i],factor))
        c.append(factor)
    c = np.array(c)
    print(c)
    # 定义目标函数
    objective = cp.Maximize(min_rate)

    min_cons=[]
    for i in range(len(users_index)):
        index=users_index[i]
        p = p_alloc[index]
        snr = tool.cal_snr_f(uavs[uav_index], users[index], p)
        lower_limit = R_min*r_c[index]/math.log2(1 + snr)
        # lower_limit=round(lower_limit, 2)
        print(lower_limit)
        min_cons.append(lower_limit)
        constraints.append(x[i] >= lower_limit)
    # 构建问题并求解
    problem = cp.Problem(objective, constraints)
    problem.solve()

    # 输出求解结果
    # print("带宽分配后最大最小速率为：", problem.value)
    # print("带宽最优解：", x.value)
    rate=[]
    for i in range(len(users_index)):
        rate.append(c[i]*x.value[i])
    # print("速率：", rate)
    b_alloc = {}
    if problem.status != cp.INFEASIBLE:
        b_alloc_list=x.value
        for i in range(len(users_index)):
            b_alloc[users_index[i]]=b_alloc_list[i]
    else:
        # todo:没有最优解的情况,将所需带宽从小到大排序，依次分配，直到剩下的带宽不够分配，然后将剩余带宽按所需比例分配其他人
        lower_cons={}
        for i in range(len(users_index)):
            lower_cons[min_cons[i]]=users_index[i]
        sorted_list=sorted(min_cons)
        sum_b=0
        cur_index=len(sorted_list)
        for i in range(len(sorted_list)):
            if sorted_list[i]+sum_b<=B_max:
                sum_b+=sorted_list[i]
                index=lower_cons[sorted_list[i]]
                b_alloc[index]=sorted_list[i]
            else:
                cur_index=i
                break
        num_b=0
        rest_b=B_max-sum_b
        for i in range(cur_index,len(sorted_list)):
            num_b+=sorted_list[i]
        for i in range(cur_index, len(sorted_list)):
            index=lower_cons[sorted_list[i]]
            b_alloc[index]=rest_b*sorted_list[i]/num_b

    return b_alloc,problem.value
def cal_p_by_CVX_u(users,uavs,uav_index,uav_to_user,r_c,b_alloc,P_rest,R_min,snr_thre,u=1.05):
    users_index = uav_to_user[uav_index]
    # 定义变量和常数
    n = len(users_index)  # 变量数量
    a = []  # 目标函数系数
    for i in users_index:
        a.append(b_alloc[i])
    a=np.array(a)
    constant_limit = P_rest  # 约束条件中的常数上限
    c=[]
    for i in users_index:
        loss = tool.forest_path_loss(uavs[uav_index], users[i])
        factor=(10 ** (-loss / 10))/((10 ** (C.noise / 10))/1000)
        c.append(factor)
    c=np.array(c)
    # 定义变量
    x = cp.Variable(n, pos=True)  # 正数变量
    # 定义目标函数
    log_terms = [cp.log(1+c[i]*x[i])/np.log(2) for i in range(n)]  # log 函数每项
    objective = cp.Maximize(a @ log_terms)
    # u=1.05
    # 定义约束条件
    constraints = [cp.sum(x) <= constant_limit]
    for i in range(len(users_index)):
        index=users_index[i]
        loss = tool.forest_path_loss(uavs[uav_index],users[index])
        lower_limit=10**(snr_thre/10)*((10 ** (C.noise / 10))/1000) / (10 ** (-loss / 10))
        constraints.append(x[i] >= lower_limit)
        number = R_min * (r_c[index])*u/ b_alloc[index]
        noise_power = (10 ** (C.noise / 10)) / 1000
        lower_limit2=(math.pow(2,number)-1)*noise_power/10 ** (-loss / 10)
        constraints.append(x[i] >= lower_limit2)

    # 构建问题并求解
    problem = cp.Problem(objective, constraints)
    problem.solve()
    p_alloc={}
    if problem.status == cp.OPTIMAL:
        # 输出求解结果
        # print("功率分配后最大值：", problem.value)
        # print("功率最优解：", x.value)
        p_alloc_list = x.value
        for i in range(len(users_index)):
            p_alloc[users_index[i]] = p_alloc_list[i]
        return p_alloc,problem.value
    else:
        sum_weight = 0
        sum_rate=0

        for i in users_index:
            sum_weight += r_c[i]
        for i in users_index:
            # p_alloc[i] = P_rest * r_c[i] / sum_weight
            p_alloc[i]=P_rest/len(users_index)
            cur_rate = tool.cal_rate(b_alloc[i], p_alloc[i], uavs[uav_index], users[i])
            sum_rate+=cur_rate

        return p_alloc,sum_rate

        # p_alloc=cal_p(users,uavs,uav_index,uav_to_user,r_c,b_alloc,P_rest,R_min,snr_thre)
    # print(p_alloc)


def cal_p(users,uavs,uav_index,uav_to_user,r_c,b_alloc,P_rest,R_min,snr_thre):
    users_index = uav_to_user[uav_index]
    # 人数
    n = len(users_index)
    p_alloc={}
    num=0
    for i in users_index:
        loss = tool.forest_path_loss(uavs[uav_index], users[i])
        noise_power = (10 ** (C.noise / 10)) / 1000
        num+=noise_power/10 ** (-loss / 10)
    sum_p=0
    cur_index=0
    for i in users_index:
        # lamda=n*b_alloc[i]/((P_rest+num)*math.log(2))
        loss = tool.forest_path_loss(uavs[uav_index], users[i])
        noise_power = (10 ** (C.noise / 10)) / 1000
        p1=(P_rest+num)/n-noise_power/10 ** (-loss / 10)
        p2=tool.get_p_by_snr_f(snr_thre,uavs[uav_index],users[i],C.noise)
        number=R_min*r_c[i]/b_alloc[i]
        p3=(math.pow(2,number)-1)*noise_power/10 ** (-loss / 10)
        max_p=max(max(p1,p2),p3)
        p_alloc[i]=max_p
        sum_p+=max_p
        if p1>p3:
            cur_index=i
        print(p1,p2,p3)
    print(sum_p)
    if sum_p-P_rest>0:
        p_alloc[cur_index]-=sum_p-P_rest
    print(p_alloc)
    return p_alloc

def resource_alloc_u(users,uavs,r_c,p_connection,uav_to_user):
    with warnings.catch_warnings():
        # 忽略特定类型的警告
        warnings.filterwarnings("ignore", message="Solution may be inaccurate.", category=UserWarning)

        best=float('-inf')
        p_rest=tool.get_rest_p(users,uavs,p_connection)
        # p_alloc={}
        # b_alloc={}
        #平分功率，带宽
        p_all = []
        b_all = []
        best_all = []
        rate=[]
        for i in range(len(uavs)):
            p_alloc={}
            b_alloc={}
            sub=1000
            for j in uav_to_user[i]:
                b_alloc[j] = C.B_max / len(uav_to_user[i])
                p_alloc[j] = p_rest[i] / 4
            while sub>1:
                # b_alloc, cur_objective = cal_b(users, uavs, i, uav_to_user, r_c, p_alloc, C.B_max, C.R_min)
                p_alloc,best=cal_p_by_CVX_u(users, uavs, i, uav_to_user, r_c, b_alloc, p_rest[i], C.R_min, C.snr_thre,1.05)
                b_alloc, cur_objective = cal_b_u(users, uavs, i, uav_to_user, r_c, p_alloc, C.B_max, C.R_min,1.05)
                sub=abs(best-cur_objective)
                # best=max(best,cur_objective)
            for j in uav_to_user[i]:
                cur_rate = tool.cal_rate(b_alloc[j], p_alloc[j], uavs[i], users[j])
                rate.append(cur_rate)
            p_all.append(p_alloc)
            b_all.append(b_alloc)
            best_all.append(best)
        # print("功率：",p_all)
        # print("带宽：",b_all)
        # print("所有速率：",best_all)
        return p_all,b_all,rate

def resource_alloc(users,uavs,r_c,p_connection,uav_to_user):
    with warnings.catch_warnings():
        # 忽略特定类型的警告
        warnings.filterwarnings("ignore", message="Solution may be inaccurate.", category=UserWarning)

        best=float('-inf')
        p_rest=tool.get_rest_p(users,uavs,p_connection)
        # p_alloc={}
        # b_alloc={}
        #平分功率，带宽
        p_all = []
        b_all = []
        best_all = []
        rate=[]
        for i in range(len(uavs)):
            p_alloc={}
            b_alloc={}
            sub=1000
            for j in uav_to_user[i]:
                b_alloc[j] = C.B_max / len(uav_to_user[i])
                p_alloc[j] = p_rest[i] / 4
            while sub>1:
                # b_alloc, cur_objective = cal_b(users, uavs, i, uav_to_user, r_c, p_alloc, C.B_max, C.R_min)
                p_alloc,best=cal_p_by_CVX_u(users, uavs, i, uav_to_user, r_c, b_alloc, p_rest[i], C.R_min, C.snr_thre,1)
                b_alloc, cur_objective = cal_b_u(users, uavs, i, uav_to_user, r_c, p_alloc, C.B_max, C.R_min,1)
                sub=abs(best-cur_objective)
                # best=max(best,cur_objective)
            for j in uav_to_user[i]:
                cur_rate = tool.cal_rate(b_alloc[j], p_alloc[j], uavs[i], users[j])
                rate.append(cur_rate)
            p_all.append(p_alloc)
            b_all.append(b_alloc)
            best_all.append(best)
        # print("功率：",p_all)
        # print("带宽：",b_all)
        # print("所有速率：",best_all)
        return p_all,b_all,rate

def resource_alloc_proportion(users,uavs,r_c,p_connection,uav_to_user):
    best = float('-inf')
    p_rest = tool.get_rest_p(users, uavs, p_connection)

    p_all=[]
    b_all=[]
    sum_weight=[0 for i in range(len(uavs))]
    rate=[]
    for i in range(len(uavs)):
        for j in uav_to_user[i]:
            sum_weight[i] += r_c[j]
    for i in range(len(uavs)):
        p_alloc = {}
        b_alloc = {}
        for j in uav_to_user[i]:
            p_alloc[j] = p_rest[i]*r_c[j] /sum_weight[i]
            b_alloc[j]=C.B_max*r_c[j] /sum_weight[i]
            cur_rate=tool.cal_rate(b_alloc[j],p_alloc[j],uavs[i],users[j])
            rate.append(cur_rate)
        p_all.append(p_alloc)
        b_all.append(b_alloc)
    # print(p_alloc,b_alloc)

    return p_all,b_all,rate

def resource_alloc_avg(users,uavs,r_c,p_connection,uav_to_user):
    best = float('-inf')
    p_rest = tool.get_rest_p(users, uavs, p_connection)

    p_all=[]
    b_all=[]
    sum_weight=[0 for i in range(len(uavs))]
    rate=[]
    for i in range(len(uavs)):
        for j in uav_to_user[i]:
            sum_weight[i] += r_c[j]
    for i in range(len(uavs)):
        p_alloc = {}
        b_alloc = {}
        for j in uav_to_user[i]:
            p_alloc[j] = p_rest[i]/len(uav_to_user[i])
            b_alloc[j]=C.B_max/len(uav_to_user[i])
            cur_rate=tool.cal_rate(b_alloc[j],p_alloc[j],uavs[i],users[j])
            rate.append(cur_rate)
        p_all.append(p_alloc)
        b_all.append(b_alloc)
    # print(p_alloc,b_alloc)

    return p_all,b_all,rate
#非凸求解不了
def resource_CVX(users,uavs,uav_index,uav_to_user,r_c,B_max,P_rest,R_min,snr_thre):
    users_index = uav_to_user[uav_index]
    # 人数
    n = len(users_index)
    b = cp.Variable(n,pos=True)  # 正数变量)
    p = cp.Variable(n,pos=True)
    c = []
    log_terms = 0  # 将 log_terms 初始化为 0
    for i in range(len(users_index)):
        index=users_index[i]
        loss = tool.forest_path_loss(uavs[uav_index], users[index])
        factor = (10 ** (-loss / 10)) / ((10 ** (C.noise / 10)) / 1000)
        c.append(factor)
        log_terms += b[i]*cp.log(1 + factor * p[i]) / np.log(2)
    c = np.array(c)
    # log_terms = [cp.log(1 + c[i] * p[i]) / np.log(2) for i in range(n)]  # log 函数每项
    # objective = cp.Maximize(cp.sum(cp.multiply(b,log_terms)))
    # log_terms = [cp.log(1 + c[i] * p[i]) / np.log(2) for i in range(n)]  # log 函数每项
    objective = cp.Maximize(log_terms)
    constraints = [cp.sum(p) <= P_rest,cp.sum(b) <= B_max]
    for i in range(len(users_index)):
        index = users_index[i]
        loss = tool.forest_path_loss(uavs[uav_index], users[index])
        lower_limit = 10 ** (snr_thre / 10) * ((10 ** (C.noise / 10)) / 1000) / (10 ** (-loss / 10))
        constraints.append(p[i] >= lower_limit)
        # number = R_min * r_c[index] / b[i]
        noise_power = (10 ** (C.noise / 10)) / 1000
        # lower_limit2 = (math.pow(2, R_min * r_c[index] / b[i]) - 1) * noise_power / 10 ** (-loss / 10)
        factor = (10 ** (-loss / 10)) / ((10 ** (C.noise / 10)) / 1000)
        constraints.append( b[i]*cp.log(1 + factor * p[i]) / np.log(2)>=R_min * r_c[index] )
        # constraints.append(p[i] >= (math.pow(2, R_min * r_c[index] / b[i]) - 1) * noise_power / 10 ** (-loss / 10))
        # snr = tool.cal_snr_f(uavs[uav_index], users[index], p[i])
        # lower_limit = R_min * r_c[index] / math.log2(1 + p[i]*(10 ** (-loss / 10))/noise_power)
        # lower_limit=round(lower_limit, 2)
        # constraints.append(b[i] >= R_min * r_c[index] / math.log2(1 + p[i]*(10 ** (-loss / 10))/noise_power))
    problem = cp.Problem(objective, constraints)
    problem.solve(verbose=True)
    return

if __name__ == '__main__':
    users = Users.getUsers("./data4.csv")
    uavs = tool.getUavInitialPos(7, users)
    c_connection = tool.get_c_association(users, uavs, math.ceil(len(users) / len(uavs)))
    p_connection = tool.get_p_association(users, uavs, c_connection, math.ceil(len(users) * 2 / len(uavs)))
    uav_to_user=tool.get_uav_to_user(c_connection, uavs)
    print(uav_to_user)
    p_alloc ={}
    r_c = {}
    r=1
    for i in range(len(uavs)):
        for j in uav_to_user[i]:
            p_alloc[j]=0.09/4
            r_c[j]=2
            r=r+1
    print(r_c)
    p_rest = tool.get_rest_p(users, uavs, p_connection)
    # cal_b_maxmin(users, uavs, 0, uav_to_user, r_c, p_alloc, C.B_max, C.R_min)
    # cal_b(users, uavs, 0, uav_to_user, r_c, p_alloc, C.B_max, C.R_min)
    # resource_alloc_avg(users, uavs, r_c, p_connection, uav_to_user)
    resource_alloc(users, uavs, r_c, p_connection, uav_to_user)
    # resource_CVX(users,uavs,0,uav_to_user,r_c,C.B_max,p_rest[0],C.R_min,C.snr_thre)
    # B_max=10/4 #10M
    # R_min=0.2
    # p_rest = tool.get_rest_p(users, uavs, p_connection)
    # print(p_rest)
    # b_alloc=cal_b(users,uavs,0,uav_to_user,r_c,p_alloc,B_max,R_min)
    # print(b_alloc)
    # p_alloc=cal_p_by_CVX(users,uavs,0,uav_to_user,r_c,b_alloc,p_rest[0],R_min,0)
    # b_alloc = cal_b(users, uavs, 0, uav_to_user, r_c, p_alloc, B_max, R_min)
    # p_alloc = cal_p_by_CVX(users, uavs, 0, uav_to_user, r_c, b_alloc, p_rest[0], R_min, 0)


