import numpy as np
from virtual_force import *

# 瑞利衰落
NOISE = -130 # w
SIGMA_E = 0
def calculate_snr(pt, g, n):   # 发射功率，增益，噪声
    snr = (pt * g) / n
    return snr

def get_gain_based_dist(dist2):
    pldB = pathlossA2G(dist2, 200)
    # pl = dB2powerratio(pldB)
    # ray_gain = np.mean(get_rayleigh_gain())
    return pldB

def calculate_common_sinr(ue_idx,uav,dist):   # 发射功率，增益，噪声
    pt = uav.common_power
    g = get_gain_based_dist(dist) # 大尺度和小尺度衰落, 目前是小尺度
    snr = (pt * g) / (sum(uav.power_allocation)+dBm2w(NOISE))
    return snr

def calculate_private_sinr(ue_idx,uav,dist):   # 发射功率，增益，噪声
    g = get_gain_based_dist(dist)
    pt = uav.power_allocation[ue_idx]
    snr = (pt * g) / (sum(uav.power_allocation) - pt+dBm2w(NOISE) + SIGMA_E*sum(uav.power_allocation)+(sum(uav.power_allocation)*SIGMA_E - pt*SIGMA_E))
    return snr

def get_rayleigh_gain(scale=1,size=1000):   # 瑞利 空地
    g_samples = np.random.rayleigh(scale, size)
    return g_samples

def get_rician_gain(K, scale=1, size=1000): # 莱斯 空空
    g_samples = np.sqrt((np.random.normal(scale=scale, size=size) + K)**2 + np.random.normal(scale=scale, size=size)**2)
    return g_samples

if __name__ == '__main__':

    # 假设的参数值
    pt = 1.0  # 发射功率为1瓦特
    K = 10.0 # Rician因子为10
    g_samples2 = get_rician_gain(K)

    # 生成1000个服从莱斯分布的信道增益样本
    g_samples = get_rayleigh_gain()

    # # 计算信噪比
    # snr_values = calculate_snr(pt, g_samples, n)
    # snr_values2 = calculate_snr(pt, g_samples2, n)

    # 打印前10个信噪比值
    # print(snr_values[:10],snr_values2[:10])