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

% video_wL1_recovery.m

%% Cleanup and Setup
close all;
clearvars;
clc; 

%% Initialize Parameters

% Read/Load Data Parameters
video_seq = 'foreman';      % foreman video sequence
blk_num   = 4;              % number of blocks to breakup each frame into
  
% Basic Experiment Parameters
frstart        = 1;         % frame to start at 
frend          = 300;       % frame to end at
mlist_fraction = [0.5];     % number of measurements to use: N*mlist_fraction 

% Measurement and Sparsity Basis Parameters
sparsityBasis.name = 'DCT'; % Discrete Cosine Transform

% Weight Method Parameters
coeffTol  = 0.1;            % coeffTol*100% of coefficients to use in support estimate 
method    = 'coeffNum';     % method to use in calculation of true probabilities 
vW        = [0.2 0.5];      % weights for weighted L1 with 1 weight
WL1_opts  = [1];            % Options for weighted L1 with 1 weight:
                              % opt 1 = Use coeffTol*100 percent of the coefficients to determine the support estimate
WL1M_opts = [1 2];          % Options for weighted L1 with multi-weights:
                              % opt 1 = Use true top coefficients from all frames (probabilities estimated exponentially)
                              % opt 2 = Estimate probability of each index being in the top (probabilities estimated exponentially)

% Save Output Parameters
saveFigs = 0;               % 1=save figures, 0=do not save figures
saveMat  = 0;               % 1=save data to mat file, 0=do not save data to mat file

% SPGL1 Parameters
sigma              = 1e-6;  % noise bound in spgl1 solver
options.verbosity  = 0;
options.optTol     = 1e-8;
options.bpTol      = 1e-8;
options.iterations = 50000;

%% Load Data

disp('Loading data from mat file');
load(sprintf('%s_video_data_%dblocks.mat',video_seq,blk_num));

%% Set up Measurement Operator

% Number of Measurements
N     = length(YQ{1}(:))/blk_num; % total number of pixels per block
mlist = N.*mlist_fraction;
m_max = max(mlist);               % maximum number of measurements

% Define Measurement Operator
disp('Defining measurement operator and sparsity basis');
[PHI, R, D] = defineMeasurementOperator(N, m_max, vp, hp, sparsityBasis);

%% Load true weights 

trueMultiweights_exp = calculateTrueCoefficientProbabilties( video_seq, blk_num, sparsityBasis, 1, 300, method, coeffTol, 0, 'exponential' );
                       
%% Process each frame via L1, Weighted-L1 with 1 weight, and Weighted-L1 with Multi-weights

