function [APP, alpha, beta, gamma] = MAP_AWGN( channel_output, sigma, next_state, prev_state_index, trellis_out )

% Maximum A-Posteriori Probability decoder for AWGN
%
% Primary assumtions:
%
% 1. Convlutionally coded BPSK transmission over AWGN
% 2. No parallel transition
% 3. Encoding from zero initial state
% 4. No terminal state known (uniform distribution on final beta)
% 5. Code trellis should be "normal". Specifically, the number transitions to each state should be equal to 2^N_IN
%    All the good convolutional code indeed satisfy this property. To check whether your code satisfies this property,
%    see the output variable A of siva_init() function. A = 0 if your code is normal. See test.m file for example;
%    IEEE 802.11a 64-state rate-1/2 code therein clearly satisfies this property
% 
% Function input arguments:
%
% 1. channel_output : sequence of (soft) channel output (row vector)
% 2. sigma : Noise standard deviation ( sigma = sqrt (sigma^2) )
% 3. next_state : Next state transition matrix ( can be obtained with trellis(G) function )
%              Note that I assume the numbering of state from 0.
%
% 4. prev_state_index : Previous state index matrix (can be obtained with siva_init( ) function )
%                    prev_state_index(i,:) is the set of state indexes whose next state transition is i-1.
%                    (i = 1 , ... , N_STATE );
%
% 5. trellis_out : (Bipolar) output sequence associated with prev_state_index (can be obtained with siva_init() function )
%
% Clearly the length of channel_output should be multiple of N_OUT since BPSK transmission is assumed
%
% Function output
%
% 1. APP : A posteriori probability for input messaage
%       APP(i,k) = Pr( x_k = i | Y_1, ... , Y_N ) ( k = 1 , ... , N = length( channel_out )/N_OUT )
% 2. alpha : forward recursion variable, alpha(i,k) = f( S_k = i , Y_1 , ... , Y_(k-1) );
% 3. beta : backward recursion variable, beta(i,k) = f( Y_(k+1) | s_(k+1) = j );
% 4. gamma : transition variable, 
%          gamma( (j-1)*2^(N_IN) + find( prev_state_index(j,:) == i ) , k ) = f( s_(k+1) = j, Y_k | s_k = i)
%          where (i,j) is a valid state transition
%          (Data structure of gamma is rather complex to avoid large use of memory)
%
% where f( ) denotes the probability density function of continuous (or mixed) random variable
% alpha, beta, gamma quantities are normalized such that
%
% sum_{i} alpha(i,k) = 1 for all k and so on.
%
% while these normalizations do not change any MAP algorithm derived in the class, it will improve
% the numerical precision, otherwise very large or extremely small recursion values (alpha, beta, gamma) occur 
% for large number of decoding sequences. These numerical instability comes from 
% the unboundness of "probatility density function" 
% (Note that alpha, beta, gamma are NOT pmf in AWGN).
%
% Written by Kee-Bong Song
%
% May/13/2003

[ N_STATE M ] = size( next_state );
N_IN = log2(M);
[ total_num_trans N_OUT ] = size( trellis_out ); 
N = length(channel_output)/N_OUT;

%%%% Variable initialization

%total_num_trans = size( trellis_out , 1 );

alpha = zeros(N_STATE, N);
% Initial condition
alpha(1,1) = 1;

beta = zeros(N_STATE,N);
% Final condition
beta(:,N) = 1/N_STATE;

gamma = zeros(total_num_trans, N);
% Gamma index extraction
for i = 1 : N_STATE

    temp_index = find( prev_state_index' == i );
    next_gamma_index((i-1)*2^(N_IN)+1:i*2^(N_IN)) = temp_index;
    
    for input = 1 : 2^N_IN
       n_state = next_state(i,input);
       gamma_index( i , input ) = n_state * 2^N_IN + find( prev_state_index(n_state+1,:) == i );
    end
   
end

%%%% Gamma update

for k = 1 : N;
    
    y_k = repmat( channel_output( (k-1)*N_OUT+1 : k*N_OUT ) , total_num_trans, 1 );
    delta =  (y_k - trellis_out).^2 ;
    temp = exp( - 1/( 2*sigma^2 ) * sum(delta,2) ); 
    gamma(:,k) = temp / sum(temp);
    
end

%%%% Forward alpha recursion

for k = 2 : N;
    
    prev_gamma = reshape( gamma( : , k-1 ) , 2^(N_IN) , N_STATE )';             % N_STATE by 2^(N_IN) matrix
    %prev_gamma(i,:) lists gamma values whose previous states are associated with the current state i
    
    prev_alpha = reshape( alpha( prev_state_index', k-1 ), 2^(N_IN), N_STATE)'; % N_STATE by 2^(N_IN) matrix
    %prev_alpha(i,:) lists alpha values whose previous states are associated with the current state i
    
    %Forward recursion
    temp = sum( prev_alpha .* prev_gamma , 2 );
    alpha(:,k) = temp / sum(temp);
    
end

%%%% Backward beta recursion

for k = N - 1 :-1: 1
    
    next_beta = reshape( beta( next_state' + 1, k+1 ), 2^(N_IN), N_STATE)'; %N_STATE by 2^(N_IN) matrix
    %next_beta(j,:) lists beta values whose next states are associated with the current state j
    
    next_gamma = reshape( gamma( next_gamma_index ,k+1), 2^(N_IN) , N_STATE )' ; %N_STATE by 2^(N_IN) matrix
    %next_beta(j,:) lists gamma values whose next states are associated with the current state j
    
    temp = sum( next_beta .* next_gamma , 2 );
    beta(:,k) = temp / sum(temp); 
    
end

%%%% Decoding

% APP calculation with alpha, beta, gamma quantities

for k = 1 : N
    
    for input = 1 : 2^N_IN
        
        APP(input,k) = sum( alpha(:,k) .* beta(next_state(:,input)+1,k) .* gamma( gamma_index(:,input) ,k ) );
        
    end

end

%Normalize for probability value
APP = APP./repmat( sum(APP), 2^N_IN, 1);