//#define USE_PRECODER
#include "itpp/itcomm.h"
#include "robprob.h"
#include "fp_turbo_equalization.h"
#include "convolutional_encoder.h"
#include "fp_block_decoder.h"
#include "eq_interleaver.h"

using namespace RobProb;
using namespace itpp;
using std::cout;
using std::cin;
using std::endl;
using std::string;

int main(void)
{

/*    // Test !!!!!!!!!!!!

    EQ_Interleaver eq_interleaver(18,true);
    eq_interleaver.generate_interleaver();
    cout << eq_interleaver.get_interleaver() << endl;

    vec test_a = "1:1:6";
    vec test_b = "11:1:19";
    vec test_c = "21:1:29";
    vec test_in = concat(test_a,test_b,test_c);

    //vec interleaved_d = eq_interleaver.interleave(test_a, test_b, test_c, true);

    cout << test_a << endl;
    cout << test_b << endl;
    cout << test_c << endl;

    vec interleaved_d = eq_interleaver.interleave(test_in, true);
    cout << interleaved_d << endl;

    vec test_out = eq_interleaver.deinterleave(interleaved_d);
    cout << test_out << endl;
    cin.get();*/


    //general parameters
    double threshold_value = 50;
    string map_metric = "maxlogMAP";
    ivec gen = "07 05";//octal notation
    int constraint_length = 3;
    int ch_nb_taps = 4;//number of channel multipaths
    int nb_errors_lim = 1e5;
    int nb_bits_lim = int(1e7);
    int perm_len = pow2i(14);//permutation length
    int nb_iter = 64; //number of iterations in the turbo equalizer
    int iteration_counts = 1; //number of iterations in the turbo decoder
    vec EbN0_dB = "3:0.5:8";
    double R = 1.0 / 3.0;//coding rate of FEC
    double Ec = 1.0;//coded bit energy

    //other parameters
    int nb_bits_tail = perm_len / gen.length();
    //int nb_bits = nb_bits_tail - (constraint_length - 1);//number of bits in a block (without tail)
    int nb_bits = 1024;
    vec sigma2 = (0.5 * Ec / R) * pow(inv_dB(EbN0_dB), -1.0);//N0/2
    int nb_blocks;//number of blocks
    int nb_errors;
    bvec bits(nb_bits);//data bits
    bvec nsc_coded_bits(perm_len);//tail is added
    bvec em_bits(perm_len);
    bmat parity_bits;
    ivec perm(perm_len);
    ivec inv_perm(perm_len);
    vec rec(perm_len);

    //SISO equalizer
    vec eq_apriori_data;
    vec eq_extrinsic_data;
    vec extrinsic_coded;

    //SISO NSC
    vec nsc_intrinsic_coded(perm_len);
    vec nsc_apriori_data(nb_bits_tail);
    nsc_apriori_data.zeros();//always zero
    vec nsc_extrinsic_coded;
    vec nsc_extrinsic_data;


    //decision
    bvec rec_bits(nb_bits_tail);
    int snr_len = EbN0_dB.length();
    mat ber(nb_iter, snr_len);
    ber.zeros();
    register int en, n;

    // EQ interleaver;
    int frame_length = (int)nb_bits*3;
    EQ_Interleaver eq_interleaver(frame_length, true);

    //CCs
    Convolutional_Code nsc;
    nsc.set_generator_polynomials(gen, constraint_length);



    //BPSK
    BPSK bpsk;

    //AWGN
    AWGN_Channel awgn;

    //multipath channel impulse response (Rayleigh fading) with real coefficients
    vec ch_imp_response(ch_nb_taps);
    vec ini_state = ones(ch_nb_taps);//initial state is zero
    MA_Filter<double, double, double> multipath_channel;

    //SISO blocks
    MY_SISO siso;
    siso.set_generators(gen, constraint_length);
    siso.set_map_metric(map_metric);
    bool siso_equalization_tail = false;

    //BER
    BERC berc;

    Convolutional_Encoder encoder;
cout << "Debug: encoder is sucessful!" << endl; //---debug
    ivec interleaver_sequence = encoder.get_interleaver_sequence(nb_bits);
cout << "Debug: getting interleaver sequence is sucessful!" << endl;

    FP_block_decoder fp_block_decoder;
cout << "Debug: decoder is sucessful!" << endl; //---debug

    Sequence_Interleaver<bin> interleaver(nb_bits);
    interleaver.set_interleaver_sequence(interleaver_sequence);
cout << "Debug: bit_interleaver is sucessful!" << endl; //---debug

    Sequence_Interleaver<double> LLR_interleaver(nb_bits);
    LLR_interleaver.set_interleaver_sequence(interleaver_sequence);
cout << "Debug: LLR_interleaver is sucessful!" << endl; //---debug

    Sequence_Interleaver<double> LLR2_interleaver(nb_bits);
    LLR2_interleaver.set_interleaver_sequence(interleaver_sequence);

    //Randomize generators
    RNG_randomize();

    //main loop
    for (en = 0;en < snr_len;en++) {
        cout << "EbN0_dB = " << EbN0_dB(en) << endl;
        awgn.set_noise(sigma2(en));
        siso.set_noise(sigma2(en));
        nb_errors = 0;
        nb_blocks = 0;
        while ((nb_errors < nb_errors_lim) && (nb_blocks*nb_bits < nb_bits_lim))//if at the last iteration the nb. of errors is inferior to lim, then process another block
        {
            cout << "Frame #" << nb_blocks << endl;
            // Generate a random frame of bits
            bvec a = randb(nb_bits);

            //BPSK modulate the systematic bits
            vec a_tx = itpp::to_vec(a)*2-1;

            // Interleave the random frame of bits
            bvec b = interleaver.interleave(a);

            //% Generate the parity bits
            vec c,e;
            encoder.encode(to_vec(a),c,e);
            vec d,f;
            encoder.encode(to_vec(b),d,f);

            //Interleaving
            //vec encoded_frame = itpp::concat(to_vec(a),c,d);
            //Sequence_Interleaver<double> frame_interleaver(encoded_frame.length());
            //frame_interleaver.randomize_interleaver_sequence();
            //vec interleaved_frame = frame_interleaver.interleave(encoded_frame);
            eq_interleaver.generate_interleaver();
            vec interleaved_frame = eq_interleaver.interleave(to_vec(a), c, d, true);

            // BPSK modulation (1->-1,0->+1)
            vec modulated_signal = bpsk.modulate_bits(to_bvec(interleaved_frame));
            //cout << modulated_signal.left(20) << endl;
            //cin.get();

            // Multipath channel + AWGN noise
            ch_imp_response = randray(ch_nb_taps);
            //ch_imp_response.zeros();
            //ch_imp_response(0) = 1;
            ch_imp_response /= sqrt(sum_sqr(ch_imp_response));//normalized power profile
            multipath_channel.set_coeffs(ch_imp_response);
            multipath_channel.set_state(ini_state);//inital state is zero

            rec = awgn(multipath_channel(modulated_signal));
            //rec = awgn((modulated_signal));

            // Initialize turbo equalizer
            eq_apriori_data.set_length(rec.length());
            eq_apriori_data.zeros();//a priori information of emitted symbols
            eq_extrinsic_data.set_length(rec.length());
            eq_extrinsic_data.zeros();
            extrinsic_coded.set_length(rec.length());
            extrinsic_coded.zeros();
            siso.set_impulse_response(ch_imp_response);
            siso.gen_chtrellis();

            // Inialise extrinsic alphas and betas of fully parallel equalizer
            int stateNb = siso.chtrellis.stateNb;
            mat extrinsic_betas=zeros(stateNb,rec.length()+1); //% We know that these are not the final state
            if (siso_equalization_tail)
            {
                extrinsic_betas.set_col(rec.length(),ones(stateNb)*(-INFINITY));
                extrinsic_betas(0,rec.length())=0;
            }
            mat extrinsic_alphas=zeros(stateNb,rec.length()+1); //% We know that these are not the final state
            extrinsic_alphas.set_col(0,ones(stateNb)*(-INFINITY));
            extrinsic_alphas(0,0)=0;

            mat apriori_alphas = extrinsic_alphas;
            mat apriori_betas = extrinsic_betas;
            mat temp_extrinsic_betas = extrinsic_betas;
            mat temp_extrinsic_alphas = extrinsic_alphas;


            //% Inialise extrinsic alphas and betas of fully parallel turbo decoder
            mat upper_extrinsic_betas=zeros(8,c.length()+1); //% We know that these are not the final state
            upper_extrinsic_betas.set_col(c.length(),ones(8)*(-100.0));
            upper_extrinsic_betas(0,c.length())=0;
            mat upper_extrinsic_alphas=zeros(8,c.length()+1); //% We know that these are not the final state
            upper_extrinsic_alphas.set_col(0, ones(8)*(-100.0));
            upper_extrinsic_alphas(0,0)=0;

            mat lower_extrinsic_betas=zeros(8,c.length()+1); //% We know that these are not the final state
            lower_extrinsic_betas.set_col(c.length(),ones(8)*(-100.0));
            lower_extrinsic_betas(0,c.length())=0;
            mat lower_extrinsic_alphas=zeros(8,c.length()+1); //% We know that these are not the final state
            lower_extrinsic_alphas.set_col(0, ones(8)*(-100.0));
            lower_extrinsic_alphas(0,0)=0;

            //nsc_intrinsic_coded = frame_interleaver.deinterleave(rec);
            nsc_intrinsic_coded = eq_interleaver.deinterleave(rec);


            for (n = 0;n < nb_iter;n++)
            {
                //cout << "Frame #" << nb_blocks << " - TE iter. #" << n << endl;

                //% Update a priori alphas and betas
                apriori_alphas = extrinsic_alphas;
                apriori_betas = extrinsic_betas;
                temp_extrinsic_betas = extrinsic_betas;
                temp_extrinsic_alphas = extrinsic_alphas;

                //parallel equalization
                int start_bit=0;
                if (itpp::mod(n,2)==1)
                    start_bit = 1;

                for (int bit_index=start_bit; bit_index<rec.length()-1; bit_index +=2)
                //for (int bit_index=0; bit_index<rec.length(); bit_index +=1)
                {
                    vec vec_temp_extrinsic_alphas = temp_extrinsic_alphas.get_col(bit_index);
                    vec vec_temp_extrinsic_betas =  temp_extrinsic_betas.get_col(bit_index+1);
                    double temp_eq_extrinsic_data;

                    // Max-Log-MAP Equalization
                    siso.fp_equalizer_maxlogMAP(rec(bit_index), apriori_alphas.get_col(bit_index), apriori_betas.get_col(bit_index+1), eq_apriori_data(bit_index),
                                                temp_eq_extrinsic_data, vec_temp_extrinsic_alphas, vec_temp_extrinsic_betas);

                    temp_extrinsic_alphas.set_col(bit_index+1, vec_temp_extrinsic_alphas);
                    temp_extrinsic_betas.set_col(bit_index, vec_temp_extrinsic_betas);
                    eq_extrinsic_data(bit_index) = temp_eq_extrinsic_data;
                }

                extrinsic_alphas = temp_extrinsic_alphas;
                extrinsic_betas = temp_extrinsic_betas;


                //eq_extrinsic_data = rec+eq_apriori_data;


                vec a_c = nsc_intrinsic_coded.left(a.length());
                vec c_c = nsc_intrinsic_coded.mid(a.length(),c.length());
                vec e_c = c_c.right(e.length());
                vec d_c = nsc_intrinsic_coded.mid(a.length()+c.length(),d.length());
                vec f_c = d_c.right(f.length());

                bvec decoded_a = bpsk.demodulate_bits((-a_c));

                //count errors
                berc.clear();
                berc.count(a, decoded_a);

                if (berc.get_errorrate()==0) break;

                //% Initialise a priori and extrinsic LLRs to zero
                vec a_e = zeros(a_c.length());
                vec a_a = zeros(a_c.length());
                vec y_e = zeros(c_c.length());
                vec z_e = zeros(d_c.length());

                vec p_e = zeros(c_c.length());
                vec q_e = zeros(d_c.length());

                //% Interleave the systematic LLRs
                vec b_c = LLR_interleaver.interleave(a_c);

                //% Perform iterative decoding
                for (int iteration_index = 0; iteration_index < iteration_counts; iteration_index++)
                {
                    //cout << "Frame #" << nb_blocks << " - TE iter. #" << n << " - TD iter. #" << iteration_index << endl;

                    //% Update a priori alphas and betas
                    mat upper_apriori_alphas = upper_extrinsic_alphas;
                    mat upper_apriori_betas = upper_extrinsic_betas;
                    mat lower_apriori_alphas = lower_extrinsic_alphas;
                    mat lower_apriori_betas = lower_extrinsic_betas;

                    //% Interleave the extrinsic LLRs
                    vec b_a = LLR_interleaver.interleave(a_e);

                    //% Obtain the uncoded a priori input for upper decoder
                    vec y_a = concat(a_a + a_c, e_c);

                    //% Obtain the uncoded a priori input for lower decoder
                    vec z_a = concat(b_a + b_c, f_c);

                    //% Operate the upper block decoders having odd indices and the lower block decoders having even indices in parallel

                    //parallel equalization
                    //for (int bit_index=start_bit; bit_index<y_a.length()-1; bit_index+=2)
                    for (int bit_index=0; bit_index<y_a.length(); bit_index +=1)
                    {
                        /*if (start_bit==0)
                            bit_index += 1;
                        else
                            bit_index -= 1;*/

                            vec temp_upper_extrinsic_betas = zeros(upper_extrinsic_betas.rows());
                            vec temp_upper_extrinsic_alphas = zeros(upper_extrinsic_alphas.rows());
                            fp_block_decoder.block_decoder(y_a(bit_index), upper_apriori_alphas.get_col(bit_index), upper_apriori_betas.get_col(bit_index+1), c_c(bit_index), y_e(bit_index), p_e(bit_index), temp_upper_extrinsic_betas, temp_upper_extrinsic_alphas);
                            upper_extrinsic_betas.set_col(bit_index, temp_upper_extrinsic_betas);
                            upper_extrinsic_alphas.set_col(bit_index+1, temp_upper_extrinsic_alphas);

                            /*if (start_bit==0)
                                bit_index -= 1;
                            else
                                bit_index += 1;*/

                            vec temp_lower_extrinsic_betas = zeros(lower_extrinsic_betas.rows());
                            vec temp_lower_extrinsic_alphas = zeros(lower_extrinsic_alphas.rows());
                            fp_block_decoder.block_decoder(z_a(bit_index), lower_apriori_alphas.get_col(bit_index), lower_apriori_betas.get_col(bit_index+1), d_c(bit_index), z_e(bit_index), q_e(bit_index), temp_lower_extrinsic_betas, temp_lower_extrinsic_alphas);
                            lower_extrinsic_betas.set_col(bit_index, temp_lower_extrinsic_betas);
                            lower_extrinsic_alphas.set_col(bit_index+1, temp_lower_extrinsic_alphas);


                    }

                    //% Remove the LLRs corresponding to the termination bits
                    a_e = y_e.left(a.length());

                    //% Remove the LLRs corresponding to the termination bits
                    vec b_e = z_e.left(b.length());

                    //% Deinterleave the extrinsic LLRs
                    a_a = LLR2_interleaver.deinterleave(b_e);


                    //*******************************************************************************

                    //% Obtain the a posteriori LLRs
                    vec a_p = a_a + a_c + a_e;

                    //Hard-decoding systematic bits
                    decoded_a =  bpsk.demodulate_bits(-a_p);

                    //count errors
                    berc.clear();
                    berc.count(a, decoded_a);

                    //cout << berc.get_errorrate() << endl;
                    if (berc.get_errorrate()==0) break;
                }

                //cin.get();
                //cout << berc.get_errorrate() << endl;
                if (berc.get_errorrate()==0)    break;

                //interleave
                extrinsic_coded = itpp::concat(a_a + a_e, p_e, q_e);
                eq_apriori_data = eq_interleaver.interleave(extrinsic_coded, true);

                //deinterleave equalized frame
                nsc_intrinsic_coded = eq_interleaver.deinterleave(eq_extrinsic_data);

                ber(n, en) += berc.get_errorrate();

            }//end iterations

            nb_errors += int(berc.get_errors());//get number of errors at the last iteration
            nb_blocks++;
        }//end blocks (while loop)

        //compute BER over all tx blocks
        ber.set_col(en, ber.get_col(en) / nb_blocks);
    }

    //save results to file
    it_file ff("fp_turbo_equalizer_bersim_multipath.it");
    ff << Name("BER") << ber;
    ff << Name("EbN0_dB") << EbN0_dB;
    ff.close();

    cout << EbN0_dB << endl;
    cout << ber<< endl;

    return 0;
}
