% Copyright (C) 2016 by Deanna Needell, Rayan Saab, and Tina Woolf

% This file is part of ''Weighted L1-Minimization for Sparse Recovery under
%                        Arbitrary Prior Information''

% This program is free software: you can redistribute it and/or modify
% it under the terms of the GNU General Public License as published by
% the Free Software Foundation, either version 3 of the License, or
% (at your option) any later version.

% This program is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
% GNU General Public License for more details.

% You should have received a copy of the GNU General Public License
% along with this program.  If not, see <http://www.gnu.org/licenses/>.

% runWeightedL1_TwoSupportEstimates.m

function [error_WL1_single, error_WL1_double] = runWeightedL1_TwoSupportEstimates(n,s,mlist,ntrials,sigma,saveMat,saveFigs,...
    rho,alpha,w,rho1,alpha1,w1,rho2,alpha2,w2)
% runWeightedL1_TwoSupportEstimates.m
%
% Inputs:
% n        = signal dimension
% s        = sparsity level
% mlist    = number of measurements (as an array)
% ntrials  = number of trials
% sigma    = measurement noise standard deviation
% saveMat  = flag to save data to mat file
% saveFigs = flag to save figures
% rho      = size of single support estimate is rho*s
% alpha    = accuracy of single support estimate
% w        = weight with a single support estimate
% rho1     = size of first of two support estimates is rho1*s
% alpha1   = accuracy of first of two support estimates
% w1       = weight for first of two support estimates
% rho2     = size of second of two support estimates is rho2*s
% alpha2   = accuracy of second of two support estimates
% w2       = weight for second of two support estimates
%
% Outputs:
% error_WL1_single = array of size length(mlist) by ntrials containing
%                    the L2-norm error between the true signal and the
%                    estimated signal using a single weight/support
%                    estimate
% error_WL1_double = array of size length(mlist) by ntrials containing
%                    the L2-norm error between the true signal and the
%                    estimated signal using two weights/support
%                    estimates

%% Initialize error results arrays
error_WL1_single = zeros(length(mlist),ntrials);
error_WL1_double = zeros(length(mlist),ntrials);

