"""This file contains various utility functions that can be used in combination with the decoders.
"""

import numpy as np
import typing
from tqdm import tqdm


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


def count_bit_errors(d: np.array, d_hat: np.array) -> int:
    """Count the number of wrong bits in a decoded codeword.

    :param d: Originally sent data
    :param d_hat: Received data
    :return: Number of bit errors
    """
    return np.sum(d != d_hat)


def test_decoder(decoder: typing.Any,
                 c: np.array,
                 SNRs: typing.Sequence[float] = np.linspace(1, 4, 7),
                 target_bit_errors=100,
                 N_max=10000) \
        -> typing.Tuple[np.array, np.array]:
    """Calculate the Bit Error Rate (BER) for a given decoder for a number of SNRs.

    This function prints its progress to stdout.

    :param decoder: Instance of the decoder to be tested
    :param c: Codeword whose transmission is to be simulated (element of [0, 1]^n)
    :param SNRs: List of SNRs for which the BER should be calculated
    :param target_bit_errors: Number of bit errors after which to stop the simulation
    :param N_max: Maximum number of iterations to perform for each SNR
    :return: Tuple of numpy arrays of the form (SNRs, BERs)
    """
    x = c * 2 - 1  # Map the codeword from [0, 1]^n to [-1, 1]^n
    BERs = []
    for SNR in tqdm(SNRs, desc="Calculating Bit-Error-Rates",
                    position=0,
                    bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):

        total_bit_errors = 0
        total_bits = 0

        for n in tqdm(range(N_max), desc=f"Simulating for SNR = {SNR} dB",
                      position=1,
                      leave=False,
                      bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):

            # TODO: Is this a valid simulation? Can we just add AWGN to the codeword, ignoring and modulation and (
            #  e.g. matched) filtering?
            y = add_awgn(x, SNR, signal_amp=(1 / np.sqrt(2)))
            y_hat = decoder.decode(y)

            total_bit_errors += count_bit_errors(c, y_hat)
            total_bits += c.size

            if total_bit_errors >= target_bit_errors:
                break

        BERs.append(total_bit_errors / total_bits)

    return np.array(SNRs), np.array(BERs)