import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import norm


def generate_noise(signal, R, EbN0):
    """Generate the noise array for a given signal and its SNR

    The array of randomly generated values is multiplied with the
    noise variance which is computed using the SNR value. The array
    thus obtained is the noise array.

    Args:
        signal (array): signal array
        R (int): symbol rate
        EbN0 (float): SNR value

    Returns:
        noise (array): The noise array
    """
    ns = len(signal)
    # Noise = Noise variance * array of random normal values
    noise = (1 / np.sqrt(2 * R * EbN0)) * (
        np.random.normal(size=ns) + np.random.normal(size=ns) * 1j
    )
    return noise


def generate_QPSK_BER_vs_EbN0_Pr(EbN0dB_range):
    """Generate practical BER for a QPSK signal

    The randomly generated message is first converted to its corresponding
    symbol array. Since this is QPSK, the symbol rate is 2. The signal array
    is generated by appending to it the corresponding complex number i.e.
    e^(j*theta) and the theta is different for each of the symbol pattern.
    For wireless channel, we multiply the signal array with the h array which
    is an array of complex numbers belonging to the normal distribution. For
    AWGN channel, signal array is kept as it is. Noise is added to the signal
    and then the receiver decodes the signal by checking the sign of both the
    real and imaginary part to obtain the received message array.

    Args:
        EbN0dB_range (range): range object containing the range of SNR
                              values used in simulations

    Returns:
        BER_AWGN (list): list of BER values of QPSK over AWGN channel
        BER_wireless (list): list of BER values of QPSK over Rayleigh channel
    """
    msg = np.random.randint(low=0, high=2, size=int(1e6))
    symbols = np.array([msg[0::2], msg[1::2]])  # Split message into symbols of two bits
    signal = []
    # QPSK modulation block
    for k in range(np.size(symbols, axis=1)):
        b_0 = symbols[0, k]
        b_1 = symbols[1, k]
        if b_0 == 0 and b_1 == 0:
            theta = 5.0 * np.math.pi / 4.0
            signal.append(np.array(np.cos(theta) + 1j * np.sin(theta)))
        elif b_0 == 0 and b_1 == 1:
            theta = 3.0 * np.math.pi / 4.0
            signal.append(np.array(np.cos(theta) + 1j * np.sin(theta)))
        elif b_0 == 1 and b_1 == 1:
            theta = 1.0 * np.math.pi / 4.0
            signal.append(np.array(np.cos(theta) + 1j * np.sin(theta)))
        elif b_0 == 1 and b_1 == 0:
            theta = 7.0 * np.math.pi / 4.0
            signal.append(np.array(np.cos(theta) + 1j * np.sin(theta)))

    BER_wireless = []
    BER_AWGN = []
    for EbN0dB in EbN0dB_range:
        EbN0 = 10 ** (EbN0dB / 10.0)
        noise = generate_noise(signal, 1, EbN0)
        # Generate RV h
        h = (1 / np.sqrt(2)) * (
            np.random.normal(size=len(signal)) + 1j * np.random.normal(size=len(signal))
        )
        # y is for wireless/Rayleigh channel
        y = h * signal + noise
        y /= h
        # y_2 is for AWGN channel
        y_2 = signal + noise
        wirelessly_received_msg = []
        received_msg = []
        # Decoder for wireless/Rayleigh channel
        for i in y:
            wirelessly_received_msg.append(np.real(i) > 0)
            wirelessly_received_msg.append(np.imag(i) > 0)
        # Decoder for AWGN channel
        for i in y_2:
            received_msg.append(np.real(i) > 0)
            received_msg.append(np.imag(i) > 0)
        Pb_pr_wireless = np.count_nonzero(msg != wirelessly_received_msg) / len(msg)
        Pb_pr = np.count_nonzero(msg != received_msg) / len(msg)
        BER_AWGN.append(Pb_pr)
        BER_wireless.append(Pb_pr_wireless)

    return BER_AWGN, BER_wireless


if __name__ == "__main__":
    EbN0dB_range = range(0, 21)
    BER_AWGN, BER_wireless = generate_QPSK_BER_vs_EbN0_Pr(EbN0dB_range)
    plt.plot(EbN0dB_range, BER_AWGN, "s-", label="QPSK over AWGN channel - Simulations")
    plt.plot(
        EbN0dB_range,
        BER_wireless,
        "D-",
        label="QPSK over Rayleigh channel - Simulations",
    )
    plt.xlabel("SNR (dB)")
    plt.xscale("linear")
    plt.ylabel("BER")
    plt.yscale("log")
    plt.legend()
    plt.grid()
    plt.show()
