% SEARCH_EFFICIENCY_TEST    A test of the efficiency of the FBMP tree 
% search in recovering all "highly probable" mixture selection vectors.  
% This program will perform an exhaustive search of all possible mixture 
% selection vectors and determine their posterior likelihood.  It will 
% compare this list to the likelihood of candidate vectors returned by FBMP
% in S_star, as well as the tree-search advocated by Larsson and Selen.
%

%
% Coded by: Justin Ziniel, Ohio State Univ.
% Email: ziniel.1@osu.edu
% Last modified: April 2, 2009
% Change summary: Added Larsson/Selen algorithm to the comparison (4/2/09)
% FBMP version 1.2 
% Copyright (c) Philip Schniter, Lee C. Potter, and Justin Ziniel, 2008
%

clear; clc
randn('state', 0);      % Return random normal generator to default state
rand('twister', 5489);  % Return random generator to default state


%% Setup an experiment

% User-defined parameters
N = 26;   			% # parameters
M = ceil(N/4);      % # observations
p1 = 0.04;			% prior probability of non-zero tap
k = round(p1*N);
vnc = 0;			% variable number of active coefs?
SNRdB = 10;			% SNR [dB]
sig2s = [0; 1];	% sparse coefficient variances [off;on]
mus = [0; 0];	% sparse coefficient means [off;on;...;on]
D = 20;				% maximum number of parallel searches
stop = 0;			% stop RGS @ this many stdevs below mean of metric

Q = length(mus);
sig2w = ((mean(mus(2:Q).^2)+sig2s(2))*p1*N)/M*10^(-SNRdB/10);
sig2s_tmp = [sig2s; sig2s(2)*ones(Q-2,1)];

% generate true mixing vector 
s_true = ones(N,1); 
if vnc == 1         % 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-1));
    end
end

