import numpy as np
from numpy.random import randint
import sk_dsp_comm.fec_conv as fec
import sk_dsp_comm.digitalcom as dc
import matplotlib.pyplot as plt
N_bits_per_frame = 10000
EbN0 = 4
total_bit_errors = 0
total_bit_count = 0
cc1 = fec.FECConv(('11101','10011'),25)
state = '0000'
while total_bit_errors < 100:
    x = randint(0,2,N_bits_per_frame)
    y,state = cc1.conv_encoder(x,state)
    yn_soft = dc.cpx_awgn(2*y-1,EbN0-3,1) # Channel SNR is 3 dB less for rate 1/2
    yn_hard = ((np.sign(yn_soft.real)+1)/2).astype(int)
    z = cc1.viterbi_decoder(yn_hard,'hard')
    bit_count, bit_errors = dc.bit_errors(x,z)
    total_bit_errors += bit_errors
    total_bit_count += bit_count
    print('Bits Received = %d, Bit errors = %d, BEP = %1.2e' %                    (total_bit_count, total_bit_errors,                    total_bit_errors/total_bit_count))
print('*****************************************************')
print('Bits Received = %d, Bit errors = %d, BEP = %1.2e' %                (total_bit_count, total_bit_errors,                total_bit_errors/total_bit_count))
# Expected:
## Rate 1/2 Object
## kmax =  0, taumax = 0
## Bits Received = 9976, Bit errors = 77, BEP = 7.72e-03
## kmax =  0, taumax = 0
## Bits Received = 19952, Bit errors = 175, BEP = 8.77e-03
## *****************************************************
## Bits Received = 19952, Bit errors = 175, BEP = 8.77e-03
#
cc1.traceback_plot()
plt.show()
#
SNRdB = np.arange(0,12,.1)
Pb_uc = fec.conv_Pb_bound(1/3,7,[4, 12, 20, 72, 225],SNRdB,2)
Pb_s_third_3 = fec.conv_Pb_bound(1/3,8,[3, 0, 15],SNRdB,1)
Pb_s_third_4 = fec.conv_Pb_bound(1/3,10,[6, 0, 6, 0],SNRdB,1)
Pb_s_third_5 = fec.conv_Pb_bound(1/3,12,[12, 0, 12, 0, 56],SNRdB,1)
Pb_s_third_6 = fec.conv_Pb_bound(1/3,13,[1, 8, 26, 20, 19, 62],SNRdB,1)
Pb_s_third_7 = fec.conv_Pb_bound(1/3,14,[1, 0, 20, 0, 53, 0, 184],SNRdB,1)
Pb_s_third_8 = fec.conv_Pb_bound(1/3,16,[1, 0, 24, 0, 113, 0, 287, 0],SNRdB,1)
Pb_s_half = fec.conv_Pb_bound(1/2,7,[4, 12, 20, 72, 225],SNRdB,1)
plt.figure(figsize=(5,5))
plt.semilogy(SNRdB,Pb_uc)
plt.semilogy(SNRdB,Pb_s_third_3,'--')
plt.semilogy(SNRdB,Pb_s_third_4,'--')
plt.semilogy(SNRdB,Pb_s_third_5,'g')
plt.semilogy(SNRdB,Pb_s_third_6,'--')
plt.semilogy(SNRdB,Pb_s_third_7,'--')
plt.semilogy(SNRdB,Pb_s_third_8,'--')
plt.semilogy([0,1,2,3,4,5],[9.08e-02,2.73e-02,6.52e-03,                                8.94e-04,8.54e-05,5e-6],'gs')
plt.axis([0,12,1e-7,1e0])
plt.title(r'Soft Decision Rate 1/2 Coding Measurements')
plt.xlabel(r'$E_b/N_0$ (dB)')
plt.ylabel(r'Symbol Error Probability')
plt.legend(('Uncoded BPSK','R=1/3, K=3, Soft',                    'R=1/3, K=4, Soft','R=1/3, K=5, Soft',                    'R=1/3, K=6, Soft','R=1/3, K=7, Soft',                    'R=1/3, K=8, Soft','R=1/3, K=5, Sim',                     'Simulation'),loc='upper right')
plt.grid();
plt.show()
#
N_bits_per_frame = 10000
EbN0 = 3
total_bit_errors = 0
total_bit_count = 0
cc2 = fec.FECConv(('11111','11011','10101'),25)
state = '0000'
while total_bit_errors < 100:
    x = randint(0,2,N_bits_per_frame)
    y,state = cc2.conv_encoder(x,state)
    yn_soft = dc.cpx_awgn(2*y-1,EbN0-10*np.log10(3),1) # Channel SNR is 10*log10(3) dB less
    yn_hard = ((np.sign(yn_soft.real)+1)/2).astype(int)
    z = cc2.viterbi_decoder(yn_hard.real,'hard')
    bit_count, bit_errors = dc.bit_errors(x,z)
    total_bit_errors += bit_errors
    total_bit_count += bit_count
    print('Bits Received = %d, Bit errors = %d, BEP = %1.2e' %                    (total_bit_count, total_bit_errors,                    total_bit_errors/total_bit_count))
print('*****************************************************')
print('Bits Received = %d, Bit errors = %d, BEP = %1.2e' %                (total_bit_count, total_bit_errors,                total_bit_errors/total_bit_count))
# Expected:
## Rate 1/3 Object
## kmax =  0, taumax = 0
## Bits Received = 9976, Bit errors = 251, BEP = 2.52e-02
## *****************************************************
## Bits Received = 9976, Bit errors = 251, BEP = 2.52e-02
#
SNRdB = np.arange(0,12,.1)
Pb_uc = fec.conv_Pb_bound(1/3,7,[4, 12, 20, 72, 225],SNRdB,2)
Pb_s_third_3_hard = fec.conv_Pb_bound(1/3,8,[3, 0, 15, 0, 58, 0, 201, 0],SNRdB,0)
Pb_s_third_5_hard = fec.conv_Pb_bound(1/3,12,[12, 0, 12, 0, 56, 0, 320, 0],SNRdB,0)
Pb_s_third_7_hard = fec.conv_Pb_bound(1/3,14,[1, 0, 20, 0, 53, 0, 184],SNRdB,0)
Pb_s_third_5_hard_sim = np.array([8.94e-04,1.11e-04,8.73e-06])
plt.figure(figsize=(5,5))
plt.semilogy(SNRdB,Pb_uc)
plt.semilogy(SNRdB,Pb_s_third_3_hard,'r--')
plt.semilogy(SNRdB,Pb_s_third_5_hard,'g--')
plt.semilogy(SNRdB,Pb_s_third_7_hard,'k--')
plt.semilogy(np.array([5,6,7]),Pb_s_third_5_hard_sim,'sg')
plt.axis([0,12,1e-7,1e0])
plt.title(r'Hard Decision Rate 1/3 Coding Measurements')
plt.xlabel(r'$E_b/N_0$ (dB)')
plt.ylabel(r'Symbol Error Probability')
plt.legend(('Uncoded BPSK','R=1/3, K=3, Hard',                    'R=1/3, K=5, Hard', 'R=1/3, K=7, Hard',                    ),loc='upper right')
plt.grid();
plt.show()
#
cc2.traceback_plot()