function varargout = saSID2Domp(y, sid, s, SAmet, bSize, verbose)
% saSID2Domp      OMP Sparse Approximation with 2D SID
%                 SAmet or Orthogonal Matching Pursiuit (OMP) is used.
% The large 2D SID dictionary is divided into blocks and OMP is used on
% corresponding signal blocks. For this task sparseapprox is used.
%
% use:   w = saSID2Domp(y, sid, s);
%        w = saSID2Domp(y, sid, w0);
%        w = saSID2Domp(y, sid, s, SAmet, bSize, verbose);  
%        [w, snr] = saSID2Domp(y, sid, s, SAmet, bSize, verbose);  
% -------------------------------------------------------------------------
%   w     The returned sparse coefficients, size Kx1
%   y     The (test/training) signal, size: Nv x Nh
%   sid   The 2D Shift Invariant FSD, as in initSID2D
%   s     Number of non-zeros to use for y
%   w0    Initial coefficients, may be given instead of s, size Kx1
%   SAmet Sparse Approximatio method as used in sparseapprox.m
%         can be 'mexOMP' (default) 
%         'PSnn', where nn is 10, 20, 50 or 100, is also possible
%   bSize block size, default [12,48; 48,12; 24,24];
%         an alternative can be:   [16,64; 64,16; 32,32];   
%   verbose    0: only warnings (default), 1: some info,  2: more,  3: much
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Made by Karl Skretting, University of Stavanger (UiS), 
% Department of Electrical Engineering and Computer Science (IDE)
% Mail:  karl.skretting@uis.no   Homepage:  http://www.ux.uis.no/~karlsk/
% 
% HISTORY:  dd.mm.yyyy
% Ver. 1.0  28.06.2016  KS: started to make file
% Ver. 1.1  25.08.2016  It is possible to use different methods
% -------------------------------------------------------------------------

%% get input (or use default)
if (~isstruct(sid) || ~isfield(sid,'Type') || ~strcmp(sid.Type, 'SID2D') || ...
    ~isfield(sid,'M') || ~isfield(sid,'Qv') || ~isfield(sid,'Qh')  )
    disp([mfilename,': Second argument not a 2D SID structure.']);
    if (nargout >= 1); varargout{1} = []; end;
    if (nargout >= 2); varargout{2} = []; end;
    return
end
if (nargin < 2)
    disp([mfilename,': There should be at least two input arguments.']);
    if (nargout >= 1); varargout{1} = []; end;
    if (nargout >= 2); varargout{2} = []; end;
    return
end
[Nv,Nh] = size(y);
if (nargin < 3) || isempty(s)
    s = floor(0.1*Nv*Nh);
end
if (nargin < 4) || isempty(SAmet)
    SAmet = 'mexOMP';
end
if (numel(SAmet) > 2) && strcmpi(SAmet(1:2),'PS')
    nComb = eval(SAmet(3:end));
end
if (nargin < 5) || isempty(bSize)
    bSize = [ min(12,Nv), min(48,Nh); ...
              min(48,Nv), min(12,Nh); ... 
              min(24,Nv), min(24,Nh) ];   
end
if nargin < 6
    verbose = 1;
end

%% check input and display information
if verbose >= 1
    fprintf('\n%s: %s on y (%i x %i) using sid of %i atoms.\n', mfilename, SAmet, Nv, Nh, sid.M);
end

if verbose >= 2
    fprintf('Image has size %ix%i. Do SA using %i blocks.\n', Nv, Nh, size(bSize,1));
    for i = 1:size(bSize,1)
        fprintf('Block %i has size %ix%i and divides image into %ix%i = %i blocks.\n', ...
            i, bSize(i,1:2), floor(Nv/bSize(i,1)), floor(Nh/bSize(i,2)), floor(Nv/bSize(i,1))*floor(Nh/bSize(i,2)));
    end
end

M = sid.M;    %  number of atoms in dictionary
PvLCM = thisLCM(sid.Pv);
PhLCM = thisLCM(sid.Ph);
if verbose >= 2
    fprintf('   Shift-Invariant-Dictionary (SID) has %i atoms and %i free variables.\n', M, sid.Q);
    fprintf(['SID atom length: Qv = [%3i',char(kron(ones(1,M-1),',%3i')),']\n'], sid.Qv);
    fprintf(['SID atom length: Qh = [%3i',char(kron(ones(1,M-1),',%3i')),']\n'], sid.Qh);
    fprintf(['SID atom shift:  Pv = [%3i',char(kron(ones(1,M-1),',%3i')),'], LCM = %2i.\n'], sid.Pv, PvLCM);
    fprintf(['SID atom shift:  Ph = [%3i',char(kron(ones(1,M-1),',%3i')),'], LCM = %2i.\n'], sid.Ph, PhLCM);
