% SEARCH_EFFICIENCY_FXN     Function for repeated execution of the search
%                           efficiency test
%
% A function that will execute the search efficiency test described in
% search_efficiency_test.m.  This function is meant to simplify the task of
% running the test repeatedly, and will return several important summary
% statistics.
%
% SYNTAX:  [MSE_fbmp_mmse, MSE_fbmp_map, MSE_ls_mmse, MSE_ls_map, ...
%           MMSE_error_fbmp, MMSE_error_ls, mass_ratio_fbmp, ...
%           mass_ratio_ls] = search_efficiency_fxn(N, M, p1, SNRdB);
%
% INPUTS:
%   N:  # of unknowns
%   M:  # of measurements
%   p1: Number of active coefficients = round(p1*N)
%   SNRdB:  Signal-to-noise ratio [dB]
%
% OUTPUTS:
%   MSE_fbmp_mmse:	MSE of FBMP's MMSE estimator
%   MSE_fbmp_map:   MSE of FBMP's MAP estimator
%   MSE_ls_mmse:    MSE of Larsson/Selen's MMSE estimator
%   MSE_ls_map:     MSE of Larsson/Selen's MAP estimator
%   MMSE_error_fbmp:    Error between true MMSE estimator and FBMP's
%   MMSE_error_ls:  Error between true MMSE estimator and Larsson/Selen's
%   mass_ratio_fbmp:    Ratio of prob. mass returned by FBMP to total
%   mass_ratio_ls:  Ratio of prob. mass returned by Larsson/Selen to total
%

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

function [MSE_fbmp_mmse, MSE_fbmp_map, MSE_ls_mmse, MSE_ls_map, ...
          MMSE_error_fbmp, MMSE_error_ls, mass_ratio_fbmp, ...
          mass_ratio_ls] = search_efficiency_fxn(N, M, p1, SNRdB);

%% Setup an experiment

% Preset parameters
vnc = 0;			% variable number of active coefs?
sig2s = [0; 1];     % sparse coefficient variances [off;on]
mus = [0; 0];       % sparse coefficient means [off;on;...;on]
D = 5;				% maximum number of RGS's
stop = 0;			% stop RGS @ this many stdevs below mean of metric

k = round(p1*N);    % Number of active coefficients

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); 
while sum(s_true > 1) < k
    s_true(ceil(rand(1)*N)) = 1 + ceil(rand(1)*(Q-1));
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
[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);
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;
end

% Normalize MMSE estimate by the scaling factor
xmmse_true = (1/scale)*xmmse_true;


%% 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
MSE_fbmp_mmse = norm(x - xmmse_fbmp)^2;
MSE_fbmp_map = norm(x - xmmse_star_fbmp{1})^2;
MSE_ls_mmse = norm(x - xmmse_ls)^2;
MSE_ls_map = norm(x - xmap_ls)^2;

% Error between true and approximated MMSE estimators
MMSE_error_fbmp = norm(xmmse_true - xmmse_fbmp)^2;
MMSE_error_ls = norm(xmmse_true - xmmse_ls)^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));