% generate an observation
A = randn(M,N);                     % mixing matrix
A = A*diag(1./sqrt(diag(A'*A)));    % ...with unit-norm columns
w = sqrt(sig2w)*randn(M,1);         % AWGN samples
x = mus(s_true) + sqrt(sig2s_tmp(s_true)).*randn(N,1);  % sparse coefficients
y = A*x + w;                        % observation


% %% Run the RGS algorithm

path(path, '../functions/')
tic
[xmmse_fbmp, xmmse_star_fbmp, psy_star, nu_star, T_star, d_tot, ...
    Cov_star, d_max] = fbmpr_fxn(y, A, p1, sig2w, sig2s, mus, D, stop);
toc
clear Cov_star

% Find the decimal representation of the mixture vectors captured in
% T_star, and keep track of them
FBMP_S_star = [];
for i = 1:length(T_star)
    FBMP_S_star = [FBMP_S_star, sum(2.^(T_star{i} - 1))];
end


%% Perform an exhaustive search and acquire probabilities for all mixture
%% vectors.  Also keep a recursively updated MMSE estimate

nu_complete = NaN*ones(1,2^N);      % Allocate space for the exhaustive
                                    % list of all mixture selection metrics
scale = 0;                          % Scaling constant, (sum exp{nu(s',y)})
xmmse_true = zeros(N,1);

time = clock;
for i = 0:(2^N - 1)
    s = flipud(str2num(dec2bin(i, N)'));	% Mixture vector under test
    s = s + ones(N,1);

    % Compute posterior for mixture vector under test
    Phi = A*diag(sig2s(s))*A' + sig2w*eye(M);
    CL = chol(Phi, 'lower');   % Cholesky decomposition
    logdet = sum(log(diag(CL)) + log(conj(diag(CL))));   
    y_c = y - A*mus(s);    % Centered y-vector
    logprob = sum(s == 1)*log(1 - p1) + sum(s > 1)*log(p1/(Q - 1));
    nu = -(1/2)*(y_c'/CL')*(CL\y_c) - (1/2)*logdet - ...
                (M/2)*log(2*pi) + logprob;     % Real-valued AWGN true metric

    % Add newly computed metric to the list, update scale
    nu_complete(i+1) = nu;
    scale = scale + exp(nu);

    % Update the MMSE estimate
    xmmse_upd = mus(s) + diag(sig2s(s))*(A'/CL')*(CL\y_c);
    xmmse_true = xmmse_true + exp(nu)*xmmse_upd;

    % Provide an update of the estimated time to completion
    if (mod(i, 10000) == 0)
        est_time = (2^N - (i + 1))*(etime(clock, time)/(i + 1))/60/60;
        disp(sprintf(['Estimated time remaining: %3.2f hours. '...
                      ' |  Printed: ' datestr(clock, 21)], est_time));
    end
end

% Normalize MMSE estimate by the scaling factor
xmmse_true = (1/scale)*xmmse_true;
   
disp(sprintf('Total runtime: %3.2f hours', etime(clock, time)/60/60));


%% Run the Larsson/Selen tree search

% Begin building the high-probability set of mixture vector decimal
% representations, including the all-active hypothesis by default
LS_S_star = [2^N - 1];

% Start with all coefficients active
s = ones(N,1) + ones(N,1);

% Continue de-activating taps until none are left active
while sum(s == 2) ~= 0
    
    best_metric = -inf;     % Best possible de-activation candidate
    best_code = [];       	% Decimal "code" of best candidate
    index = [];             % Index to de-activate for given depth
    
    for j = 1:N     % Cycle through, de-activating taps
        s_new = s;
        if s_new(j) == 2    % Feasible deactivation found
            s_new(j) = 1;               % Deactivate
            s_new = s_new - ones(N,1);  % Convert to binary
            % Compute the unique decimal representation of the mixture
            % vector under test
            code = bin2dec(num2str(flipud(s_new))');
            % Offset the code to index into the existing exhaustive list of
            % metrics, and if it is the best, update to reflect this
            if nu_complete(code + 1) > best_metric
                best_metric = nu_complete(code + 1);
                best_code = code;
                index = j;
            end
        end
    end
    
    % Entire depth has been explored.  Update current s-vector and add the
    % best code to the high-probability set list
    s(index) = 1;
    LS_S_star = [LS_S_star, best_code];
end

% Larsson/Selen search is complete.  Now construct the MAP and MMSE
% estimates
nu_star_ls = nu_complete(LS_S_star + 1);        % Larsson/Selen metrics
psy_star_ls = exp(nu_star_ls)/sum(exp(nu_star_ls));     % Normalized posteriors
xmmse_ls = zeros(N,1);
for j = 1:(N + 1)
    s = flipud(str2num(dec2bin(LS_S_star(j), N)'));	% Re-build mixture vector
    s = s + ones(N,1);

    % Compute quantities needed for conditional estimates
    Phi = A*diag(sig2s(s))*A' + sig2w*eye(M);
    CL = chol(Phi, 'lower');   % Cholesky decomposition
    y_c = y - A*mus(s);    % Centered y-vector

    % Update the LS MMSE estimate
    xmmse_upd = mus(s) + diag(sig2s(s))*(A'/CL')*(CL\y_c);
    xmmse_ls = xmmse_ls + psy_star_ls(j)*xmmse_upd;
    
    % If MAP, store solution
    if psy_star_ls(j) == max(psy_star_ls)
        xmap_ls = xmmse_upd;
    end
end


%% Calculate important performance metrics

% Normalized mean-squared errors
NMSE_fbmp_mmse = norm(x - xmmse_fbmp)^2/norm(x)^2;
NMSE_fbmp_map = norm(x - xmmse_star_fbmp{1})^2/norm(x)^2;
NMSE_ls_mmse = norm(x - xmmse_ls)^2/norm(x)^2;
NMSE_ls_map = norm(x - xmap_ls)^2/norm(x)^2;

% Error between true and approximated MMSE estimators
MMSE_error_fbmp = norm(xmmse_true - xmmse_fbmp)^2/norm(xmmse_true)^2;
MMSE_error_ls = norm(xmmse_true - xmmse_ls)^2/norm(xmmse_true)^2;

% Ratio of probability mass captured by FBMP and LS
mass_ratio_fbmp = sum(exp(nu_complete(FBMP_S_star + 1)))/sum(exp(nu_complete));
mass_ratio_ls = sum(exp(nu_complete(LS_S_star + 1)))/sum(exp(nu_complete));


%% Plotting

exp_length = length(nu_star);       % Number of explored mixture vectors
unexp_length = length(best_missed); % Number of unexplored mixture vectors

nu_max = max([best_missed, nu_star]);

psy_missed = exp(best_missed - nu_max)/sum(exp([best_missed, nu_star] - nu_max));
psy_found = exp(nu_star - nu_max)/sum(exp([best_missed, nu_star] - nu_max));

% Generate figure
handle = semilogy(1:unexp_length, psy_missed, 'r', 'LineWidth', 2); hold on
semilogy(unexp_length + 1:unexp_length + exp_length, fliplr(psy_found), 'b', 'LineWidth', 2); hold off
a_handle = gca;
xtick = get(a_handle, 'XTick');
set(a_handle, 'XTick', linspace((2^N - max(xtick)), 2^N, length(xtick)))
a = axis;
l = line([unexp_length + 0.5; unexp_length + 0.5], a(3:4));
set(l, 'Color', 'black', 'LineStyle', '--', 'Linewidth', 1);
gtext({'Top', 'Unexplored', 'Mixture', 'Vectors'}, 'Color', 'black', ...
        'HorizontalAlignment', 'center', 'FontSize', 18)
gtext({'Top', 'Explored', 'Mixture', 'Vectors'}, 'Color', 'black', ...
        'HorizontalAlignment', 'center', 'FontSize', 18)
xlabel('Rank-ordered mixture vector indices'); ylabel('p(s|y)')
ttl_str = ['N = ' num2str(N) ', M = ' num2str(M) ', SNR = ' num2str(SNRdB) ...
            ' dB, \lambda_1 = ' num2str(p1) ', \sigma_1^2 = '...
            num2str(sig2s(2)) ', \mu_1 = ' num2str(mus(2)) ', D_{max} = ' ...
            num2str(D)];
title(ttl_str)
