% SINUSOID_DEMO     This file demonstrates how a signal that is not sparse
% in the time domain can be transformed into a signal that is sparse in the
% frequency domain, at which point FBMP can be used to recover the signal,
% with a GEM iterative refinement algorithm re-estimating required
% hyperparameters
%

%
% Coded by: Justin Ziniel, Ohio State Univ.
% Email: ziniel.1@osu.edu
% Last modified: April 2, 2009
% Change summary:   Corrected form of the operator passed to FBMP (4/2/09)
%                   Switched from fbmpr_gem_refine_fxn to
%                       fbmpc_gem_refine_fxn (8/13/09)
% FBMP version 1.2 
% Copyright (c) Philip Schniter, Lee C. Potter, and Justin Ziniel, 2008
%

clear;


%% Setup an experiment

% User-defined parameters
M = 128;           	% # observations
N = 512;   		% # parameters
SNRdB = 40;		% SNR (in dB)
D = 20;			% maximum number of allowable repeated greedy searches
stop = 0;		% stop RGS @ this many stdevs below mean of metric
E = 10;             	% maximum number of allowable parameter refinement stages
r_stop = 0.02;      	% threshold for early termination of parameter refinement
                    	% (i.e. percentage change threshold)

% Initial estimates of hyperparameters
p1 = 0.01;		% prior probability of non-zero tap
sig2w_est = 0.1;  	% noise variance
sig2s = [0; 1];     	% sparse coefficient variances [off;on]
mus = [0; 0];       	% sparse coefficient means [off;on;...;on]

% Construct a periodic signal
k1 = round((1/16)*N);   % First frequency component (DFT index)
k2 = round((1/8)*N);    % Second frequency component (DFT index)
k3 = round((7/16)*N);   % Third frequency component (DFT index)
n = [1:N]';     	% Discrete-time index
x = (2/N)*(cos(2*pi*k1*n/N + 0.5*pi) + cos(2*pi*k2*n/N) + sin(2*pi*k3*n/N)); % Time domain signal
T = dftmtx(N);  	% Sparsifying transform (DFT matrix)
X = (1/sqrt(N))*T*x;    % Sparse signal in frequency domain


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

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


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


%% Run the FBMP algorithm w/ GEM-based re-estimation of hyperparameters

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

tic
% Call the FBMPR_GEM_REFINE function
[Xmmse, Xmmse_star, psy_star, nu_star, T_star, d_tot, d_max, e_tot] = ...
    fbmpc_gem_refine_fxn(y, A*(1/sqrt(N))*T', p1, sig2w_est, sig2s, mus, D, stop, E, r_stop);
toc

% Reconstruct time-domain signal from recovered freq. domain signal
xmmse = (1/sqrt(N))*T'*Xmmse;
for i = 1:length(Xmmse_star)    % Reconstruct each conditional estimate
    xmmse_star{i} = (1/sqrt(N))*T'*Xmmse_star{i};
end

% Performance metric
NMSE = (norm(xmmse - x)^2)/(norm(x)^2);


%% Plotting results

subplot(311)
handy = plot([x, real(xmmse)]);
legend('x (time-domain signal)', 'x_{MMSE}', 'Location', 'Best')
ttl = ['Sinusoid signal recovery, (NMSE = ' num2str(10*log10(NMSE)) ' dB)'];
title(ttl); ylabel('x[n]'); xlabel('n')

subplot(312)
handy = stem([abs(X), abs(Xmmse), abs(Xmmse_star{1})]);
set(handy(3),'Marker','o');
set(handy(3),'Marker','+');
set(handy(3),'Marker','x');
legend(handy,'X_{true}', 'X_{MMSE}', 'X_{MAP}', 'Location', 'Best')
ttl = ['Sparse domain recovery, (NMSE = ' num2str(10*log10(NMSE)) ' dB)'];
title(ttl); ylabel('magnitude X[k]'); xlabel('k')

subplot(313)
mask = abs(X) > 1e-8;   % Only display non-negligible values
handy = stem([angle(X).*mask, angle(Xmmse), angle(Xmmse_star{1})]);
set(handy(3),'Marker','o');
set(handy(3),'Marker','+');
set(handy(3),'Marker','x');
legend(handy,'X_{true}', 'X_{MMSE}', 'X_{MAP}', 'Location', 'Best')
ttl = ['Sparse domain recovery, (NMSE = ' num2str(10*log10(NMSE)) ' dB)'];
title(ttl); ylabel('angle X[k]'); xlabel('k')
