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);