% Loop on frames
for ifrm = frstart:frend
        
    disp(['Beginning loop on frame ' num2str(ifrm)]);

    %%% Loop on the number of measurements
    for im = 1:length(mlist)
        
        m  = mlist(im);
        A  = PHI(1:m,:)/sqrt(m);
        RA = R(1:m,:)/sqrt(m);
        
        % Loop on blocks
        for j = 1:blk_num
            
            disp(['Frame ' num2str(ifrm) ', Meas ' num2str(im) ', blk_num ' num2str(j)]); % display iteration

            % Sample
            b = RA*YQp{ifrm}{j}(:);   
        
            % Recover with L1
            disp('Recovering with L1');
            options.weights = [];
            x_sp_rec_k1L1   = spgl1(A, b, 0, sigma, zeros(N,1), options);

            % reshape reconstructed vector
            YdctQprec_L1{ifrm}{im}{j} = reshape(x_sp_rec_k1L1, size(YQp{ifrm}{j}));
            YQprec_L1{ifrm}{im}{j}    = reshape(double(uint8(round(D\YdctQprec_L1{ifrm}{im}{j}(:)))),size(YdctQprec_L1{ifrm}{im}{j}));

            % Solve with Weighted-L1 with 1 weight
            % Loop on weights
            for iw = 1:length(vW)  
                
                for iwo = 1:length(WL1_opts)
                    disp(['Weight ' num2str(vW(iw)) ' Opt ' num2str(WL1_opts(iwo))]);
                    
                    switch WL1_opts(iwo)
                        
                        case 1
                            
                            if ifrm == frstart % use L1 for first frame reconstruction                        
                                YdctQprec_WL1{ifrm}{im}{j}{iw}{iwo} = YdctQprec_L1{ifrm}{im}{j};
                                YQprec_WL1{ifrm}{im}{j}{iw}{iwo}    = YQprec_L1{ifrm}{im}{j};
                                continue;
                            end
                            
                            % Use all previous frames for support estimate
                            if ifrm < frstart + 2 % only second frame should go here
                               WL1_totalSupport_prev_coeff{im}{j}{iw} = []; % support from all previous frames
                            end
                            [val idx] = sort(abs(YdctQprec_WL1{ifrm-1}{im}{j}{iw}{iwo}(:)), 'descend');
                            k1  = round(coeffTol*length(val));
                            Tk1 = idx(1:k1);                 
                            Tk1 = union(Tk1,WL1_totalSupport_prev_coeff{im}{j}{iw});
                            disp(['Tk1 size: ' num2str(length(Tk1))]); % display set size
                            WL1_totalSupport_prev_coeff{im}{j}{iw} = Tk1; % update the support estimate
                        
                        otherwise
                            error(['The WL1 method iwo ' num2str(iwo) ' is not defined!']);
                    
                    end % end switch(WL1_opts(iwo))

                    % use constant weighting on Tk1
                    W      = ones(N,1);
                    w      = vW(iw);  
                    W(Tk1) = w+0.001;

                    % Recover with WL1
                    disp('Recovering with WL1');
                    options.weights = W;
                    x_sp_rec_k1L1   = spgl1(A, b, 0, sigma, zeros(N,1), options);

                    % reshape reconstructed vector
                    YdctQprec_WL1{ifrm}{im}{j}{iw}{iwo} = reshape(x_sp_rec_k1L1, size(YQp{ifrm}{j}));
                    YQprec_WL1{ifrm}{im}{j}{iw}{iwo}    = reshape(double(uint8(round(D\YdctQprec_WL1{ifrm}{im}{j}{iw}{iwo}(:)))),size(YdctQprec_WL1{ifrm}{im}{j}{iw}{iwo}));

                end % for iwo on WL1_opts
                
            end % for iw
            
            %%% Solve with Weighted-L1 with mulitple weights
            % Loop on multi-weight options
            for iwo = 1:length(WL1M_opts)  
                disp(['Multi-Weight Opt ' num2str(WL1M_opts(iwo))]);
                
                switch(WL1M_opts(iwo))
                        
                    case 1 
                        % Use true support from all frames (probabilities estimated exponentially)
                        W = trueMultiweights_exp{j} + 0.001;
                        
                    case 2
                        
                        if ifrm == frstart % use L1 for first frame reconstruction                        
                            YdctQprec_WL1M{ifrm}{im}{j}{iwo} = YdctQprec_L1{ifrm}{im}{j};
                            YQprec_WL1M{ifrm}{im}{j}{iwo}    = YQprec_L1{ifrm}{im}{j};
                            continue;
                        end
                        
                        % Initialize counts if this is the second frame
                        if (ifrm == frstart + 1)
                           idxCounts_coeff_exp{im}{j} = zeros(N,1); 
                        end
                        
                        % Use coeffTol * length(val) coefficients in the support estimate
                        [val idx] = sort(abs(YdctQprec_WL1M{ifrm-1}{im}{j}{iwo}(:)), 'descend');
                        k1 = round(coeffTol*length(val)); % number of coefficients to use
                        
                        % Update counts and set weights
                        idxCounts_coeff_exp{im}{j}(idx(1:k1)) = idxCounts_coeff_exp{im}{j}(idx(1:k1)) + 1; 
                        W = (exp(-5*idxCounts_coeff_exp{im}{j}./(ifrm-frstart))-exp(-5))/(1-exp(-5)) + 0.001;
                        
                    otherwise
                        error(['The multi-weight method iwo ' num2str(iwo) ' is not defined!']);
                
                end % switch(iwo)

                % Recover with WL1M
                disp('Recovering with WL1M');
                weights_WL1M{ifrm}{im}{j}{iwo} = W; % store the weights for WL1M
                options.weights = W;
                x_sp_rec_k1L1   = spgl1(A, b, 0, sigma, zeros(N,1), options);

                % reshape reconstructed vector
                YdctQprec_WL1M{ifrm}{im}{j}{iwo} = reshape(x_sp_rec_k1L1, size(YQp{ifrm}{j}));
                YQprec_WL1M{ifrm}{im}{j}{iwo}    = reshape(double(uint8(round(D\YdctQprec_WL1M{ifrm}{im}{j}{iwo}(:)))),size(YdctQprec_WL1M{ifrm}{im}{j}{iwo}));
            end % for iwo on WL1M_opts         

        end % for j = 1:blk_num
        
        % rebuild frames
        if (ifrm == frstart)
            YQrec_L1 = []; YQrec_WL1 = []; YQrec_WL1M = []; % initialize these variables since this is the first frame
        end
        [ YQrec_L1, YQrec_WL1, YQrec_WL1M ] = rebuildFrame( blk_num, ifrm, im, vW, WL1_opts, WL1M_opts, YQprec_L1, YQrec_L1, YQprec_WL1, YQrec_WL1, YQprec_WL1M, YQrec_WL1M );

        % calculate PSNR
        PSNR_Y_L1(ifrm,im) = psnr2(YQrec_L1{ifrm}{im}, YQ{ifrm});
        for iw = 1:length(vW)
            for iwo = 1:length(WL1_opts)
                PSNR_Y_WL1(ifrm,im,iw,iwo) = psnr2(YQrec_WL1{ifrm}{im}{iw}{iwo}, YQ{ifrm});
            end
        end
        for iwo = 1:length(WL1M_opts)
            PSNR_Y_WL1M(ifrm,im,iwo) = psnr2(YQrec_WL1M{ifrm}{im}{iwo}, YQ{ifrm});
        end

    end % for im
    
    if (saveMat)
        filenameMat = sprintf('videoRecovery_%s_%dblocks_%dfrstart_%dfrend_%s',video_seq,blk_num,frstart,frend,sparsityBasis.name);
        save([filenameMat '.mat']);
    end
    
    disp(['Done with frame ' num2str(ifrm)]);
    
