function varargout = saSIDbmp2(y, sid, s, tre, verbose)
% saSIDbmp2       MP Sparse Approximation with Shift Invariant Dictionary 
%                 (SID) and BMP Basic Matching Pursiuit (BMP) is used
%
% There may be 3 to 5 input arguments and 1 to 3 output arguments.
% The SID should be normalized when two or more output arguments are used.
% 
% use:   w = saSIDbmp2(y, sid, s);
%        w = saSIDbmp2(y, sid, s, tre, verbose);  
%        [AtA, Atb] = saSIDbmp2(y, sid, s, tre);  
%        [w, AtA, Atb] = saSIDbmp2(y, sid, s, tre);  
% -------------------------------------------------------------------------
%   y     The (training) signal, size Nx1, or NxL
%   sid   The Shift Invariant Dictionary, as in initSID
%   s     Number of non-zeros to use in each column of y
%   tre   Target realtive error  ||r|| < tre*||y||, 0 <= tre < 1
%         (only implemented in mex-file)
%   w     The sparse coefficients, size Kx1, or KxL
%   AtA   the (new part of) Q-by-Q matrix A'*A, as in LS equation: Ax = b
%   Atb   the (new part of) Q-by-1 vector A'*b, as in LS equation: Ax = b
% -------------------------------------------------------------------------

%--------------------------------------------------------------------------
% 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  30.10.2015  KS: made file (based on ex03.thisBMPst) 
% Ver. 1.1  04.11.2015  KS: may use faster mex-file
% Ver. 1.2  10.11.2015  KS: new version based on saSIDbmp
%                       and use new mex-file (ver. 1.4) 
%--------------------------------------------------------------------------

if ~isstruct(sid) || ~isfield(sid,'Type') || ~strcmp(sid.Type, 'SID')
    disp([mfilename,': Second argument not a SID structure.']);
    if nargout >= 1; varargout{1} = []; end
    if nargout >= 2; varargout{2} = []; end
    if nargout >= 3; varargout{3} = []; end       
    return
end
if nargin < 5
    verbose = false;
end
if ((nargin < 4) || isempty(tre))
    tre = 0;
end

[N,L] = size(y);
M = sid.M;    %  number of atoms in dictionary
Q = sid.Q;    %  total number of free variables, i.e. sum of Qm
Qm = sid.Qm;  %  length of each atom, size Mx1
Pm = sid.Pm;  %  shift for each atom, size Mx1
% Qmax = max(Qm);
Km = floor( (N+Pm-Qm)./Pm );
K = sum(Km);
k2 = cumsum(Km);
k1 = [1; k2(1:(M-1))+1];
if verbose
    fprintf('%s: BMP on y (%i x %i), up to %i nonzeros in each column, tre = %6.4f.\n', ...
        mfilename, N, L, s, tre);
    fprintf('  SID is (%i x %i) %i atoms, %i  free variables.\n', N, K, M, Q);
end

% check normalization
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));
    na = atom'*atom;
    if (abs(na - 1) > 1e-8) 
        wasNormalized = false;
        scaleFactor(m) = 1/sqrt(na);
        sid.freeVar(sid.i1(m):sid.i2(m)) = atom*scaleFactor(m);
        if verbose
            fprintf(' atom %3i has length %5i and norm %6.2f. (normalize it) \n', m, numel(atom), sqrt(na));
        end
    end
end
% now sid is normalized

if ( ((N*K > 1e5) || (tre > 0)) && (exist('saSIDbmp_mex','file') == 3))
    if verbose
        disp('  Now use saSIDbmp_mex. '); 
    end
    % fast (and the same results as other methods)
    if nargout == 1
        w = sparse(saSIDbmp_mex(y, sid.freeVar, sid.Qm, sid.Pm, s, tre)); % ver 1.4
    elseif nargout == 2
        [AtAc, Atb] = saSIDbmp_mex(y, sid.freeVar, sid.Qm, sid.Pm, s, tre); 
    elseif nargout == 3
        [w, AtAc, Atb] = saSIDbmp_mex(y, sid.freeVar, sid.Qm, sid.Pm, s, tre); 
        w = sparse(w);
    end
    if nargout > 1
        AtA = zeros(Q);
        AtA(:,sid.i1) = AtAc;
        AtA(sid.i1,:) = AtAc';
        qq = setdiff(1:Q, sid.i1);
        for q1 = qq
            for q2 = qq
                AtA(q1,q2) = AtA(q1-1,q2-1);
            end
        end
    end
    
