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

% calculateTrueCoefficientProbabilities.m

function [ w ] = calculateTrueCoefficientProbabilties( videoName, blk_num, sparsityBasis, frstart, frend, method, coeffTol, saveFigs, weightProbType )
% Inputs:
% videoName          = name of video sequence
% blk_num            = number of blocks each frame is broken into
% sparsityBasis.name = name of sparsity basis
% frstart            = starting frame
% frend              = ending frame
% method             = 'coeffNum', method for determining the "top" pixels in each frame
% coeffTol           = coeffTol*100% of coefficients are used to estimate probabilities 
% saveFigs           = flag to save figures
% weightProbType     = relationship between weights and probabilities, 'linear' or 'exponential'
%
% Outputs:
% w = weights determined from "true" probabilities (listed as a column vector)

% Load Video Data
switch videoName
    
    case 'foreman'
        
        switch sparsityBasis.name
            
            case 'DCT'
                load(sprintf('foreman_video_data_%dblocks.mat',blk_num));
                                
            otherwise
                error('Invalid sparsityBasis.name');        
        end
        
    otherwise
        error('Invalid videoName: Must be foreman.');
        
end

% Get Size of Frames
[vp, hp] = size(YQp{1}{1}); % vp = number of rows, hp = number of columns

% Define Sparsity Basis
switch sparsityBasis.name
    
    case 'DCT'
        D = opDCT2(vp, hp);
        titlestr = 'DCT';
                
    otherwise
        error('Invalid sparsityBasis.name: Must be DCT');
    
end

switch method
        
    case 'coeffNum'
        
        % Calculate Probabilities
        % loop on blocks
        for jj = 1:blk_num
            % loop on frames
            for ifrm = frstart:frend
                % Get largest coefficients  
                coeff = D*YQp{ifrm}{jj}(:); 
                [val idx] = sort(abs(coeff), 'descend');
                k1 = round(coeffTol*length(val));
                Tk1{ifrm}{jj} = idx(1:k1); % top indices containing the energy 
                if ifrm > 1
                    A{jj}(Tk1{ifrm}{jj}) = A{jj}(Tk1{ifrm}{jj}) + 1;
                else
                    A{jj} = zeros(size(YQp{ifrm}{jj}));
                    A{jj}(Tk1{ifrm}{jj}) = 1;
                end
            end
            switch weightProbType
                
                case 'linear'
                    w{jj} = 1-A{jj}(:)/(frend-frstart+1); % weights, listed as a column vector, determined linearly from probabilities
                    
                case 'exponential'
                    w{jj} = (exp(-5*A{jj}(:)/(frend-frstart+1))-exp(-5))/(1-exp(-5)); % weights determined exponentially from probabilities
                    
                otherwise
                    error('Invalid weightProbType');
            
            end
            % Plots counts and weights
            hc = figure; 
            imagesc(A{jj}); colorbar;
            title({['Counts, ' titlestr ', Block ' num2str(jj) ];...
                ['coeffTol=' num2str(coeffTol*100) '%, Frames ' num2str(frstart) ' to ' num2str(frend)]});
            hw = figure; 
            imagesc(reshape(w{jj},size(A{jj}))); colorbar;
            title({['Weights, ' titlestr ', Block ' num2str(jj) ', ' weightProbType ];...
                ['coeffTol=' num2str(coeffTol*100) '%, Frames ' num2str(frstart) ' to ' num2str(frend)]});

            % Do you want to save figures of the counts and weights?
            if (saveFigs) 
                filename_hc = sprintf('frameCoefficients_counts_%svideo_%ssparsity_%dblock_%dfrstart_%dfrend_%1.3coeffTol',videoName,sparsityBasis.name,jj,frstart,frend,coeffTol);
                print(hc,'-depsc','-r0',[filename_hc '.eps']) 
                saveas(hc,[filename_hc '.png'],'png'); 
                saveas(hc,[filename_hc '.fig'],'fig');

                filename_hw = sprintf('frameCoefficients_weights_%svideo_%ssparsity_%dblock_%dfrstart_%dfrend_%1.3coeffTol',videoName,sparsityBasis.name,jj,frstart,frend,coeffTol);
                print(hw,'-depsc','-r0',[filename_hc '.eps']) 
                saveas(hw,[filename_hw '.png'],'png'); 
                saveas(hw,[filename_hw '.fig'],'fig');
            end
        end
        
    otherwise
        error('Invalid method: must be coeffNum');

end


end

