# Type your code here 
import numpy as np
import functools
from scipy.special import erfc
from multiprocessing import Process, Manager

class Parameters:

    k = 8
    N = 16
    G = np.array([[1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0],
        [1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0],
        [1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])

def encode(message, Parameters) -> np.array:
    '''
    Method to encode the message using the generator matrix G
    '''
    G = Parameters.G
    message = np.array(message)
    encoded_message = np.dot(message, G) % 2
    return encoded_message

def bpsk_modulation(bit_sequence: np.ndarray) -> np.ndarray:
    '''
    Method that receives a bit sequence as input and returns the bpsk modulation
    '''
    return 2*bit_sequence-1

def add_noise(signal: np.ndarray, EbOverN0: float) -> np.ndarray:
    '''
    Method that receives a signal and adds noise to it
    '''
    noise_std = np.sqrt(1 / (2 * EbOverN0))

    # Add AWGN
    noise = noise_std * np.random.randn(len(signal))
    received_signal = signal + noise

    return received_signal

def polar_map_decode(received_signal, Parameters, codebook) -> np.array:
    '''
    Method to decode the received signal using the polar map decoding algorithm. Signal is uniformly distributed so we can use ML
    '''
    G = Parameters.G

    distances = []

    for u in codebook:

        u = np.array([int(i) for i in u])

        # Encode u
        encoded_u = encode(u, Parameters)

        # BPSK modulation
        modulated_u = bpsk_modulation(encoded_u)

        # Compute distance
        distances.append(np.sum((received_signal - modulated_u)**2))

    decoded_message = codebook[np.argmin(distances)]

    # Convert to numpy array
    return np.array([int(i) for i in decoded_message])

def bpsk_theoretical_BER(Eb_N0_linear: np.ndarray) -> np.ndarray:
    '''
    Theoretical BER for BPSK over AWGN
    '''
    return 1/2 * erfc(np.sqrt(Eb_N0_linear))


def monte_carlo_simulation(EbOverN0: np.ndarray, num_frames: int, Parameters, BEP, codebook) -> None:

    BER = []

    EbOverN0Linear = 10**(EbOverN0/10)

    block_error = 0

    
    for frame in range(num_frames):

        # Generate random message
        message = np.random.randint(2, size=Parameters.k)

        # Encode message
        encoded_message = encode(message, Parameters)

        # BPSK modulation
        modulated_message = bpsk_modulation(encoded_message)

        # Add noise
        received_signal = add_noise(modulated_message, EbOverN0Linear*(Parameters.k/Parameters.N))

        # Decode received signal
        decoded_message = polar_map_decode(received_signal, Parameters, codebook)
        
        # Compute BER
        BER.append(np.sum(np.abs(message - decoded_message)))

        block_error += int(np.sum(np.abs(message - decoded_message)) > 0)

    

    BEP[float(EbOverN0)] = np.sum(BER) / (num_frames * Parameters.k)

    block_error_rate = block_error / num_frames

    print(f'Eb/N0: {EbOverN0} dB, Bit Error Probability: {BEP[EbOverN0]}, Block Error Rate: {block_error_rate}')


def main():

    codebook = []

    for i in range(2**8):
        codebook.append(bin(i)[2:].zfill(8))


    #Define EbOverN0 range

    num_eboverno = 2

    EbOverN0dB = np.linspace(8, 9, num_eboverno)

    num_frames = 20000

    workers = []

    manager = Manager()

    BEP = manager.dict()

    for EbOverN0 in EbOverN0dB:

        print(f'Starting simulation for Eb/N0: {EbOverN0} dB')

        p = Process(target=monte_carlo_simulation, args=(EbOverN0, num_frames, Parameters, BEP, codebook))
        p.start()
        workers.append(p)

    for worker in workers:
        worker.join()

    # Theoretical BER
    theoretical_BER = bpsk_theoretical_BER(10**(EbOverN0dB/10))

    # Convert dict to np array
    #BEP = np.array([BEP[float(eb)] for eb in EbOverN0dB])

    # Plot
    import matplotlib.pyplot as plt

    #EbOverN0dB = np.linspace(-10, 5, 16)
    """
    BEP = np.array([0.4575875, 0.4447875, 0.4297125, 0.4108125, 
                    0.38385, 0.36095, 0.323275, 0.2825, 
                    0.2329, 0.179025, 0.1251, 0.0793125, 
                    0.0419, 0.0197625, 0.006425, 0.0018625])"""

    plt.figure()
    plt.semilogy(EbOverN0dB, BEP, 'o-', label='ML Simulated')

    plt.semilogy(EbOverN0dB, theoretical_BER, 'o-', label='BPSK Theoretical')

    plt.xlabel('Eb/N0 (dB)')
    plt.ylabel('BER')

    plt.legend()
    plt.grid()
    plt.show()


if __name__ == '__main__':
    main()