function res = ex315(transform, K, targetPSNR, noIT, testLena, verbose)
% ex315           Training of dictionary for images using separable ILS-DLA
%                 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).
% The results are saved in 'ex315mmmddhhmm.mat'.
% 
% use:
%  res = ex315(transform, K, targetPSNR, noIT, testLena, verbose)
%-------------------------------------------------------------------------
% arguments:
%   transform   may be:  'none', 'dct', 'lot', 'elt', 'db79', 'm79'
%               see myim2col.m for more
%   K           a 1x2 vector [K1,K2] giving K = K1*K2 as number of vectors 
%               in dictionary. K1 (for vertical part, Dv) and K2 (for
%               horizontal part, Dh) should be smaller than 40 .
%   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. (or 2 for very verbose)
%-------------------------------------------------------------------------
% example:
%  res = ex315('none', [21,21], 36, 50);         % simple example
%  res = ex315('m79', [21,21], 36, 200, 1, 1);   

%----------------------------------------------------------------------
% 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  06.11.2009  KS: Function made and works
% Ver. 1.1  20.01.2010  KS: use getXfrom8images (not dataXimage)
% Ver. 1.2  09.08.2011  KS: Simplified a little bit
%----------------------------------------------

mfile = 'ex315';

if ((nargin == 1) && strcmpi(transform,'many'))
    K = [21,21];  % gives K=441;
    noit = 2000;
    target = 36;
    for i=1
        ex315( 'm79', K, target, noit);
        ex315('none', K, target, noit);
    end
    res = 'done';
    return;
end

if (nargin < 3)
   error([mfile,': wrong number of arguments, see help.']);
end
if (numel(K) ~= 2)
   error([mfile,': K should be 1x2, see help.']);
end
if (nargin < 4); noIT = 500; end;
if (nargin < 5); testLena = 0; end;
if (nargin < 6); verbose = 0; end;

K1 = min(abs(K(1)),40);
K2 = min(abs(K(2)),40);
K = K1*K2;
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.2*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-S training for images, ',datestr(now()),...
    ', N=64, K=',int2str(K),', L=',int2str(L),', noIT=',int2str(noIT),...
    ', target PSNR = ',num2str(targetPSNR) ]);

% always select the DC vector, thus 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
xnorm = sqrt(xsquared);
sumXX = sum(xsquared);
I = find(xnorm > maxError);  % for the rest DC is enough
DCrr = sumXX - sum(xsquared(I)); % the errors for the ones where only DC is selected

if strcmpi(transform, 'none')
    % initial dictionary, the K random training vectors made for sine
    % D is given by vertical (Dv) and horizontal (Dh) factors
    % NOTE that DC componet is first column of D (and of Dv and Dh)
    Dh = get8x21haar();  % just some extra column vectors to use
    Dv = [get8x21sine(36),Dh(:,2:end)];
else
    Dv = [eye(8), fliplr(get8x21haar()), ...
        toeplitz([1,0.1,0.25,0,0,0,0,0]), eye(8)+0.3*randn(8)];
end
Dv = Dv+0.1*randn(size(Dv));   % add some noise
Dv = Dv.*(ones(8,1)*(1./sqrt(sum(Dv.*Dv))));
Dh = Dv(:,1:K2);
Dv = Dv(:,1:K1);
D4d = zeros(8,8,K1,K2);     % 4-D dictionary, all outer-products
for k1=1:K1
    for k2=1:K2
        D4d(:,:,k1,k2) = Dv(:,k1)*Dh(:,k2)';
    end
