% FBMPC_REFINE_FXN   Function for executing the complex-valued Fast Bayesian 
% Matching Pursuit algorithm using a repeated greedy search (RGS), accompanied 
% by an iterative refinement of hyperparameter estimates using the surrogate 
% equations instead of the GEM equations in the hyperparameter update step
%
% This function allows one to implement the algorithm described in "Fast 
% Bayesian Matching Pursuit: Model Uncertainty and Parameter Estimation 
% for Sparse Linear Models," P. Schniter, L. C. Potter, and J. Ziniel.  
% FBMPC_REFINE_FXN will find a high-probability set of sparse solutions "x"
% to the underdetermined system of linear equations
%                           y = Ax + w,
% where both x and w are complex Gaussian.
%
% The refinement stages will progressively refine the given parameters using 
% the outputs of the FBMPC_REDUCED_FXN function in an attempt to
% converge to parameter choices that minimize NMSE.  This is useful when
% the parameters are not known with high accuracy.
%
% SYNTAX: [xmmse, xmmse_star, psy_star, nu_star, T_star, d_tot, d_max, ...
%            e_tot, est_diff, one_shot] = fbmpc_refine_fxn(y, A, p1, ...
%            	sig2w, sig2s, mus, D, stop, E, r_stop, est, update)
%
% INPUTS:
%
% y -       Vector of observations of length M
%
% A -       M-by-N measurement matrix 
%
% p1 -      Initial estimate of the prior probability of active taps 
%           (i.e. Pr{s_n = q} = p1/(Q-1), q > 0)
%
% sig2w -   Initial estimate of noise variance.  The authors find that
%           better performance is obtained when the noise is over-estimated
%
% sig2s -   2-by-1 vector of initially estimated sparse coefficient 
%           variances, ([off; on]).  Note that the variances of all Q-1 
%           active Gaussian mixture densities will be equal
%
% mus -     Q-by-1 vector of initially estimated means for the Q-ary 
%           Gaussian mixture densities, ([off; on_1; on_2; ...; on_Q-1]).  
%           Note that the first element of the vector should be zero, to 
%           model sparsity
%
% D -       Maximum number of repeated greedy searches (default 5)
%
% stop -    Threshold for the number of standard deviations below
%           E[nu(s|y)] at which FBMP may terminate upon finding a metric
%           that exceeds it, (default 0)
%
% E -       Maximum number of refinement stages (default 10)
%
% r_stop -  Threshold for early termination of parameter refinement, a
%           number between 0 and 1.  Once the percentage change in
%           parameter estimates for all parameters is below r_stop,
%           FBMP_REFINE_FXN will terminate, (default 0.02)
%
% est -     A list of parameters to re-estimate at each refinement stage,
%           expressed as a cell array of strings, (default {'p1', 'sig2w',
%           'sig2s'}).  User can estimate p1, sig2w, sig2s, and mus.  If 
%           user wishes to estimate both sig2s and mus, pass the single 
%           string 'both' instead of individual strings 'sig2s' and 'mus',
%           for faster execution
%
% update -  String specifying whether the MAP estimate or MMSE estimate 
%           will be used to refine parameter estimates, (default 'mmse')
%
%
% OUTPUTS:
%
% xmmse -       Approximate sparse MMSE estimate of x, given y
%
% xmmse_star -	Significant conditional MMSE estimates of x, given s and y
%
% psy_star -    Corresponding significant conditional mixture vector
%               posterior probabilities, p(s|y)
%
% nu_star -     Corresponding metric values, nu(s,y)
%
% T_star -      Indices of active taps for the conditional MMSE estimates 
%               of x returned in xmmse_star
%
% d_tot -       Average number of required RGS iterations to find a mixture
%               vector that exceeded E[nu(s,y)] - stop*stddev(nu(s,y)).  If
%               the returned value of d_tot is equal to D, then it is
%               likely that no mixture vectors ever exceeded the threshold
%
% d_max -       RGS iteration at which the maximum metric was found in 
%               the last iteration of the FBMP algorithm, (1 <= d_max <= D)
%
% e_tot -       Number of refinement iterations required, (1 <= e_tot <= E)
%
% est_diff -    A structure that contains the percentage difference between
%               starting and ending estimates of parameters
%
% one_shot -    Elapsed runtime for the final iteration of FBMP

%
% Coded by: Justin Ziniel, The Ohio State Univ.
% FBMP algorithm by: Philip Schniter, The Ohio State Univ.
% E-mail: ziniel.1@osu.edu, schniter@ece.osu.edu
% Last change: June 30, 2009
% Change summary: Made minor changes to update this function to FBMPv1.2 
%		function naming conventions.
% FBMP version 1.2 
% Copyright (c) Philip Schniter, Lee C. Potter, and Justin Ziniel, 2008
%

function [xmmse, xmmse_star, psy_star, nu_star, T_star, d_tot, d_max, ...
            e_tot, est_diff, one_shot] = fbmpc_refine_fxn(y, A, p1, ...
                                               sig2w, sig2s, mus, D, ...
                                               stop, E, r_stop, ...
                                               est, update)
                            
%% Fill in missing variables and initialize hyperparameter estimates

if nargin < 12
    update = 'mmse';
    if nargin < 11
        est = {'p1', 'sig2w', 'sig2s'};
        if nargin < 10
            r_stop = 0.02;
            if nargin < 9
                E = 10;
                if nargin < 8
                    stop = 0;
                    if nargin < 7
                        D = 5;
                        if nargin < 6
                            error('FMBPC_REFINE_FXN requires at least 6 arguments')
                        end
                    end
                end
            end
        end
    end