end

% find how coefficients are organized for SID expanded to
im = thisKfun(sid, Nv, Nh, (verbose >= 2)); 

% check normalization for SID
wasNormalized = true;       % if input sid was normalized or not
scaleFactor = ones(M,1);
for m = 1:M       % in this loop: c = D'*r 
    atom = sid.freeVar(sid.i1(m):sid.i2(m));
    noa = atom'*atom;   % norm of atom
    if (abs(noa - 1) > 1e-8) 
        wasNormalized = false;
        scaleFactor(m) = 1/sqrt(noa);
        sid.freeVar(sid.i1(m):sid.i2(m)) = atom*scaleFactor(m);
        if verbose >= 2
            fprintf(' atom %3i has length %5i and norm %6.2f. (normalize it) \n', m, numel(atom), sqrt(noa));
        end
    end
end
if (verbose >= 1) && ~wasNormalized
    fprintf(' Dictionary was NOT normalized initially.\n');
end
% now sid is normalized

%
if (Nv*Nh) <= 600
    fprintf('Image is small, so %s is done on the whole image.\n', SAmet);
    Dim = thisDblk(sid, im);
    if numel(s) == im.K;
        s = nnz(s);
    else
        s = s(1);
    end
    if exist('nComb','var')
        w = sparseapprox(y(:), Dim, 'javaPS', 'nComb',nComb, 'tnz',s);
    else
        w = sparseapprox(y(:), Dim, SAmet, 'tnz',s);
    end
    if (nargout >= 1); varargout{1} = w; end;
    if (nargout >= 2); varargout{2} = []; end;
    return;
end

% check sparseness s (or input coefficients w0)
if (numel(s)  == im.K)
    w = reshape(full(s), im.K, 1);
    if (verbose >= 1)
        fprintf('Use input coefficients w of size %ix%i and %i non-zeros.\n', size(w), nnz(w) );
    end
else
    if (numel(s) > 1) 
        s = s(1);
    end
    bv = min(ceil(max(max(sid.Qv),16)/PvLCM)*PvLCM, Nv);
    bh = min(ceil(max(max(sid.Qh),16)/PhLCM)*PhLCM, Nh);
    if (verbose >= 1)
        fprintf('\nTRY to find some intial coefficients, (use %ix%i block).\n', bv,bh);
    end
    w = zeros(im.K,1);
    blk = thisKfun(sid, bv, bh, false);
    if (verbose >= 1)
        fprintf('This SID2D on block of size %ix%i gives N = %i, K = %i.\n',... 
            blk.Nv, blk.Nh, blk.N, blk.K );
    end
    %
    Lv = floor(im.Nv/blk.Nv);
    Lh = floor(im.Nh/blk.Nh);
    L = Lv*Lh;
    if (verbose >= 1)
        fprintf('This divides image into (%ix%i) %i non-overlapping blocks.\n', Lv, Lh, L);
    end
    Dblk = thisDblk(sid, blk);
    if ((Lv*blk.Nv == im.Nv) && (Lh*blk.Nh == im.Nh))
        X = im2col(y, [blk.Nv, blk.Nh], 'distinct');
    else
        X = im2col(y(1:(Lv*blk.Nv),1:(Lh*blk.Nh)), [blk.Nv, blk.Nh], 'distinct');
    end
    % uniform
    % sL = floor(s/L)*ones(1,L);  
    % alternative
    X2 = sqrt(sum(X.*X));   sL = floor( X2*s/sum(X2) );  % X is signal
    a = min(L,s-sum(sL)); if (a > 0); sL(1:a) = sL(1:a)+1; end;
    sL(sL > 0.2*blk.N) = floor(0.2*blk.N);
    a = min(L,s-sum(sL)); if (a > 0); sL(1:a) = sL(1:a)+1; end;
    a = min(L,s-sum(sL)); if (a > 0); sL(1:a) = sL(1:a)+1; end;
    if (verbose >= 1)
        fprintf('Total number of non-zero coefficients should be  s = %i, it is sum(sL) = %i.\n', s, sum(sL));
    end
    if exist('nComb','var')
        w = sparseapprox(X, Dblk, 'javaPS', 'nComb',nComb, 'tnz',sL);
    else
        W = sparseapprox(X, Dblk, SAmet, 'tnz',sL);
    end
    i = 1;
    for olc = 0:blk.Nh:(blk.Nh*(Lh-1))
        for our = 0:blk.Nv:(blk.Nv*(Lv-1))
            w0 = w(thisMap(blk, im, our, olc));
            % update only if there is improvement
            if (sum((X(:,i) - Dblk*W(:,i)).^2) < sum((X(:,i) - Dblk*w0).^2)) 
                w(thisMap(blk, im, our, olc)) = W(:,i);
            end
            i = i+1;
        end
    end
    % and (just for debug) find and display SNR
    R = X - Dblk*W;
    X2 = sum(sum(X.^2));
    R2 = sum(sum(R.^2));
    snr = 10*log10(X2/R2);
    if (verbose >= 1)
        fprintf('  sum X squared = %12e, sum R (=X-XR) squared = %12e,  SNR = %5.2f.\n', X2, R2, snr);
    end
    % check if w (and multSID2D.m) are matched
    yr = multSID2D(sid, w, Nv, Nh);
    R = y - yr;
    X2 = sum(sum(y.^2));
    R2 = sum(sum(R.^2));
    snr = 10*log10(X2/R2);
    if (verbose >= 1)
        fprintf('  sum y squared = %12e, sum R (=y-yr) squared = %12e,  SNR = %5.2f.\n', X2, R2, snr);
    end
    clear X W R yr
