function res = ex311(transform, K, targetPSNR, noIT, testLena, verbose)
% ex311           Training of dictionary for images using ILS-DLA (MOD)
%                 Each 8x8 image block, or 8x8 coefficient block, are made
% into a column vector of length N = 64.
% The training set is generated by getXfrom8images (12000 vectors), 
% and the mean is subtracted before the set is used.
% The results are saved in 'ex311mmmddhhmm.mat'.
% 
% use:
%  res = ex311(transform, K, targetPSNR, noIT, testLena, verbose)
%-------------------------------------------------------------------------
% arguments:
%   transform   may be:  'none', 'dct', 'lot', 'elt', 'db79', 'm79'
%               see myim2col.m for more
%   K           number of vectors in dictionary
%   targetPSNR  target Peak Signal to Noise Ratio, should be >= 30
%               R = X - D*W;  % not an image (R ~= A - Ar)
%               sumRR = sum(sum(R.*R));
%               PSNR = 10*log10( numel(R)*255^2 / sumRR )
%   noIT        number of iterations through the training set
%   testLena    0 or 1, default 0. (imageapprox.m is used)
%   verbose     0 or 1, default 0.
%-------------------------------------------------------------------------
% example:
%  res = ex311('m79', 128, 30, 50);         % the fast example
%  res = ex311('m79', 440, 36, 200, 1, 1);  % more realistic example
%  res = ex311('lot', 128, 35, 500);        % a LOT example
%  res = ex311('many');     % adjust m-file to learn many dictionaries                            

%-------------------------------------------------------------------------
% Copyright (c) 2009.  Karl Skretting.  All rights reserved.
% University of Stavanger.
% Mail:  karl.skretting@uis.no   Homepage:  http://www.ux.uis.no/~karlsk/
% 
% HISTORY:  dd.mm.yyyy
% Ver. 1.0  30.06.2009  KS: test started
% Ver. 1.1  08.08.2009  KS: dictionaries do work
% Ver. 1.2  11.08.2009  KS: Make as function
% Ver. 1.3  12.08.2009  KS: Made the function simpler
% Ver. 1.4  18.08.2009  KS: more changes
% Ver. 1.5  29.10.2009  KS: testLena part is mainly in sparseImage
% Ver. 1.6  02.11.2009  KS: use ex314 to plot SNR during training
% Ver. 1.7  20.01.2010  KS: use getXfrom8images (not dataXimage)
% Ver. 1.8  09.08.2011  KS: Simplified a little bit
%-------------------------------------------------------------------------

mfile = 'ex311';

if ((nargin == 1) && strcmpi(transform,'many'))
    K = 440;   % 128, 440 (DC excluded), 512
    noit = 2000;
    target = 36;
    for i = 1
        ex311( 'm79', K, target, noit);
        ex311('none', K, target, noit);
    end
    res = 'done';
    return;
end

if (nargin < 3)
   error([mfile,': wrong number of arguments, see help.']);
end
if (nargin < 4); noIT = 500; end;
if (nargin < 5); testLena = 0; end;
if (nargin < 6); verbose = 0; end;

maxS = 40;
N = 64;
substep = 5;
noIT = substep*ceil(noIT/substep);

% limit for absolute error, ORMP end when ||r|| < maxError
% PSNR = 10*log10( numel(.)*255^2 / sum( ||r_i||^2 ) )
% setting maxError, the average error will be smaller
maxError = 1.5*sqrt(N)*255*10^(-targetPSNR/20);  
relLim = 1e-6;

X = getXfrom8images('t',transform, 'getFixedSet',1, 'v',1);
[N,L] = size(X);
%
disp(' ');
disp([mfile,': start ILS-DLA (MOD) training for images, ',datestr(now()),...
    ', N=64, K=',int2str(K),', L=',int2str(L),', noIT=',int2str(noIT),...
    ', target PSNR = ',num2str(targetPSNR) ]);

% always remove DC here
if strcmpi(transform,'none')
    X = X - ones(64,1)*mean(X);
else
    X(1,:) = zeros(1,L);
end
%
% for many training vectors the DC element will be enough
% the dictionary is trained for the rest
xsquared = sum(X.*X);   % this is ||r||^2 after DC is selected
sumXX = sum(xsquared);
I = find(xsquared > (maxError^2));  % for the rest DC is enough
DCrr = sumXX - sum(xsquared(I)); % the errors for the ones where only DC is selected
I = I(randperm(numel(I)));  % permute the elements of I, important only for D0