end

Q = length(mus);
[M, N] = size(A);

% Initial parameter estimates
p1_est = p1;
sig2w_est = sig2w;
sig2s_est = sig2s;
mus_est = mus;


%% Begin refinement procedure

% Run the FBMP algorithm and update estimates of the hyperparameters using
% surrogate update equations iteratively, until the parameter estimates
% converge to steady values, or 'E' iterations have occurred
for e = 1:E
    
    % Recover signal based on current parameter estimates
    one_shot = cputime;     % Start stopwatch
    [xmmse, xmmse_star, psy_star, nu_star, T_star, d_tot(e), d_max] = ...
        fbmpc_fxn_reduced(y, A, p1_est, sig2w_est, sig2s_est, mus_est, ...
                         D, stop);
    one_shot = cputime - one_shot;  % Capture algorithm runtime
    
    % Transfer estimates from new to old
    p1_old = p1_est;
    sig2w_old = sig2w_est;
    sig2s_old = sig2s_est;
    mus_old = mus_est;
    
    % Identify the appropriate estimator for updating
    switch lower(update)
        case 'mmse'
            x_r = xmmse;
        case 'map'
            x_r = xmmse_star{1};
        otherwise
            error('%s is not a valid estimator for parameter refinement', ...
                            update)
    end
    
    % Refine parameters using newest recovery results
    for p = 1:length(est)
        switch lower(est{p})        % Current parameter being refined
            case 'p1'
                p1_est = sum(x_r ~= 0)/N;        % Updated estimate of p1
            case 'sig2w'
                sig2w_est = var(y - A*x_r);   	% Updated estimate of sig2w
            case 'sig2s'
                if Q == 2
                    % Single mean, thus a conventional estimate of variance
                    % should suffice
                    sig2s_est = [0; var(x_r(x_r ~= 0))];    % Updated estimate of sig2s
                else
                    % Multiple means are in play, therefore use the k-means
                    % algorithm to find the means, and re-estimate the
                    % variance using distances from the clustered means.
                    % MAP estimator seems to work better for this purpose
                    [idx, centers, dummy, dists] = kmeans(xmmse_star{1}(xmmse_star{1} ~= 0), ...
                                                          (Q - 1), 'start', mus_est);
                    dists = min(dists, [], 2);      % Grab the distance from point-to-center
                    var_active = max([dists'*dists/(sum(xmmse_star{1} ~= 0) - 1), 1e-6]);
                    sig2s_est = [0; var_active];   % Updated estimate of sig2s
                end
            case 'mus'
                if Q == 2
                    % Single mean
                    mus_est = [0; mean(x_r(x_r ~= 0))];     % Updated estimate of mus
                else
                    % Multiple means, thus, use k-means clustering
                    % algorithm to identify possible means.  Using the MAP
                    % estimate seems to be superior for this purpose
                    [idx, centers, dummy, dists] = kmeans(xmmse_star{1}(xmmse_star{1} ~= 0), ...
                                                          (Q - 1), 'start', mus_est);
                    mus_est = [0; centers];     % Updated estimate of mus
                end
            case 'both'
                % Both sig2s and mus are to be re-estimated
                if Q == 2
                    mus_est = [0; mean(x_r(x_r ~= 0))];     % Updated estimate of mus
                    sig2s_est = [0; var(x_r(x_r ~= 0))];    % Updated estimate of sig2s
                else
                    % Multiple means are in play, therefore use the k-means
                    % algorithm to find the re-estimated means, and 
                    % re-estimate the variance using distances from the 
                    % clustered means.  MAP estimator seems to work better 
                    % for this purpose
                    [idx, centers, dummy, dists] = kmeans(xmmse_star{1}(xmmse_star{1} ~= 0), ...
                                                          (Q - 1), 'start', mus_est(2:end));
                    dists = min(dists, [], 2);      % Grab the distance from point-to-center
                    var_active = max([dists'*dists/(sum(xmmse_star{1} ~= 0) - 1), 1e-6]);
                    sig2s_est = [0; var_active];   % Updated estimate of sig2s
                    mus_est = [0; centers];     % Updated estimate of mus
                end
            otherwise
                error(sprintf('%s is not a parameter that can currently be refined', ...
                                est{p}))
        end
    end
    
    % Check for early termination
    percent_change = [abs(p1_est - p1_old)/p1_old; ...
                        abs(sig2w_est - sig2w_old)/sig2w_old; ...
                        abs(sig2s_est(2) - sig2s_old(2))/sig2s_old(2); ...
                        max(abs(mus_est(mus_est ~= 0) - mus_old(mus_old ~= 0))./...
                            abs(mus_old(mus_old ~= 0)))];
    if max(percent_change) < r_stop
        break
    end
end

d_tot(d_tot == inf) = D;    % Remove infinity flag and replace with maximum
d_tot = mean(d_tot);        % Return the average number of RGS iterations req'd
e_tot = e;                  % Return the total number of refinement stages
est_diff.p1 = abs(p1_est - p1)/p1;      % Percentage change in p1 estimate
est_diff.sig2w = abs(sig2w_est - sig2w)/sig2w;  % Percentage change in sig2w estimate
est_diff.sig2s = abs(sig2s_est(2) - sig2s(2))/sig2s(2); % Percentage change in sig2s estimate