end

%% now we should have initial coefficients  w of size im.K x 1 
%  check atom in block, and atom in image, and mapping

y2 = sum(sum(y.^2));
if (verbose >= 2)
    fprintf('\nPvLCM = %i, PhLCM = %i\n', PvLCM, PhLCM);
end
nofIterations = size(bSize,1);
for iteration = 1:nofIterations    % do some iterations for image
    bno = rem(iteration-1, size(bSize,1)) + 1;   % i.e bno = iteration
    blk = thisKfun(sid, bSize(bno,1), bSize(bno,2), false);
    Dblk = thisDblk(sid, blk);
    % select offset (vertical and horizontal) with PxLCM as factors up to block size.
    % this is offset for first block, then as many blocks as possible are
    % extracted from the image. Atoms covering more than one block can not
    % be changed.
    ov = floor(rand(1)*blk.Nv/PvLCM)*PvLCM; 
    oh = floor(rand(1)*blk.Nh/PhLCM)*PhLCM; 
    if (verbose >= 2)
        fprintf('Iteration %i use block of size %ix%i (%i), ov=%i, oh=%i\n', iteration, blk.Nv, blk.Nh, blk.N, ov, oh);
    end
    % Total number of blocks is Lv*Lh
    Lv = floor((im.Nv-ov)/blk.Nv);
    Lh = floor((im.Nh-oh)/blk.Nh);
    L = Lv*Lh;
    if (verbose >= 2)
        fprintf('This divides image into (%ix%i) %i non-overlapping blocks.\n', Lv, Lh, L);
    end
    w0 = w;     % keep the coefficients used last time
    sL = zeros(1,L);   % number of non-zero atoms within each block
    i = 1;
    for olc = oh + (0:(Lh-1))*blk.Nh        % number of columns left of current block
        for our = ov + (0:(Lv-1))*blk.Nv    % number of rows above of current block
            map = thisMap(blk, im, our, olc);
            sL(i) = nnz(w(map));
            w(map) = 0;
            i = i+1;
        end
    end
    yr = y - multSID2D(sid, w, Nv, Nh);    % subtract atoms that are in several blocks (and non-zero coefficient)
    X = im2col(yr((ov+1):(ov+blk.Nv*Lv),(oh+1):(oh+blk.Nh*Lh)), [blk.Nv, blk.Nh], 'distinct');
    if (verbose >= 2)
        fprintf('Use yr(%i:%i,%i:%i) to form X of size %ix%i.\n', ov+1,ov+blk.Nv*Lv,oh+1,oh+blk.Nh*Lh,size(X));
    end
    if exist('nComb','var')
        W = sparseapprox(X, Dblk, 'javaPS', 'nComb',nComb, 'tnz',sL);
    else
        W = sparseapprox(X, Dblk, SAmet, 'tnz',sL);
    end
    i = 1;
    for olc = oh + (0:(Lh-1))*blk.Nh        
        for our = ov + (0:(Lv-1))*blk.Nv
            wi = w0(thisMap(blk, im, our, olc));  
            % update only if there is improvement
            if (sum((X(:,i) - Dblk*W(:,i)).^2) < sum((X(:,i) - Dblk*wi).^2)) 
                w(thisMap(blk, im, our, olc)) = W(:,i);
            else
                w(thisMap(blk, im, our, olc)) = wi;
            end
            i = i+1;
        end
    end
    yr = multSID2D(sid, w, Nv, Nh);
    R = y - yr;
    R2 = sum(sum(R.^2));
    snr = 10*log10(y2/R2);
    if (verbose >= 2) || ((verbose >= 1) && (iteration == nofIterations))
        fprintf('  It. %3i: sum y squared = %12e, sum R (=y-yr) squared = %12e,  SNR = %5.2f.\n', iteration, y2, R2, snr);
    end
