% 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/>.

% runPriorDensity.m

%% Cleanup and Setup
clc;
close all;
clearvars;
rng('default') % seed the random number generator to produce a predictable sequence of numbers

% Basic Inputs
n            = 256;                 % signal dimension
s            = 24;                  % sparsity (only used for tree ptype signals)
ntrials      = 50;                  % number of trials
mlist        = [20:10:250];         % number of measurements
sigma        = 0.01;                % noise standard deviation
w_single     = [0 0.3 0.5 0.7 0.9]; % weights for weighted L1 with a single weight
w_multi_type = 'exponential';       % linear or exponential relationship between weights and probabilities
ptype        = 'tree';              % signal density type: uniform, block, power, sqrt, tree
punif        = 0.06;                % nonzero probability for uniform density
pthresh      = 0.025;               % if nonzero probability is less than pthresh, just set it to 0 (not used for uniform ptype)
pblock       = [1 0.9 0.7 0.5 0.2]; % nonzero probabilities of blocks in block ptype; length(pblock)= number of blocks, and each block contains round(s/nblock) indices
saveFigs     = 0;                   % flag for saving figures
saveMat      = 0;                   % flag for saving variables to mat file

%% Define Nonzero Probabilities per Index

switch (ptype)

    case 'uniform'
        % Uniform Probabilities
        p = ones(1,n)*punif; 
        figTitle = sprintf('ptype=%s, punif=%1.3f',ptype,punif);
        
    case 'block'
        % Block-Sparse Probabilities
        p = zeros(1,n);
        N = length(pblock);
        validIdx = 1:round(s/N):(n-round(s/N));
        tmp = randperm(length(validIdx));
        for ii = 1:N
           p( validIdx(tmp(ii)):(validIdx(tmp(ii))+round(s/N)-1) ) = pblock(ii); 
        end
        figTitle = sprintf('ptype=%s, nBlocks=%d, nPerBlock=%d',ptype,N,round(s/N));

    case 'power'
        % Power Law Probabilities
        for ii = 1:n
           p(ii) = 1/ii;
           if ( p(ii) < pthresh)
               p(ii) = 0;
           end
        end
        figTitle = sprintf('ptype=%s, pthresh=%1.3f',ptype,pthresh);
        
    case 'sqrt'
        % Square Root Power Law Probabilities
        for ii = 1:n
           p(ii) = 1/sqrt(ii);
           if ( p(ii) < pthresh)
               p(ii) = 0;
           end
        end
        figTitle = sprintf('ptype=%s, pthresh=%1.3f',ptype,pthresh);
        
    case 'tree'
        % Binary Tree Probabilities
        tree_trials = 10000;
        p = zeros(1,n);
        for ii = 1:tree_trials
           a = genSignal(n,s);
           p = (abs(a)>0) + p;
        end
        p = p/tree_trials;
        p(p < pthresh) = 0;
        figTitle = sprintf('ptype=%s, pthresh=%1.3f, s=%d',ptype,pthresh,s);

    otherwise
        error(['ptype ' ptype ' is not defined; use either uniform, block, power, sqrt, or tree']);

end

%% Define Weights for Multi-Weighted L1
switch w_multi_type
    
    case 'linear'
        w_multi = 1 - p;
    case 'exponential'
        w_multi = (exp(-5*p)-exp(-5))/(1-exp(-5));
        figure; set(gcf,'WindowStyle','normal'); set(gcf,'PaperPositionMode','Auto');
        plot(0:0.05:1,(exp(-5*(0:0.05:1))-exp(-5))/(1-exp(-5)),'linewidth',2);
        xlabel('Probability of Nonzero','fontsize',20);
        ylabel('Weight','fontsize',20);
        title('Exponential Relationship Between Weights and Probabilites','fontsize',20);
        set(gca,'FontSize',20);
        set(gcf,'Position',[400 300 600 430]);
        if (saveFigs) 
           filename = sprintf('PriorDensity_WeightsProb_%s',w_multi_type);
           print('-depsc','-r0',[filename '.eps']) 
           saveas(gcf,[filename '.png'],'png'); 
           saveas(gcf,[filename '.fig'],'fig');
        end
    otherwise
        error(['w_multi_type ' w_multi_type ' is not defined!']);

end

W_multi = diag(w_multi);

%% Loop on Measurements and Trials

% Initialize error result array
errorL1        = zeros(length(mlist),ntrials);
errorWL1       = zeros(length(w_single),length(mlist),ntrials);
errorWL1_multi = zeros(length(mlist),ntrials);
trueSparsity   = zeros(length(mlist),ntrials);