% to selecte some or all of the basis vectors of the transform
% may be a good idea, i.e. include eye(N) in D0, but we risk that some
% will not be used at all. Here we rather select the first vectors
% initial dictionary, the K random training vectors
D0 = dictnormalize( X(:,I(1:K)) );

java_access;
timestart = now();
jD0 = mpv2.SimpleMatrix(D0);
jDicLea  = mpv2.DictionaryLearning(jD0, 1);
jDicLea.setORMP(int32(maxS), relLim, maxError);
jDicLea.setVerbose(0);

tabIT = zeros((noIT/substep),1);
tabPSNR = zeros((noIT/substep),1);
tabNNZ = zeros((noIT/substep),1);  % number of non-zeros
for subiteration = 1:(noIT/substep)
    %
    disp(' ');
    tic;
    jDicLea.ilsdla( reshape(X(:,I),N*numel(I),1), substep );
    %  use D to find new SNR and number of non-zero weights
    jD =  jDicLea.getDictionary();
    D = reshape(jD.getAll(), N, K);
    jDD = mpv2.SymmetricMatrix(K,K);
    jDD.eqInnerProductMatrix(jD);
    jMP = mpv2.MatchingPursuit(jD,jDD);
    W = zeros(K,L);
    nzW = ones(1,L);  % the DC componet
    for j=I
        sumxx = sum(X(:,j).^2);
        W(:,j) = jMP.vsORMP(X(:,j), int32(maxS), max(maxError/sqrt(sumxx), relLim));
        nzW(j) = 1 + sum(W(:,j) ~= 0);
    end
    R = X(:,I) - D*W(:,I);  % excluding the ones with only DC
    sumRR = DCrr + sum(sum(R.*R));
    newPSNR = 10*log10( ((L*N)*255^2)/sumRR );
    disp([mfile,': ',int2str(subiteration*substep),' iterations:',...
        ' maxError = ',num2str(maxError),...
        ' PSNR = ',num2str(newPSNR),...
        '  sparseness = ',num2str( sum(nzW)/(L*N) ),...
        '  NOF non-zeros is ',int2str(sum(nzW)),'.']);
    timeleft = (ceil(noIT/substep)-subiteration)*(now()-timestart)/subiteration;
    disp(['Time for ',int2str(substep),' iterations is ',num2str(toc),' sec. ', ...
        'Estimated finish time is ',datestr(now()+timeleft)]);
    tabIT(subiteration) = subiteration*substep;
    tabPSNR(subiteration) = newPSNR;
    tabNNZ(subiteration) = sum(nzW);
    %
    % we may want to adjust maxError
    % factor = 1 + 0.1*(1-subiteration/(noIT/substep))*abs(targetPSNR-newSNR);
    factor = 1 + 0.2*((1-subiteration/(noIT/substep))^2)*min(abs(targetPSNR-newPSNR),2);
    if (newPSNR > targetPSNR)
        % we want to select fewer, i.e. increase maxError
        maxError = maxError*factor;
    else
        % we want to select more, i.e. deccrease maxError
        maxError = maxError/factor;
    end
    I = find(xsquared > (maxError^2));  % for the rest DC is enough
    DCrr = sumXX - sum(xsquared(I)); % the errors for the ones where only DC is selected
end
%
jD =  jDicLea.getDictionary();
D = reshape(jD.getAll(), N, K);
dstr = datestr(now());
ResultFile = [mfile,dstr([4:6,1,2,13,14,16,17]),'.mat'];
if (verbose >= 0); disp(['Save results in ',ResultFile]); end;
save(ResultFile, 'D','tabNNZ','tabPSNR','tabIT','ResultFile',...
    'targetPSNR','N','K','L','transform');
%
res = struct('D',D,...
             'tabNNZ',tabNNZ,...
             'tabIT',tabIT,...
             'tabPSNR',tabPSNR,...
             'ResultFile',ResultFile,...
             'targetPSNR',targetPSNR,...
             'N',N,'K',K,'L',L,'transform',transform );


% display some properties for the results
ex31prop(ResultFile);
    
if testLena
    % sparse representation of lena using trained dictionary
    targetPSNRtab = [32, 34, 36, 38];
    r2 = cell(size(targetPSNRtab));
    for i = 1:numel(targetPSNRtab)
        r2{i} = imageapprox(double(imread('lena.bmp'))-128, ...
            'Transform',res.transform, ...
            'Dictionary',res.D, ...
            'targetPSNR',targetPSNRtab(i), ...
            'peak',255, ...
            'delta',0, ...
            'verbose', 1);
    end
    res.r2 = r2;
end

return