end

if (nargout >= 1); varargout{1} = sparse(w); end;
if (nargout >= 2); varargout{2} = snr; end;

end

  
%% subfunctions
function y = thisLCM(x)
% x may be a vector of several elements (positive integers)
% matlab lcm function allows only two input arguments, i.e it works
% elementwise on vectors and matrix inputs!
M = numel(x);
if (M == 1) || (min(x) == max(x))
    y = x(1);
else
    y = lcm(x(1), x(2));
    for m = 3:M
        y = lcm(y, x(m));
    end
end
end

function out = thisKfun(sid, Nv, Nh, verbose)
% find how coefficients are organized
Kv = ceil( (Nv+1-sid.Qv)./sid.Pv );
Kh = ceil( (Nh+1-sid.Qh)./sid.Ph );
Km = Kv.*Kh;
K = sum(Km);
k2 = cumsum(Km);
k1 = [1, reshape(k2(1:(sid.M-1)),1,sid.M-1)+1];
out = struct('Qv',sid.Qv, 'Qh',sid.Qh, ...
             'Pv',sid.Pv, 'Ph',sid.Ph, ...
             'N',Nv*Nh, 'K',K, 'M',sid.M, ...
             'Nv',Nv,     'Nh',Nh, ...
             'Kv',Kv,     'Kh',Kh,    'Km',Km, ...
             'k1',k1,     'k2',k2 );             
if verbose
    fprintf('\nExpand SID2D to image/block of size %ix%i has %i atoms.\n', Nv, Nh, K);
    for m = 1:out.M
        fprintf(' Atom %2i use coefficients %7i to %7i (%ix%i).\n', ...
            m, k1(m), k2(m), Kv(m), Kh(m) );
    end
end        
end

% this map coefficients from block to image using SID2D structure
% Dim is the huge expanded dictionary for image y (if y is 512x512, size 262144x1184035)
% Dblk is the smaller expanded dictionary for a block blk (if blk is 12x48, size 576x1958)
% For image   y:  y   = Dim*w + Rim,      Rim is reconstruction error of image
% For block blk:  blk = Dblk*wblk + Rblk, Rblk is reconstruction error of block. 
% where blk is a patch of the image:  blk = y( (our+1):(our+blk.Nv), (olc+1):(olc+blk.Nh) )
%       our is offset upper row     in  0:PvLCM:(im.Nv-im.Qv(m))
%       olc is offset left column   in  0:PhLCM:(im.Nh-im.Qh(m))
% The map is now such that:   w(map) = wblk
function map = thisMap(blk, im, our, olc)
map = zeros(blk.K,1);
for m=1:blk.M
    a = im.Kv(m)-blk.Kv(m);
    b = our/blk.Pv(m) + (olc/blk.Ph(m))*im.Kv(m) + im.k1(m);
    map(blk.k1(m):blk.k2(m)) = (b:(b+blk.Km(m)-1)) + ...
        kron( 0:a:(a*(blk.Kh(m)-1)), ones(1,blk.Kv(m)) );
end
end

function Dblk = thisDblk(sid, blk)
Dblk = zeros(blk.N, blk.K);
k = 1;
for m=1:blk.M
    atom = sid.freeVar(sid.i1(m):sid.i2(m));
    b = blk.Nv - blk.Qv(m);
    idx = (1:(blk.Qv(m)*blk.Qh(m))) + kron( 0:b:(b*(blk.Qh(m)-1)), ones(1,blk.Qv(m)) );
    % idx is indexes (in block) for atom in upper left position of block
    % reshape(idx, blk.Qv(m), blk.Qh(m))
    % This atom is shifted to all possible positions completly inside block
    for j = 0:blk.Ph(m):(blk.Nh-blk.Qh(m)); 
        for i = 0:blk.Pv(m):(blk.Nv-blk.Qv(m))
            b = j*blk.Nv + i; 
            Dblk(idx+b,k) = atom;
            k = k+1;
        end
    end
end

end