%% Loop through Measurements and Trials
% Loop through number of measurements
for mm = 1:length(mlist)
   m = mlist(mm);
   % Loop through trials
   for nn = 1:ntrials
      
        % Generate true support and signal
        temp = randperm(n);
        T    = temp(1:s);
        x    = zeros(n,1);
        x(T) = randn(s, 1);
        
        % Generate single support estimate with accuracy alpha
        idxOfT  = randperm(length(T));
        idxOfT  = T(idxOfT(1:round(rho*alpha*s))); % support estimate indices from T
        Tc      = setdiff(1:n,T);                  % T-complement
        idxOfTc = randperm(length(Tc));
        idxOfTc = Tc(idxOfTc(1:(rho*s-round(rho*alpha*s)))); % support estimate indices from Tc
        Ttilde  = union(idxOfT,idxOfTc);           % support estimate
        Ttildec = setdiff(1:n,Ttilde);             % Ttilde-complement 
        
        % Check size of Ttilde
        if ~( length(Ttilde) == round(rho*s) )
            error('Ttilde should be of length rho*s');
        end
       
        % Check size of Ttilde
        if ~( length(Ttilde) == round(rho*s) )
            error('Ttilde should be of length rho*s');
        end
       
        % If rho1+rho2=rho, generate two support estimates with accuracies 
        % alpha1 and alpha2 such that Ttilde = union(Ttilde1,Ttilde2).
        % Otherwise, simply generate two disjoint support estimates with  
        % accuracies alpha1 and alpha2.
        if (rho1 + rho2 == rho)
            idxOfT_Ttilde = randperm(length(idxOfT));
            idxOfT_forTtilde1 = idxOfT(idxOfT_Ttilde(1:round(rho1*alpha1*s)));
            idxOfT_forTtilde2 = idxOfT(idxOfT_Ttilde(round(rho1*alpha1*s)+1:end));
            idxOfT_Ttildec = randperm(length(idxOfTc));
            idxOfTc_forTtilde1 = idxOfTc(idxOfT_Ttildec(1:(round(rho1*s)-round(rho1*alpha1*s))));
            idxOfTc_forTtilde2 = idxOfTc(idxOfT_Ttildec((round(rho1*s)-round(rho1*alpha1*s))+1:end));
            Ttilde1 = union(idxOfT_forTtilde1,idxOfTc_forTtilde1);
            Ttilde2 = union(idxOfT_forTtilde2,idxOfTc_forTtilde2);
            
            if ~( isequal(Ttilde, union(Ttilde1,Ttilde2)) )
               error('Ttilde and union(Ttilde1,Ttilde2) should be the same!'); 
            end
        else
            idxOfT  = randperm(length(T));
            idxOfT_forTtilde1 = T(idxOfT(1:round(rho1*alpha1*s)));
            idxOfT_forTtilde2 = T(idxOfT(round(rho1*alpha1*s+1):round(rho1*alpha1*s+rho2*alpha2*s)));
            idxOfTc = randperm(length(Tc));
            idxOfTc_forTtilde1 = Tc(idxOfTc(1:(round(rho1*s)-round(rho1*alpha1*s))));
            idxOfTc_forTtilde2 = Tc( idxOfTc( (round(rho1*s)-round(rho1*alpha1*s)+1): (round(rho1*s)-round(rho1*alpha1*s)+round(rho2*s)-round(rho2*alpha2*s))));
            Ttilde1 = union(idxOfT_forTtilde1,idxOfTc_forTtilde1);
            Ttilde2 = union(idxOfT_forTtilde2,idxOfTc_forTtilde2);
        end
        
        % Check size of Ttilde1
        if ~( length(Ttilde1) == round(rho1*s) )
            error('Ttilde1 should be of length rho1*s');
        end
        
        % Check size of Ttilde2
        if ~( length(Ttilde2) == round(rho2*s) )
            error('Ttilde2 should be of length rho2*s');
        end
        
        % Check that size of Ttilde and Ttilde1+Ttilde match
        if ~( length(Ttilde) == length(Ttilde1) + length(Ttilde2) )
           error('Ttilde and Ttilde1+Ttilde2 have different sizes!') 
        end
               
        % Measurement matrix
        A = 1/sqrt(m)*randn(m,n);
        
        % Noise
        z = randn(m,1)*sigma;
        e = norm(z);
        
        % Measurements
        y = A*x + z;
        
        % Weights for single support estimate
        weights         = ones(n,1);
        weights(Ttilde) = w;
        W               = diag(weights);
        
        % Weights for two support estimates
        weights2          = ones(n,1);
        weights2(Ttilde1) = w1;
        weights2(Ttilde2) = w2;
        W2                = diag(weights2);
        
        % Solve weighted L1 with single support estimate
        cvx_begin 
        cvx_quiet(true)
            variable xhat1(n); 
            minimize( norm(W*xhat1,1) ); 
            subject to
                norm(y - A*xhat1) <= e;
        cvx_end
        
        % Solve weighted L2 with two support estimates
        cvx_begin 
        cvx_quiet(true)
            variable xhat2(n); 
            minimize( norm(W2*xhat2,1) ); 
            subject to
                norm(y - A*xhat2) <= e;
        cvx_end
        
        % Calculate error
        error_WL1_single(mm,nn) = norm(xhat1-x);
        error_WL1_double(mm,nn) = norm(xhat2-x);
       
   end
   disp(['m ' num2str(m)]);
end

%% Save and Plot Data

% Save data
filename = sprintf('weightedL1_TwoSupportEstimates_%ds_%dn_%dtrials_%1.2fsigma_%1.2frho_%1.2falpha_%1.2fw_%1.2frho1_%1.2falpha1_%1.2fw1_%1.2frho2_%1.2falpha2_%1.2fw2',...
    s,n,ntrials,sigma,rho,alpha,w,rho1,alpha1,w1,rho2,alpha2,w2);
if(saveMat)
   save([filename '.mat']); 
end

% Plot
figure; 
set(gcf,'WindowStyle','normal')
semilogy(mlist,mean(error_WL1_single,2),'r','linewidth',2); hold on;
semilogy(mlist,mean(error_WL1_double,2),'g','linewidth',2);
legend('Single','Double');
xlabel('Measurements','fontsize',14);
ylabel('Mean Recovery Error','fontsize',14);
set(gca,'fontsize',12);
title({['n=' num2str(n) ', s=' num2str(s) ', trials=' num2str(ntrials) ', rho=' num2str(rho) ', alpha=' num2str(alpha) ', w=' num2str(w)];...
    [ 'rho1=' num2str(rho1) ', alpha1=' num2str(alpha1) ', w1=' num2str(w1) ', rho2=' num2str(rho2) ', alpha2=' num2str(alpha2) ', w2=' num2str(w2)]},...
    'fontsize',14);
grid on;
set(gcf,'PaperPositionMode','Auto');
filename = sprintf('weightedL1_TwoSupportEstimates_cvx_%ds_%dn_%dtrials_%1.2fsigma_%1.2frho_%1.2falpha_%1.2fw_%1.2frho1_%1.2falpha1_%1.2fw1_%1.2frho2_%1.2falpha2_%1.2fw2'...
    ,s,n,ntrials,sigma,rho,alpha,w,rho1,alpha1,w1,rho2,alpha2,w2);
if (saveFigs)
    print('-depsc','-r0',[filename '.eps']);
    saveas(gcf,[filename '.png'],'png');
    saveas(gcf,[filename '.fig'],'fig');
end
