% TEST_FBMPR     Demonstrates how to construct a signal according to the
% signal model, and recover the signal using the FBMP algorithm by calling 
% the function FBMPR_FXN
%

%
% Coded by: Phil Schniter and Justin Ziniel, The Ohio State Univ.
% Email: schniter@ece.osu.edu, ziniel.1@osu.edu
% Last modified: June 30, 2009
% FBMP version 1.2 
% Copyright (c) Philip Schniter, Lee C. Potter, and Justin Ziniel, 2008
%

clear; 
%randn('state',1);rand('twister',0); 
%randn('state',state_randn);
%rand('twister',state_rand);


%% Setup an experiment

% User-defined parameters
M = 128;           	% # observations
N = 512;   			% # parameters
p1 = 0.04;			% prior probability of non-zero tap
vnc = 0;			% variable number of active coefs?
SNRdB = 15;			% SNR (in dB)
D = 20;				% maximum number of allowable repeated greedy searches
stop = 0;			% stop RGS @ this many stdevs below mean of metric

% *************************************************************************
% SOME EXAMPLE SIGNAL MODELS (un-comment the desired model)

% Zero-mean binary model
sig2s = [0; 1];     % sparse coefficient variances [off;on]
mus = [0; 0];       % sparse coefficient means [off;on;...;on]

% % Non-zero-mean binary model
% sig2s = [0; 1];     % sparse coefficient variances [off;on]
% mus = [0; 1];       % sparse coefficient means [off;on;...;on]

% % Non-zero-mean ternary model
% sig2s = [0; 0.05];     % sparse coefficient variances [off;on]
% mus = [0; 1; -1];      % sparse coefficient means [off;on;...;on]

% *************************************************************************


%% Use user-provided signal model and other variables to conduct an
%% experiment

% generate true mixing vector 
Q = length(mus)-1;
ps = [1 - p1; ones(Q,1)*p1/Q];      % All active Gaussians are equiprobable
intps = cumsum([0;ps]);
sig2s_tmp = [sig2s; sig2s(2)*ones(Q-1,1)];
s_true = ones(N,1); 
if vnc ~= 0         % if variable number of active coefs...
    v = rand(N,1);
    for i = 2:length(ps)
        s_true = s_true + (v > intps(i));
    end
else                % if fixed number of active coefs...
    while sum(s_true > 1) < round(p1*N)
        s_true(ceil(rand(1)*N)) = 1 + ceil(rand(1)*Q);
    end
end

% generate an observation
A = randn(M,N);                     % random normal mixing matrix
A = A*diag(1./sqrt(diag(A'*A)));    % ...with unit-norm columns
x = mus(s_true) + sqrt(sig2s_tmp(s_true)).*randn(N,1);  % sparse coefficients
sig2w = norm(A*x)^2/M*10^(-SNRdB/10);   % Noise variance from SNR and signal realization
w = sqrt(sig2w)*randn(M,1);         % AWGN samples
y = A*x + w;                        % observation


%% Run the RGS algorithm to recover MMSE estimate of x

path(path, '../functions/');        % Add FBMP functions to current path

tic
% Call the FBMPR function
[xmmse, xmmse_star, psy_star, nu_star, T_star, d_tot, d_max] = ...
    fbmpr_fxn_reduced(y, A, p1, sig2w, sig2s, mus, D, stop);
toc

NMSE = (norm(xmmse - x)^2)/(norm(x)^2)


%% Plotting results

ttl = ['||x||_0 = ' int2str(sum(x ~= 0)) ', NMSE = ' num2str(NMSE) ', N = ' int2str(N) ...
            ', M = ' int2str(M) ', p_1 = ' num2str(p1) ', P_{signal} = ' num2str(norm(A*x)^2) ...
            ', P_{noise} = ' num2str(norm(w)^2)];
subplot(211)
handy = stem([x, xmmse, xmmse_star{1}]);
set(handy(1),'Marker','o');
set(handy(2),'Marker','+');
set(handy(3),'Marker','x');
legend(handy,'x_{true}', 'x_{MMSE}', 'x_{MAP}', 'Location', 'Best')
title(ttl); ylabel('x[n]'); xlabel('n')
subplot(212)
plot(psy_star);
title('Posterior probabilities of the dominant mixture vectors')
xlabel('Top mixture vectors'); ylabel('p(s|y)')
