import torch

def add_frequency_noise(fft, snr=10, vacant_snr=15, mask=None):
    ### 根据SNR确定noise的放大比例
    num_pixels = fft.numel()

    fft_magnitude = torch.abs(fft)
    fft_phase = torch.angle(fft)

    # fft_magnitude
    mag_psr = torch.mean(torch.abs(fft_magnitude) ** 2)
    mag_pnr = mag_psr / (10 ** (snr / 10))  # Calculate noise power
    noise_mag = torch.randn_like(fft_magnitude) * torch.sqrt(mag_pnr)

    mag_psr_vacant = mag_psr / (10 ** (vacant_snr / 10))
    noise_mag_vacant = torch.randn_like(fft_magnitude) * torch.sqrt(mag_psr_vacant)

    fft_magnitude = fft_magnitude + \
                    noise_mag * fft_magnitude * mask + \
                    noise_mag_vacant * (1- mask)
    fft_magnitude = torch.abs(fft_magnitude)

    # fft_phase
    pha_psr = torch.mean(torch.abs(fft_phase) ** 2)
    pha_pnr = pha_psr / (10 ** (snr / 10))  # Calculate noise power for phase
    noise_pha = torch.randn_like(fft_phase) * torch.sqrt(pha_pnr)

    pha_psr_vacant = pha_psr / (10 ** (vacant_snr / 10))
    noise_pha_vacant = torch.randn_like(fft_phase) * torch.sqrt(pha_psr_vacant)

    fft_phase = fft_phase + \
                noise_pha * fft_phase * mask + \
                noise_pha_vacant * (1- mask)

    noise_fft = fft_magnitude * torch.exp(1j * fft_phase)

    return noise_fft