#include "fp_block_decoder.h"

FP_block_decoder::FP_block_decoder()
{
}

void FP_block_decoder::block_decoder(double &apriori_uncoded, vec apriori_alphas, vec apriori_betas, double &apriori_encoded, double &extrinsic_uncoded, vec &extrinsic_betas, vec &extrinsic_alphas)
{
    /*cout << "Debug: checking input" << endl;
    cout << "apriori_alphas = " << apriori_alphas << endl;
    cout << "apriori_betas = " << apriori_betas << endl;
    cout << "extrinsic_alphas = " << extrinsic_alphas << endl;
    cout << "extrinsic_betas = " << extrinsic_betas << endl;*/

    // Set Log-MAP approximation method
    max_star.initialize(1);

    //% Matrix to describe the trellis
    //% Each row describes one transition in the trellis
    //% Each state is allocated an index 1,2,3,... Note that this list starts
    //% from 1 rather than 0.
    //%               FromState,        ToState,        UncodedBit,     EncodedBit
    /*mat transitions = "1                1               0               0;
                       2                5               0               0";
                       3                6               0               1;
                       4                2               0               1;
                       5                3               0               1;
                       6                7               0               1;
                       7                8               0               0;
                       8                4               0               0;
                       1                5               1               1;
                       2                1               1               1;
                       3                2               1               0;
                       4                6               1               0;
                       5                7               1               0;
                       6                3               1               0;
                       7                4               1               1;
                       8                8               1               1";
                       */
    mat transitions =  "1 1 0 0; 2 5 0 0; 3 6 0 1; 4 2 0 1; 5 3 0 1; 6 7 0 1; 7 8 0 0; 8 4 0 0; 1 5 1 1; 2 1 1 1; 3 2 1 0; 4 6 1 0; 5 7 1 0; 6 3 1 0; 7 4 1 1; 8 8 1 1";

    // Find the largest state index in the transitions matrix
    state_count = std::max(itpp::max(transitions.get_col(0)),itpp::max(transitions.get_col(1)));
    
    //% Calculate a posteriori transition log-probabilities
    vec deltas = zeros(transitions.rows());
    for (int i=0; i<transitions.rows(); i++)
    {
        deltas(i) = apriori_alphas(transitions(i,0)-1) + apriori_betas(transitions(i,1)-1); //% 16 additions
        deltas(i) = deltas(i) + transitions(i,2)*apriori_uncoded + transitions(i,3)*apriori_encoded;
    };
//cout << "Debug: deltas = " << deltas << endl;

    /*//% Calculate a posteriori transition log-probabilities
    deltas = apriori_alphas(transitions(:,1)) + apriori_betas(transitions(:,2)); //% 16 additions
    deltas(transitions(:,3)==0 & transitions(:,4)==1) = deltas(transitions(:,3)==0 & transitions(:,4)==1) + apriori_encoded; //% 4 additions
    deltas(transitions(:,3)==1 & transitions(:,4)==0) = deltas(transitions(:,3)==1 & transitions(:,4)==0) + apriori_uncoded; //% 4 additions
    deltas(transitions(:,3)==1 & transitions(:,4)==1) = deltas(transitions(:,3)==1 & transitions(:,4)==1) + (apriori_uncoded + apriori_encoded); //% 5 additions
    */

    //% Calculate the uncoded extrinsic LLR
    double log_p0 = max_star.approximation(deltas.left(8));
    double log_p1 = max_star.approximation(deltas.right(8));
    extrinsic_uncoded = log_p1 - log_p0 - apriori_uncoded;
 //cout << "Debug: extrinsic_uncoded = " << extrinsic_uncoded << endl;

        /*% Calculate the uncoded extrinsic LLR
        log_p0 = maxstar(deltas(transitions(:,3) == 0)); % 7 maxstars
        log_p1 = maxstar(deltas(transitions(:,3) == 1)); % 7 maxstars
        extrinsic_uncoded = log_p1 - log_p0 - apriori_uncoded; % 2 additions
        */

    //% Calculate the extrinsic backward state log-probabilities
    vec vec_log_p1 = zeros(apriori_alphas.length());
    for (int state_index = 0; state_index<state_count; state_index++)
    {
            vec temp1(2);
            int temp1_index=0;
            for (int i=0; i<transitions.rows(); i++)
                if (transitions(i,0) == state_index+1)
                {
                    temp1(temp1_index) = deltas(i);
                    temp1_index++;
                }
            //cout << "temp1_index=" << temp1_index << " temp1 = "<< temp1 << endl;
            vec_log_p1(state_index) = max_star.approximation(temp1.left(temp1_index));
    };
 //cout << "Debug: vec_log_p1 = " << vec_log_p1 << endl;
 //cout << "Debug: apriori_alphas = " << apriori_alphas << endl;
    extrinsic_betas = vec_log_p1 - apriori_alphas;
// cout << "Debug: extrinsic_betas = " << extrinsic_betas << endl;
    
        /*% Calculate the extrinsic backward state log-probabilities
        log_p1 = zeros(size(apriori_alphas));
        for state_index = 1:state_count
            log_p1(state_index) = maxstar(deltas(transitions(:,1) == state_index)); % 1 maxstar per iteration of for loop = 8 maxstars
        end
        extrinsic_betas = log_p1 - apriori_alphas; % 8 additions
        */

    //% Calculate the extrinsic forward state log-probabilities
    vec vec_log_p0 = zeros(apriori_betas.length());
    for (int state_index = 0; state_index<state_count; state_index++)
    {
            vec temp2(transitions.rows());
            int temp2_index=0;
            for (int i=0; i<transitions.rows(); i++)
                if (transitions(i,1) == state_index+1)
                {
                    temp2(temp2_index) = deltas(i);
                    temp2_index++;
                }
            vec_log_p0(state_index) = max_star.approximation(temp2.left(temp2_index));
    };
    extrinsic_alphas = vec_log_p0 - apriori_betas;
//cout << "Debug: extrinsic_alphas = " << extrinsic_alphas << endl;

        /*% Calculate the extrinsic forward state log-probabilities
        log_p1 = zeros(size(apriori_betas));
        for state_index = 1:state_count
            log_p1(state_index) = maxstar(deltas(transitions(:,2) == state_index)); % 1 maxstar per iteration of for loop = 8 maxstars
        end
        extrinsic_alphas = log_p1 - apriori_betas; % 8 additions
        */

}

