"""This file contains utility functions relating to tests and simulations of the decoders."""

import numpy as np
import typing
from tqdm import tqdm
from timeit import default_timer

from utility import noise


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(n: int,
                 k: int,
                 decoder: typing.Any,
                 SNRs: typing.Sequence[float] = np.linspace(1, 7, 7),
                 target_frame_errors: int = 100) \
        -> typing.Tuple[np.array, np.array]:
    """Calculate the Bit Error Rate (BER) for a given decoder for a number of SNRs.

    This function assumes the all-zeros assumption holds. Progress is printed to stdout.

    :param n: Length of a codeword of the used code
    :param k: Length of a dataword of the used code
    :param decoder: Instance of the decoder to be tested
    :param SNRs: List of SNRs for which the BER should be calculated
    :param target_frame_errors: Number of frame 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 = np.zeros(n)
    x_bpsk = 1 - 2 * x  # Map x from [0, 1]^n to [-1, 1]^n

    BERs = []
    for SNR in tqdm(SNRs,
                    desc=f"Calculating BERs for {decoder.__class__.__name__}",
                    position=1,
                    leave=False,
                    bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):

        total_bit_errors = 0
        total_bits = 0
        total_frame_errors = 0

        pbar = tqdm(total=target_frame_errors,
                    desc=f"Simulating for SNR = {SNR} dB",
                    position=2,
                    leave=False,
                    bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]")

        while total_frame_errors < target_frame_errors:
            # Simulate channel
            y = noise.add_awgn(x_bpsk, SNR, n, k)

            # Decode received frame
            x_hat = decoder.decode(y)

            # Calculate statistics
            bit_errors = count_bit_errors(x, x_hat)
            total_bits += x.size

            if bit_errors > 0:
                total_frame_errors += 1
                total_bit_errors += bit_errors
                pbar.update(1)

        pbar.close()

        BERs.append(total_bit_errors / total_bits)

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


def test_decoders(n: int,
                  k: int,
                  decoders: typing.List,
                  SNRs: typing.Sequence[float] = np.linspace(1, 7, 7),
                  target_frame_errors: int = 100) \
        -> typing.Tuple[np.array, np.array]:
    """Calculate the Bit Error Rate (BER) for a number of given decoders for a number of SNRs.

    This function assumes the all-zeros assumption holds. Progress is printed to stdout.

    :param n: Length of a codeword of the used code
    :param k: Length of a dataword of the used code
    :param decoders: List of decoder objects to be tested
    :param SNRs: List of SNRs for which the BER should be calculated
    :param target_frame_errors: Number of frame errors after which to stop the simulation
    :return: Tuple of the form (SNRs, [BERs_1, BERs_2, ...]) where SNR and BERs_x are numpy arrays
    """
    result_BERs = []

    start_time = default_timer()

    for decoder in tqdm(decoders,
                        desc="Calculating the answer to life, the universe and everything",
                        position=0,
                        leave=False,
                        bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
        _, BERs = test_decoder(n, k, decoder, SNRs, target_frame_errors)
        result_BERs.append(BERs)

    end_time = default_timer()
    print(f"Elapsed time: {end_time - start_time:.2f}s")

    return SNRs, result_BERs