% Perform orthogonal matching pursuit in order to find the representation of a number of samples (X) in terms of a 
% given (overcomplete) basis (C)
%
%function [P, R] = momp(X, C, k, noise_level);
%
%Input:
% X : samples to be encoded (one sample per column)
% C : basis that is used to encode the samples in X (one basis element per column)
% k : maximum number of OMP iterations (i.e., maximum number of non-zero coefficients per sample in P)
%     if noise_level does not equal zero, this is set to dim(X)
% noise_level : perform OMP iterations until the norm of the residual drops below noise_level (default : 0)
%
%Output:
% P  : the coefficients (one coefficient vector per column)
% R  : the residual of each sample
%
%Note: All columns of C have to be set to unit norm
%      X can be obtained from P,R and C by: X = C*P + R
%
%
% Copyright (C) 2009 Kai Labusch
%
% 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 2
% of the License, or (at your option) any later version.
%
function [P, R] = momp(X, C, k, noise_level)

C = C';

if ~exist('noise_level','var'),
    noise_level = 0;
else 
    k=size(X,1);
end

% check if all rows of C have unit norm
n = sqrt(sum(C.^2, 2));
if ~any(n ~= 1) > 0,
    disp('columns of C do not have unit norm ...');
    C  = C ./ (n * ones(1,size(C,2)));
end

P = zeros(size(C, 1), size(X,2));
R = zeros(size(X));

candidates = 1:size(C, 1);

for i=1:size(X,2),
    [R(:,i), P(:,i)] = omp_rec(X(:,i), C, candidates, k, noise_level);
end

return;

function [r_final, p, next_win] = omp_rec(r, C, candidates, k, noise_level)

y = C(candidates,:)*r; % find current winning basis element
[t, inds] = sort(-(y.^2));

y_win  = y(inds(1));
win    = candidates(inds(1));

winner = C(win,:);

r = r - y_win*winner'; % remove projection to winning basis element from sample

% remove projection to winning basis element from C
p_to_win = C*winner';    
C  = C - diag(p_to_win)*repmat(winner, size(C,1), 1);

candidates = sort(candidates(inds(2:end))); % remove winner from list of candidates

if (k > 1) && (norm(r) > noise_level),
    % set norm of remaining rows of C to 1
    n                     = zeros(1,size(C,1));
    n(candidates)         = sqrt(sum(C(candidates,:).^2, 2));

    C(candidates,:) = diag(1./n(candidates))*C(candidates,:);
    
    % next level of recursion
    [r_final, p, next_win] = omp_rec(r, C, candidates, k-1, noise_level);

    for i=1:length(next_win),
        p(win) = p(win) - p(next_win(i))*p_to_win(next_win(i))/n(next_win(i));
        p(next_win(i)) = p(next_win(i))/n(next_win(i));
    end
    p(win) = p(win) + y_win;
    next_win = [ next_win , win];
else % stop recursion after k coefficients have been used ...
    p        = zeros(size(C,1), 1);
    p(win)   = y_win;
    next_win = win;
    r_final  = r;
end

return;