void FP_block_decoder::block_decoder(double &apriori_uncoded, vec apriori_alphas, vec apriori_betas, double &apriori_encoded, double &extrinsic_uncoded, double &extrinsic_encoded, vec &extrinsic_betas, vec &extrinsic_alphas)
{
    /*cout << "Debug: checking input" << endl;
    cout << "apriori_alphas = " << apriori_alphas << endl;
    cout << "apriori_betas = " << apriori_betas << endl;
    cout << "extrinsic_alphas = " << extrinsic_alphas << endl;
    cout << "extrinsic_betas = " << extrinsic_betas << endl;*/

    // Set Log-MAP approximation method
    max_star.initialize(1);

    //% Matrix to describe the trellis
    //% Each row describes one transition in the trellis
    //% Each state is allocated an index 1,2,3,... Note that this list starts
    //% from 1 rather than 0.
    //%               FromState,        ToState,        UncodedBit,     EncodedBit
    /*mat transitions = "1                1               0               0;
                       2                5               0               0";
                       3                6               0               1;
                       4                2               0               1;
                       5                3               0               1;
                       6                7               0               1;
                       7                8               0               0;
                       8                4               0               0;
                       1                5               1               1;
                       2                1               1               1;
                       3                2               1               0;
                       4                6               1               0;
                       5                7               1               0;
                       6                3               1               0;
                       7                4               1               1;
                       8                8               1               1";
                       */
    mat transitions =  "1 1 0 0; 2 5 0 0; 3 6 0 1; 4 2 0 1; 5 3 0 1; 6 7 0 1; 7 8 0 0; 8 4 0 0; 1 5 1 1; 2 1 1 1; 3 2 1 0; 4 6 1 0; 5 7 1 0; 6 3 1 0; 7 4 1 1; 8 8 1 1";

    // Find the largest state index in the transitions matrix
    state_count = std::max(itpp::max(transitions.get_col(0)),itpp::max(transitions.get_col(1)));

    //% Calculate a posteriori transition log-probabilities
    vec deltas = zeros(transitions.rows());
    for (int i=0; i<transitions.rows(); i++)
    {
        deltas(i) = apriori_alphas(transitions(i,0)-1) + apriori_betas(transitions(i,1)-1); //% 16 additions
        deltas(i) = deltas(i) + transitions(i,2)*apriori_uncoded + transitions(i,3)*apriori_encoded;
    };


    // Calculate the uncoded extrinsic LLR
    double log_p0 = max_star.approximation(deltas.left(8));
    double log_p1 = max_star.approximation(deltas.right(8));
    extrinsic_uncoded = log_p1 - log_p0 - apriori_uncoded;

    // Calculate the encoded extrinsic LLR
    vec encoded_log_p0 = zeros(1);
    vec encoded_log_p1 = zeros(1);
    for (int i=0; i<transitions.rows(); i++)
    {
        if (transitions(i,3)==0)
            encoded_log_p0.ins(0,deltas(i));
        else if (transitions(i,3)==1)
            encoded_log_p1.ins(0,deltas(i));
    }
    encoded_log_p0.del(encoded_log_p0.length()-1);
    encoded_log_p1.del(encoded_log_p1.length()-1);
    log_p0 = max_star.approximation(encoded_log_p0);
    log_p1 = max_star.approximation(encoded_log_p1);
    extrinsic_encoded = log_p1 - log_p0 - apriori_encoded;

    //% Calculate the extrinsic backward state log-probabilities
    vec vec_log_p1 = zeros(apriori_alphas.length());
    for (int state_index = 0; state_index<state_count; state_index++)
    {
            vec temp1(2);
            int temp1_index=0;
            for (int i=0; i<transitions.rows(); i++)
                if (transitions(i,0) == state_index+1)
                {
                    temp1(temp1_index) = deltas(i);
                    temp1_index++;
                }
            //cout << "temp1_index=" << temp1_index << " temp1 = "<< temp1 << endl;
            vec_log_p1(state_index) = max_star.approximation(temp1.left(temp1_index));
    };
    extrinsic_betas = vec_log_p1 - apriori_alphas;


    //% Calculate the extrinsic forward state log-probabilities
    vec vec_log_p0 = zeros(apriori_betas.length());
    for (int state_index = 0; state_index<state_count; state_index++)
    {
            vec temp2(transitions.rows());
            int temp2_index=0;
            for (int i=0; i<transitions.rows(); i++)
                if (transitions(i,1) == state_index+1)
                {
                    temp2(temp2_index) = deltas(i);
                    temp2_index++;
                }
            vec_log_p0(state_index) = max_star.approximation(temp2.left(temp2_index));
    };
    extrinsic_alphas = vec_log_p0 - apriori_betas;


}