% Loop on number of measurements
for mm = 1:length(mlist)
    m = mlist(mm);
    % Loop on trials
    for nn = 1:ntrials
        
        % Generate a support given these probabilities
        switch (ptype)
            
            case {'uniform','block','power','sqrt'}
            randNum = rand(1,n); % generate uniformly random values between 0 and 1
            suppIdx = randNum < p;
            
            case 'tree'
                a = genSignal(n,s);
                suppIdx = abs(a) > 0;
                
            otherwise
                error(['ptype ' ptype ' is not defined; use either uniform, block, power, sqrt, or tree']);
        end
        
        trueSparsity(mm,nn) = sum(suppIdx);
        T                   = find(suppIdx);
        
        % Generate signal supported on T
        x       = zeros(n,1);
        x(T)    = randn(length(T),1);
        
        % Measurement Matrix
        A = 1/sqrt(m)*randn(m,n);
        
        % Noise
        z = randn(m,1)*sigma;
        e = norm(z);
        
        % Measurements
        y = A*x + z;
        
        % Solve via L1-Minimization
        cvx_begin 
        cvx_quiet(true)
            variable xhat(n); 
            minimize( norm(xhat,1) ); 
            subject to
                norm(y - A*xhat) <= e;
        cvx_end
        errorL1(mm,nn)        = norm(xhat-x);
        
        % Solve via Weighted L1-Minimization with a single weight
        for ii = 1:length(w_single)
            w = ones(n,1);
            w(p>0) = w_single(ii);
            W = diag(w);
            
            cvx_begin 
            cvx_quiet(true)
                variable xhatw(n); 
                minimize( norm(W*xhatw,1) ); 
                subject to
                    norm(y - A*xhatw) <= e;
            cvx_end
            errorWL1(ii,mm,nn) = norm(xhatw-x);
        end
        
        % Solve via Weighted L1-Minimization with multi-weights
        cvx_begin 
        cvx_quiet(true)
            variable xhatw_multi(n); 
            minimize( norm(W_multi*xhatw_multi,1) ); 
            subject to
                norm(y - A*xhatw_multi) <= e;
        cvx_end
        errorWL1_multi(mm,nn) = norm(xhatw_multi-x);

    end
    disp(['m ' num2str(m)]);
end
disp(['Average True Total Sparsity: ' num2str(mean(mean(trueSparsity)),4)]);

switch ptype
    case {'uniform','block','power','sqrt'}
        figTitle = sprintf('%s, Avg True s=%1.3f',figTitle,mean(mean(trueSparsity)));
        
    case 'tree'
        %do nothing
        
    otherwise
        error(['ptype ' ptype ' is not defined; use either uniform, block, power, sqrt, or tree']);
end

%% Plots
myLineStyles = {'-cv','-g^','-kx','-ms', '-y+','--cv','--g^','--kx','--ms', '--y+'};

figure; set(gcf,'WindowStyle','normal'); set(gcf,'PaperPositionMode','Auto');
LegendString{1} = 'L1 (w=1)'; LegendString{2} = 'Non-uniform Weights';
semilogy(mlist,mean(errorL1,2),'r','linewidth',2); hold on;
semilogy(mlist,mean(errorWL1_multi,2),'--b','linewidth',2);
for ii = 1:length(w_single)
   semilogy(mlist,mean(squeeze(errorWL1(ii,:,:)),2),myLineStyles{ii},'linewidth',2); 
   LegendString{2+ii} = ['w=' num2str(w_single(ii))];
end
hl = legend(LegendString,1);
set(hl,'fontsize',20);
xlabel('Measurements','fontsize',20);
ylabel('Mean Recovery Error','fontsize',20);
%title({['n=' num2str(n) ', trials=' num2str(ntrials) ', sigma=' num2str(sigma)];figTitle},'fontsize',20);
set(gca,'FontSize',20);
xlim([mlist(1) mlist(end)]);
ylim([min( [min(mean(errorL1,2)),min(mean(errorWL1_multi,2)),min(min(mean(errorWL1,3)))] ) ...
      max( [max(mean(errorL1,2)),max(mean(errorWL1_multi,2)),max(max(mean(errorWL1,3)))] )]);
grid on;
set(gcf,'Position',[370 350 680 420]);
if (saveFigs) 
   filename = sprintf('PriorDensity_ErrorvsMeasurements_%dn_%dtrials_%1.4fsigma_%sptype',n,ntrials,sigma,ptype);
   print('-depsc','-r0',[filename '.eps']) 
   saveas(gcf,[filename '.png'],'png'); 
   saveas(gcf,[filename '.fig'],'fig');
end

figure; set(gcf,'WindowStyle','normal'); set(gcf,'PaperPositionMode','Auto');
plot(w_multi,'bx','markersize',8); hold on;
plot(p,'r.','markersize',12);
hl = legend('Non-uniform Weights','Prob Nonzero',1);
xlabel('Support Index','fontsize',20);
ylabel('Weight or Probability of Nonzero','fontsize',20);
%title('Multi-Weights and Probability of Nonzero','fontsize',20);
set(gca,'FontSize',20);
xlim([1 n]);
ylim([0 1]);
if (saveFigs) 
   filename = sprintf('PriorDensity_WeightsProb_%dn_%dtrials_%1.4fsigma_%sptype',n,ntrials,sigma,ptype);
   print('-depsc','-r0',[filename '.eps']) 
   saveas(gcf,[filename '.png'],'png'); 
   saveas(gcf,[filename '.fig'],'fig');
end

if (saveMat)
   filenameMat = sprintf('PriorDensity_%dn_%dtrials_%1.4fsigma_%sptype',n,ntrials,sigma,ptype);
   save([filenameMat '.mat']);
end