end %for ifrm

%% Plot Results

im = 1; % plot first measurement level
myLineStyles = {'-','--','-.'};
myMarkers    = {'o','+','^','<','>','d','*','s','p'};
markerEveryNPoints = 5; % plot a marker every 5 data points (otherwise plot is too crowded)

figure; hold on; set(gcf,'WindowStyle','normal'); set(gcf,'PaperPositionMode','Auto');
legendCounter = 0;
plotHandles   = []; % store plot handles to be plotted
% Plot L1
h_L1 = plot(frstart:frend, PSNR_Y_L1(frstart:frend,im),'-r','linewidth',2);
legendCounter               = legendCounter+1;
legendString{legendCounter} = sprintf('L1 (w=1)');
plotHandles(legendCounter)  = h_L1;
% Plot WL1
for iw = 1:length(vW)
    for iwo = 1:length(WL1_opts)
        h_WL1_line(im,iw,iwo)   = plot(frstart:frend, PSNR_Y_WL1(frstart:frend,im,iw,iwo),'-k','linewidth',1);
        h_WL1_marker(im,iw,iwo) = plot(frstart:markerEveryNPoints:frend, PSNR_Y_WL1(frstart:markerEveryNPoints:frend,im,iw,iwo),sprintf('%sk',myMarkers{iw}),'linewidth',1,'markersize',8);
        legendCounter               = legendCounter+1;
        legendString{legendCounter} = sprintf('w=%1.1f',vW(iw));
        plotHandles(legendCounter)  = h_WL1_marker(im,iw,iwo);
    end
end
% Plot WL1M
iwo = 2; % this is for opt 2
h_WL1M_2_line   = plot(frstart:frend, PSNR_Y_WL1M(frstart:frend,im,iwo),'-m','linewidth',1);
h_WL1M_2_marker = plot(frstart:markerEveryNPoints:frend, PSNR_Y_WL1M(frstart:markerEveryNPoints:frend,im,iwo),'m>','linewidth',1,'markersize',8);
legendCounter               = legendCounter+1;
legendString{legendCounter} = sprintf('Multi-Weight');
plotHandles(legendCounter)  = h_WL1M_2_marker;

iwo = 1; % this is for opt 1
h_WL1M_1_line = plot(frstart:frend, PSNR_Y_WL1M(frstart:frend,im,iwo),'-.b','linewidth',2);
legendCounter               = legendCounter+1;
legendString{legendCounter} = sprintf('Multi-Weight Truth');
plotHandles(legendCounter)  = h_WL1M_1_line;       

% Add frames to plot
keyFrames  = [   86        153       177        189           197        208         222       275];    % which frames to plot
yLocations = {[26 27.7],[32.8 34.5],[25.3 27],[35 36.7],[26.8 28.5],[35.4 37.1],[28.8 30.5],[28.3 30]}; % the y position of the frames
xLocations = {[70 100],   [137 167],[153 183],  [164 194],  [184 214],[197 227],  [221 251] [259 289]}; % the x position of the frames
for ii = 1:length(keyFrames)
   imagesc(xLocations{ii}, yLocations{ii},flipud(YQ{keyFrames(ii)})); colormap(gray); 
   text(xLocations{ii}(1)+4,yLocations{ii}(2)+0.2,sprintf('Frame %d',keyFrames(ii)),'fontsize',14); % add text for the frame number
   hframe = plot(keyFrames(ii),PSNR_Y_L1(keyFrames(ii),im),'g.','markersize',20);     % add marker for the frame
end
legendCounter = legendCounter+1;
legendString{legendCounter} = sprintf('True Frames');
plotHandles(legendCounter) = hframe;

hl = legend(plotHandles,legendString,2);
set(hl,'fontsize',20);
set(gca,'fontsize',20);
ylim([25,37.5]);
xlabel('Frame Number','fontsize',20);
ylabel('PSNR','fontsize',20);
grid on;
set(gcf,'Position',[300 0 1100 600]);

if (saveFigs) 
   filename = sprintf('videoRecovery_%s_%dblocks_%dfrstart_%dfrend_%s',video_seq,blk_num,frstart,frend,sparsityBasis.name);
   print('-depsc','-r0',[filename '.eps']) 
   saveas(gcf,[filename '.png'],'png'); 
   saveas(gcf,[filename '.fig'],'fig');
end

