"""Utility functions relating to noise and SNR calculations."""


import numpy as np


def get_noise_amp_from_SNR(SNR: float, signal_amp: float = 1) -> float:
    """Calculate the amplitude of the noise from an SNR and the signal amplitude.

    :param SNR: Signal-to-Noise-Ratio in dB
    :param signal_amp: Signal Amplitude (linear)
    :return: Noise Amplitude (linear)
    """
    SNR_linear = 10 ** (SNR / 10)
    noise_amp = (1 / np.sqrt(SNR_linear)) * signal_amp

    return noise_amp


def add_awgn(c: np.array, SNR: float, signal_amp: float = 1) -> np.array:
    """Add Additive White Gaussian Noise to a data vector. As this function adds random noise to
    the input, the output changes, even if it is called multiple times with the same input.

    :param c: Binary vector representing the data to be transmitted
    :param SNR: Signal-to-Noise-Ratio in dB
    :param signal_amp: Amplitude of the signal. Used for the noise amplitude calculation
    :return: Data vector with added noise
    """
    noise_amp = get_noise_amp_from_SNR(SNR, signal_amp=signal_amp)
    y = c + np.random.normal(scale=noise_amp, size=c.size)
    return y