end
D = reshape(D4d, N, K);     % the dictionary as a NxK matrix 
% DC-component is number 1, DC is removed from X so the weight for 
% this should always be zero. But the atom can not be removed since we want
% DC-component both in Dv and Dh!
java_access;
timestart = now();
jD = mpv2.SimpleMatrix(D(:,2:end));       % no DC componet needed

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;
    for substepnumber = 1:substep
        % 1. vector selection, D and jD are current dictionary
        jDD = mpv2.SymmetricMatrix.innerProductMatrix(jD);
        jMP = mpv2.MatchingPursuit(jD, jDD);
        W = zeros(K,L);    % include DC here
        for i = I   % this is those where DC is not enough
            W(2:K,i) = jMP.vsORMP(X(:,i), int32(maxS), ...
                max(relLim, maxError/xnorm(i)));
        end
        % 2. Dictionary update Separable frame
        for i=1:12    % approx 0.2 second per iteration
            Dvi = Dv; Dhi = Dh;
            V = Dh*reshape(permute(reshape(W,K1,K2,L),[2,1,3]),K2,K1*L);
            V = reshape(permute(reshape(V,8,K1,L),[2,1,3]),K1,8*L);
            Dv = (reshape(X,8,8*L)*V')/(V*V');
            V = Dv*reshape(W,K1,K2*L);
            V = reshape(permute(reshape(V,8,K2,L),[1,3,2]),8*L,K2);
            Dh = reshape(permute(reshape(X,8,8,L),[2,1,3]),8,8*L)*V/(V'*V);
            % check convergence
            nv = norm(Dvi-Dv);
            nh = norm(Dhi-Dh);
            if verbose > 1
                disp([' iteration ',...
                    int2str((subiteration-1)*substep+substepnumber),...
                    '-',int2str(i),' : ||Dvi-Dv|| = ',num2str(nv,'%10.8f'),...
                    '  ||Dhi-Dh|| = ',num2str(nh,'%10.8f')]);  
            end
            if ((nv+nh)<1e-6); break; end;
        end
        % first is DC, and normalize
        if strcmpi(transform,'none')
            Dv(:,1) = ones(8,1)/sqrt(8);
            Dh(:,1) = ones(8,1)/sqrt(8);
        else
            Dv(:,1) = [1; zeros(7,1)];
            Dh(:,1) = [1; zeros(7,1)];
        end
        if (K1 == K2) && (subiteration < (noIT/(3*substep))) 
            % make Dv and Dh equal (which partly destroy convergence)
            Dv = Dv+Dh; Dh = Dv;
        end
        % this just normalize (which can be compensated for in W)
        Dv = Dv.*(ones(8,1)*(1./sqrt(sum(Dv.*Dv))));
        Dh = Dh.*(ones(8,1)*(1./sqrt(sum(Dh.*Dh))));
        for k1=1:K1
            for k2=1:K2
                D4d(:,:,k1,k2) = Dv(:,k1)*Dh(:,k2)';
            end
        end
        D = reshape(D4d, N, K);     % the dictionary as a matrix
        jD = mpv2.SimpleMatrix(D(:,2:end));    % no DC componet needed
    end
    %  use D to find W and new SNR and number of non-zero weights
    jDD = mpv2.SymmetricMatrix.innerProductMatrix(jD);
    jMP = mpv2.MatchingPursuit(jD, jDD);
    W = zeros(K,L);    % include DC here
    for i = I   % this is those where DC is not enough
        W(2:K,i) = jMP.vsORMP(X(:,i), int32(maxS), ...
            max(relLim, maxError/xnorm(i)));
    end
    nzW = ones(1,L) + sum(W ~= 0);  % the DC componet + W
    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(xnorm > maxError);  % for the rest DC is enough
    DCrr = sumXX - sum(xsquared(I)); % the errors for the ones where only DC is selected
end
%
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','Dv','Dh');
%

res = struct('D',D,'Dv',Dv,'Dh',Dh,...
             'tabNNZ',tabNNZ,...
             'tabIT',tabIT,...
             'tabPSNR',tabPSNR,...
             'ResultFile',ResultFile,...
             'targetPSNR',targetPSNR,...
             'N',N,'K',K,'L',L,'transform',transform );


if isPC()        % make a figure
    h=figure(1);clf;hold on;
    set(h,'Color',0.95*[1,1,1]);
    sc = 0.6;
    k = 1;
    for y0 = 3:(-1):1
        for x0 = 1:7;
            % plot dictionary elements in Dv
            x = x0 + (2:(size(res.Dv,1)+1))/12; 
            h = text(x0+1/12,y0,int2str(k));
            set(h,'HorizontalAlignment','right');
            plot([x(1),x(end)], [y0, y0], 'b-');
            for n=1:size(res.Dv,1)
                plot([x(n), x(n)], [y0, y0+sc*res.Dv(n,k)], 'b-');
                plot(x(n), y0+sc*res.Dv(n,k), 'b.');
            end
            k = k +1;
            if (k>size(res.Dv,2)); break; end;
        end
        if (k>size(res.Dv,2)); break; end;
    end
    title('The vectors in vertical factor of dictionary, Dv.');
    axis off;
    h=figure(2);clf;hold on;
    set(h,'Color',0.95*[1,1,1]);
    k = 1;
    for y0 = 3:(-1):1
        for x0 = 1:7;
            % plot dictionary elements in Dh
            x = x0 + (2:(size(res.Dh,1)+1))/12; 
            h = text(x0+1/12,y0,int2str(k));
            set(h,'HorizontalAlignment','right');
            plot([x(1),x(end)], [y0, y0], 'b-');
            for n=1:size(res.Dh,1)
                plot([x(n), x(n)], [y0, y0+sc*res.Dh(n,k)], 'b-');
                plot(x(n), y0+sc*res.Dh(n,k), 'b.');
            end
            k = k +1;
            if (k>size(res.Dh,2)); break; end;
        end
        if (k>size(res.Dh,2)); break; end;
    end
    title('The vectors in horizontal factor of dictionary, Dh.');
    axis off;
    %
end

% 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

% % test: reshaping of multidimensional matrices
% N1 = 2; N2 = 2; L1 = 5; L2 = 6; K1 = 3; K2 = 3;
% N = N1*N2; L = L1*L2; K = K1*K2;
% A = ((1:(N1*L1))'+randn(N1*L1,1))*(1:(N2*L2));  % image A
% A = A + randn(N1*L1,N2*L2);      % N1*L1 x N2*L2
% disp(A(1:2,1:2));                % upper left corner
% 
% X = reshape(A, N1, L1, N2, L2);  % expanded to all dimensions
% X = permute(X,[1,3,2,4]);        % and ordered to N1xN2xL1xL2
% Xm = reshape(X, N, L);           % as a NxL matrix
% 
% a = [10, 70, 130]*pi/180;
% Dv = [cos(a); sin(a)];           % vertical factor of Dm
% Dh = [cos(a+pi/6); sin(a+pi/6)]; % horizontal factor of Dm
% D = zeros(N1,N2,K1,K2);
% for k1=1:K1
%     for k2=1:K2
%         D(:,:,k1,k2) = Dv(:,k1)*Dh(:,k2)';
%     end
% end
% Dm = reshape(D, N, K);             % the dictionary as a matrix
% k1 = 3; k2 = 1; k = k1+(k2-1)*K1;
% disp([squeeze(D(:,:,k1,k2)), reshape(Dm(:,k),N1,N2)]);  % same atom twice
% 
% % sparse representation
% java_access();
% jD = mpv2.SimpleMatrix(Dm);
% jDD = mpv2.SymmetricMatrix.innerProductMatrix(jD);
% jMP = mpv2.MatchingPursuit(jD, jDD);
% Wm = zeros(K,L);
% for i = 1:L
%     Wm(:,i) = jMP.vsORMP(Xm(:,i), int32(2), 1e-6);
% end
% W = reshape(Wm, K1, K2, L1, L2);  % expanded to all dimensions
% 
% % reconstruction, R, and error, E
% Rm = Dm*Wm;
% Em = Xm-Rm;
% R = zeros(size(X));
% for i1=1:L1
%     for i2=1:L2
%         R(:,:,i1,i2) = Dv*W(:,:,i1,i2)*Dh';
%     end
% end
% E = X-R;
% disp('Signal X and Xm for upper left corner.');
% disp([squeeze(X(:,:,1,1)), reshape(Xm(:,1),N1,N2)]);
% disp('Reconstruction R and Rm for upper left corner.');
% disp([squeeze(R(:,:,1,1)), reshape(Rm(:,1),N1,N2)]);
% disp('Error E and Em for upper left corner.');
% disp([squeeze(E(:,:,1,1)), reshape(Em(:,1),N1,N2)]);
% 
% % set X from Xm : X = reshape(Xm, N1, N2, L1, L2);
% % set Xm from X : Xm = reshape(X, N, L);