else   % do not use mex-file
    if ((N*K) < 1e9)  %  if K = 3N   limit == 1e9 ==> N approx 20000
        if verbose
            fprintf('  %s: Build the dictionary matrix and use it.\n', mfilename);
        end
        D = makeSIDmatrix(sid, N);    % size NxK
        w = sparse(K,L);
        for column = 1:L
            noit = 0;
            r = y(:,column);
            wcol = zeros(K,1);
            nnzwcol = 0;
            while (nnzwcol < s) && (noit < 2*s)
                % Next line gives algorithm O(N*K*s*L)
                % (but lower since sparseness is D exploited)
                % even so, it does a lot of work for each vector to select
                c = D'*r;
                [temp,k] = max(abs(c));
                if wcol(k)
                    wcol(k) = wcol(k) + c(k);
                else
                    wcol(k) = c(k);
                    nnzwcol = nnzwcol + 1;
                end
                r = r - c(k)*D(:,k);
                noit = noit + 1;
            end
            w(:,column) = wcol; %#ok<SPRIX>
            if verbose
                fprintf(' Col %4i: noit = %5i, nnz(w) = %5i, ||y|| = %12.3f and ||r|| = %12.3f.\n', ...
                    column, noit, nnz(wcol), sqrt(y(:,column)'*y(:,column)), sqrt(r'*r) );
            end
        end
    else
        % this method is surprisingly slow and used only for very large
        % dictionaries when mex-file is not available.
        if verbose
            fprintf('  %s: Use dictionary structure (slow). \n', mfilename);
        end
        % Solve this part more like GMP (Global MP) in sparseapprox
        % than thisBMPmat in ex03. Oct 2015 KS, or another way
        w = sparse(K,L);
        for column = 1:L
            noit = 0;
            r = y(:,column);
            c = full(multSIDt(sid, r)); % all inner products
            wcol = zeros(K,1);
            nnzwcol = 0;
            while (nnzwcol < s) && (noit < 2*s)
                % for large K, implemetation using priority queue could be considered
                [temp,k] = max(abs(c));    % in   1,2,...,K
                m = find((k <= k2), 1, 'first');    % which atom is used, in  1,2,...,M
                km = k - k1(m);     % in  0,1,...Km(m)-1
                atom = sid.freeVar(sid.i1(m):sid.i2(m));
                % update coefficient k
                if wcol(k)
                    wcol(k) = wcol(k) + c(k);
                else
                    wcol(k) = c(k);
                    nnzwcol = nnzwcol + 1;
                end
                % update r
                n1 = 1 + km*Pm(m);
                n2 = n1 + Qm(m) - 1;
                r(n1:n2) = r(n1:n2) - c(k)*atom;
                % update elements in c (use only n1 and n2 from above)
                % n1 and n2 may be "anywhere" in 1,2,...,N
                for m = 1:M
                    atom = sid.freeVar(sid.i1(m):sid.i2(m));
                    % find atoms in this block related to n1:n2
                    % k3, k4  in  0,1,...,Km(m)-1
                    k3 = max(0, ceil((n1-Qm(m))/Pm(m)));
                    k4 = min(Km(m)-1, floor(n2/Pm(m))-1);
                    n3 = k3*Pm(m) + 1;
                    n4 = n3 + Qm(m) - 1;  % for atom starting in n3
                    % inner products k, from k3 to k4 in block m, should be updated
                    %                 for k = (k3+k1(m)):(k4+k1(m))
                    %                     c(k) = atom'*r(n3:n4);
                    %                     n3 = n3 + Pm(m);
                    %                     n4 = n4 + Pm(m);
                    %                 end
                    % or
                    nok = k4-k3+1;
                    non = n4-n3+1;
                    n = repmat((n3:n4)',1,nok) + repmat(0:Pm(m):((nok-1)*Pm(m)),non,1);
                    c((k3+k1(m)):(k4+k1(m))) = atom'*reshape(r(n(:)),non,nok);
                end
                noit = noit + 1;
            end
            w(:,column) = wcol; %#ok<SPRIX>
        end
    end
    
    % here, the algorithm has found sparse w (KxL)
    if nargout > 1
        % now finw AtA and Atb from w, and k1, k2, sid
        if L > 3  % use compact form (usually slower)
            AtAc = zeros(Q,M);
            for column = 1:L
                AtAc = AtAc + thisAtAcompact(w(:,column), sid, N);
            end
            % compact form to full form
            AtA = zeros(Q);
            AtA(:,sid.i1) = AtAc;
            AtA(sid.i1,:) = AtAc';
            qq = setdiff(1:Q, sid.i1);
            for q1 = qq
                for q2 = qq
                    AtA(q1,q2) = AtA(q1-1,q2-1);
                end
            end
        else
            AtA = zeros(Q);
            for column = 1:L
                AtA = AtA + thisAtA(w(:,column), sid, N);
            end
        end
        Atb = zeros(Q,1);
        for column = 1:L
            Atb = Atb + thisAtb(y(:,column),w(:,column), sid);
        end
    end
    
end

if ~wasNormalized    % sid is not normalized
    for m = 1:M
        w(k1(m):k2(m),:) = w(k1(m):k2(m),:)*scaleFactor(m); %#ok<SPRIX>
    end
    if verbose
        disp(['  ',mfilename,': The coefficients are scaled to match the not normalized SID.']); 
    end
    if (nargout > 1)
        disp(['  ',mfilename,': AtA and Atb are NOT scaled to match the not normalized SID.']); 
    end
end

%
if nargout == 1
    varargout{1} = w;
elseif nargout == 2
    varargout{1} = AtA;
    varargout{2} = Atb;
elseif nargout == 3
    varargout{1} = w;
    varargout{2} = AtA;
    varargout{3} = Atb;
end

end


%% subfunctions (some from ex03)

% Map to Aij from wi and wj for Shift-Invariant G (implicit given) Sparse Xcorr
% Map to Aij from wi and wj for Shift-Invariant G (implicit given) Xcorr
function A = thisMapAijXc(wi, wj, sz)
    r = xcorr(full(wi),full(wj),max(sz)-1);
    m = (numel(r)+1)/2;   % middle
    i1 = m - sz(1) + 1;
    i2 = m + sz(2) - 1;
    A = toeplitz(r(m:(-1):i1),r(m:i2));
end

% function A = thisMapAijSp(wi, wj, sz)
%     r = thisXcorrSp(wi,wj,sz);
%     A = toeplitz(r(sz(1):(-1):1),r(sz(1):end));
% end
% 
% function r = thisXcorrSp(x,y,sz)
%     ix = find(x);  
%     iy = find(y);  
%     r = zeros(sum(sz)-1,1);
%     for m = (1-sz(2)):(sz(1)-1)
%         i = intersect(ix+m,iy);
%         if numel(i)
%             r(sz(1)-m) = x(i-m)'*y(i);
%         end
%     end
% end    

% Map to A from w for Shift Invariant Dictionary
function A = thisAtA(w, sid, N)
    M = sid.M;    %  number of atoms in dictionary
    Qm = sid.Qm;  %  length of each atom, size Mx1
    Q = sid.Q;
    Pm = sid.Pm;  %  shift for each atom, size Mx1
    Km = floor( (N+Pm-Qm)./Pm );
    K = sum(Km);
    if K == numel(w)
        k2 = cumsum(Km);          % last element of each segment i w 
        k1 = [1; k2(1:(M-1))+1];  % first element of each segment i w 
        A = zeros(Q);             % QxQ
        for i = 1:M
            wi = w(k1(i):k2(i));
            if Pm(i) > 1
                temp = zeros(Pm(i),1); temp(1) = 1;
                wi = kron(wi,temp);
            end
            ii = sid.i1(i):sid.i2(i);   % indexes for atom i
            for j = i:M
                wj = w(k1(j):k2(j));
                if Pm(j) > 1
                    temp = zeros(Pm(j),1); temp(1) = 1;
                    wj = kron(wj,temp);
                end
                jj = sid.i1(j):sid.i2(j);   % indexes for atom j
                % may try different alternatives.
                A(ii,jj) = thisMapAijXc(wi, wj, [Qm(i),Qm(j)]);   
                % A(ii,jj) = thisMapAijSp(wi, wj, [Qm(i),Qm(j)]);   
                if j > i
                    A(jj,ii) = A(ii,jj)';
                end
            end
        end
    else
        disp('Size of w does not match sum(K).');
        A = [];
    end
end

% Map to b from x and w for Shift-Invariant G (implicit given)
% Q has lengths of each atom in Shift-Invariant dictionary
function b = thisAtb(y, w, sid)
    N = numel(y);
    M = sid.M;    %  number of atoms in dictionary
    Qm = sid.Qm;  %  length of each atom, size Mx1
    Q = sid.Q;
    Pm = sid.Pm;  %  shift for each atom, size Mx1
    Km = floor( (N+Pm-Qm)./Pm );
    K = sum(Km);
    if K == numel(w)
        k2 = cumsum(Km);         % last element in each block of w
        k1 = [1; k2(1:(M-1))+1]; % first element in each block of w
        b = zeros(Q,1);
        for i = 1:M
            wi = w(k1(i):k2(i));
            if Pm(i) > 1
                temp = zeros(Pm(i),1); temp(1) = 1;
                wi = kron(wi,temp);
            end
            r = xcorr(full(wi), y, Qm(i)-1);
            b(sid.i1(i):sid.i2(i)) = r(Qm(i):(-1):1);
        end
        
    else
        disp('Size of w does not match sum(K).');
        b = [];
    end   
end


% Map to A from w for Shift Invariant Dictionary, compact form
function A = thisAtAcompact(w, sid, N)
    M = sid.M;    %  number of atoms in dictionary
    Qm = sid.Qm;  %  length of each atom, size Mx1
    Q = sid.Q;
    Pm = sid.Pm;  %  shift for each atom, size Mx1
    Km = floor( (N+Pm-Qm)./Pm );
    K = sum(Km);
    if K == numel(w)
        k2 = cumsum(Km);          % last element of each segment i w 
        k1 = [1; k2(1:(M-1))+1];  % first element of each segment i w 
        A = zeros(Q,M);           % QxM
        for i = 1:M
            wi = w(k1(i):k2(i));
            wii = find(wi);
            wiv = wi(wii);
            if Pm(i) > 1
                wii = Pm(i)*(wii-1) + 1;
            end
            ii = sid.i1(i):sid.i2(i);   % indexes for atom i
            for j = i:M
                wj = w(k1(j):k2(j));
                wji = find(wj);
                wjv = wj(wji);
                if Pm(j) > 1
                    wji = Pm(j)*(wji-1) + 1;
                end
                jj = sid.i1(j):sid.i2(j);   % indexes for atom j
                %
                r = thisXcorrSparse(wiv,wii,wjv,wji,[Qm(i),Qm(j)]);
                A(jj,i) = r(Qm(i):end)';
                if j > i
                    A(ii,j) = r(Qm(i):(-1):1)';
                end
            end
        end
    else
        disp('Size of w does not match sum(K).');
        A = [];
    end
end

function r = thisXcorrSparse(xv,xi,yv,yi,sz)
    % x = rand(200,1); x(x<0.8) = 0; xi = find(x); xv = x(xi); xs = sparse(x);
    % y = rand(200,1); y(y<0.8) = 0; yi = find(y); yv = y(yi); ys = sparse(y);
    % sz = [3,4]; rtrue = xcorr(x,y,3); disp(rtrue'); disp([x, y]);
    r = zeros(sum(sz)-1,1);   % 1-sz(1),...,0,..,sz(2)-1
    n1 = numel(xv);  
    n2 = numel(yv);
    if (n1 == 0) || (n2 == 0)
        return;
    end
    offset = sz(1);
    limit1 = 1-sz(1);
    limit2 = sz(2)-1;
    i1 = 1; 
    i2 = 1;
    i3 = 1;
    while (i1 <= n1)
        m = xi(i1) - yi(i2);
        % disp([' i1=',int2str(i1),', i2=',int2str(i2),' m=',int2str(m)]);
        if (m < limit1)  
            i1 = i1 + 1;
            i2 = i3;
        elseif (m > limit2)
            i3 = i3 + 1;
            i2 = i3;
            if (i2 > n2); break; end;
        else
            rm = m + offset;
            r(rm) = r(rm) + xv(i1)*yv(i2);
            i2 = i2 + 1;
            if (i2 > n2); 
                i1 = i1 + 1;
                i2 = i3;
            end
        end
    end
end    
