% EXP_DECAY_TEST    This code will evaluate the performance of multiple 
% sparse reconstruction algorithms in solving a problem involving a 
% random-sign, random-permutation exponentially decaying signal, for 
% varying values of rho, the decay parameter.  FBMP will be tested using a 
% refinement strategy, given coarse initial parameter estimates.  To choose 
% between FBMP's GEM-based and surrogate hyperparameter updates, please see
% compare_fxn_refine.m
%
% * Note that the following sparse reconstruction algorithms must be
% on MATLAB's search path for full functionality: SparseBayes, StOMP, OMP,
% GPSR, BCS, and VB-BCS *
%

%
% Coded by: Justin Ziniel, Ohio State Univ.
% Email: ziniel.1@osu.edu
% Last modified: April 2, 2009
% Change summary: Added VB-BCS to comparisons (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


%% Initialize static variables

T = 2000;              	% # of trials for each value of lambda
M = 128;                % # observations
N = 512;                % # parameters
SNRdB = 15;             % SNR (in dB)
D = 5;                  % Max # of RGS iterations allowed
stop = 0;               % Stop RGS @ this many stdevs below mean of metric
E_max = 20;             % Maximum number of refinement stages
est = {'p1', 'sig2w', 'both'};     % Parameters to refine (not used by GEM)


% Coarse initial parameter estimates for FBMP
p1_est = 0.10;
sig2w_est = 0.05;
sig2s_est = [0; 0.5];
mus_est = [0; 0];

rho = 0.10 : 0.05 : 0.85;        % Range of decay values to evaluate
Nexp = length(rho)*T;            % Total # of experiments to run

Algs = 7;                           % # of algorithms under test

% Create a cell array that will house the MSEs, and runtimes for each trial 
% of each algorithm.
NMSE = cell(1, Algs);
RUNTIME = cell(1, Algs);
SPARSE = cell(1, Algs);


%% Begin testing

path(path, '../functions');     % Add path for FBMP solvers

time = clock;
for ind = 1:length(rho)
    % Conduct T trials w/ a decay of rho
    rho_i = rho(ind);
    
    for t = 1:T   

        % Generate an experiment
        % Generation of test signal
        x = exp(-rho_i*[1:N]');     % Unknown signal
        [dummy, order] = sort(rand(N,1));
        [dummy, order2] = sort(rand(N,1));
        signs = [ones(N/2,1); -ones(N/2,1)];
        x = signs(order2).*x(order);             % Random permutation & sign
        clear dummy order order2 signs

        % Formation of observations
        A = randn(M,N);                     % mixing matrix
        A = A*diag(1./sqrt(diag(A'*A)));    % ...with unit-norm columns
        sig2w = ((norm(x)^2)/M)*10^(-SNRdB/10);    % noise variance
        w = sqrt(sig2w)*randn(M,1);         % AWGN samples
        y = A*x + w;                        % observation
        

        % Obtain the solutions
        [NMSEs Times Sparse trial] = compare_fxn_refine(y, A, sig2s_est, ...
                                                        p1_est, mus_est, ...
                                                        sig2w_est, D, x, ...
                                                        sig2w, E_max, ...
                                                        est);
                                 
        % Distribute results to appropriate cells
        nmse_names = fieldnames(NMSEs);
        time_names = fieldnames(Times);
        sparse_names = fieldnames(Sparse);
        for c = 1:Algs
            nmse = getfield(NMSEs, nmse_names{c});
            count = getfield(Sparse, sparse_names{c});
            runtime = getfield(Times, time_names{c});
            
            if c == 1   % FBMP algorithm
                NMSE{c}(t,ind,1) = nmse(1);     % NMSE for MMSE est.
                NMSE{c}(t,ind,2) = nmse(2);     % NMSE for MAP est.
                SPARSE{c}(t,ind,1) = count(1);  % Sparseness for MMSE est.
                SPARSE{c}(t,ind,2) = count(2);  % Sparseness for MAP est.
            else
                NMSE{c}(t,ind) = nmse;
                SPARSE{c}(t,ind) = count;
            end

            RUNTIME{c}(t,ind) = runtime;
        end
        FBMP.d_tot(t,ind) = trial.d_tot;    % Avg. number of RGS iterations
        FBMP.e_tot(t,ind) = trial.e_tot;    % Number of refinement stages
        FBMP.p1_diff(t,ind) = trial.p1_diff;    % Difference in p1 estimate
        FBMP.sig2w_diff(t,ind) = trial.sig2w_diff;  % " "
        FBMP.sig2s_diff(t,ind) = trial.sig2s_diff;  % " "
        FBMP.one_shot(t,ind) = trial.one_shot;  % Single iteration runtime
        
        
        % Provide an update of the estimated time to completion
        if (mod(t, 5) == 0)
            est_time = (Nexp - ((T*(ind - 1)) + t))*...
                        (etime(clock, time)/((T*(ind - 1)) + t))/60/60;
            disp(sprintf('Estimated time remaining: %3.2f hours', est_time));
        end
    end
end
clear A NMSEs Times Sparse trial


%% Plotting results

marker = {'bo-' 'gv-' 'r.-' 'c*-' 'md-' 'ys-' 'kx-'};

figure(1); clf; axis off
figure(2); clf; axis off
figure(3); clf; axis off

for i = 1:Algs
    figure(1);
    if i == 1   % FBMP algorithm
        h1 = plot(rho, 10*log10(mean([NMSE{i}(:,:,1)],1)), marker{i}); hold on
        h1 = plot(rho, 10*log10(mean([NMSE{i}(:,:,2)],1)), 'b^-'); hold on
    else
        h1 = plot(rho, 10*log10(mean([NMSE{i}],1)), marker{i}); hold on
    end
    figure(2);
    h2 = semilogy(rho, mean([RUNTIME{i}],1), marker{i}); hold on
    if i == 1
        semilogy(rho, mean(FBMP.one_shot,1), 'b^-');
    end
    figure(3);
    if i == 1   % FBMP algorithm
        h3 = plot(rho, mean([SPARSE{i}(:,:,1)],1), marker{i}); hold on
        h3 = plot(rho, mean([SPARSE{i}(:,:,2)],1), 'b^-'); hold on
    elseif i ~= 1 && i ~=7
        h3 = plot(rho, mean([SPARSE{i}],1), marker{i}); hold on
    end
end
figure(1); hold off;
figure(2); hold off;
figure(3); hold off;

figure(1)
xlabel('\rho'); ylabel('NMSE [dB]')
ttl_str = ['N = ' num2str(N) ', M = ' num2str(M) ', SNR = ' num2str(SNRdB) ...
            ' dB, D_{max} = ' num2str(D) ', E_{max} = ' num2str(E_max) ...
            ', T = ' num2str(T)];
title(ttl_str)
legend('FBMP_{mmse} (w/ EM update)', 'FBMP_{map} (w/ EM update)', ...
        'SparseBayes', 'OMP', 'StOMP', 'GPSR', 'BCS', 'VB-BCS', ...
        'Location', 'best');
    
figure(2)
xlabel('\rho'); ylabel('Runtime [s]')
title(ttl_str)
legend('FBMP (w/ EM update)', 'FBMP (w/o EM update)', ...
        'SparseBayes', 'OMP', 'StOMP', 'GPSR', 'BCS', 'VB-BCS', ...
        'Location', 'best');

    
figure(3)
xlabel('\rho'); ylabel('||x_{recovery}||_0')
title(ttl_str)
legend('FBMP_{mmse} (w/ EM update)', 'FBMP_{map} (w/ EM update)', ...
        'SparseBayes', 'OMP', 'StOMP', 'GPSR', 'BCS', ...
        'Location', 'best');
