function [A_est,B_est,S_est, E] = RoneSp_ADMM(D,p,num_sub,omega)
% % % % % % % % % % % % % % % %
% This code handles the following optimization problem 
%  min_{A_est,B_est,S_est} |D - A_est - B_est*S_est|_{omega} + |S_est|_1 
%  s.t. A is subject-wise rank-one 
% ========== INPUT =======
% D is the data matrix, whose columns correspond to each image. 
%    For instance, D is a 4800x1000 matrix for 1000 face images of size 80x60 pixel dimension  
%    (This code assumes that each subject have the same number of training
%    samples, and the training labels for D should be arranged in this way [1 1 1... 2 2 2... ])
% p is the number of atoms to be learned in B_est
% num_sub is the number of subjects in D; 
% considering presence of occlusion (omega=1) or noise (omega=2)
% ========== OUTPUT ======
% A_est is the rank-one identity matrix; in general, it denotes the mean
% face of each subject
% B_est is the dictionary
% S_est is the sparse coefficients
% E is the error (sparse for omega=1; dense for omega=2)
%**************************************************************************
% Written by Tsung-Han Chan (thchan@ieee.org), Nov. 12, 2013
%

D = bsxfun(@times, D, 1./sqrt(sum(D.^2))); % Data normalization

[m,n] = size(D);
q = n/num_sub; % number of images per subject
Maxiter = 2000;
OPTS.disp = 0;

A_est = zeros(m,num_sub);
if omega == 1
    mu = 3000/svds(D,1);
    E = zeros(m,n);
    Y = zeros(m,n);
    error = 1; 
    iter_count = 0;

    while error>1e-5 && iter_count<Maxiter
          % step 1
          template = D - E + 1/mu*Y;
          for i=1:num_sub
              A_est(:,i) = mean(template(:,(i-1)*q+1:i*q),2);
          end
          U = D - kron(A_est,ones(1,q)) - E + 1/mu*Y;
          [C D] = eigs(U*U',p,'LM',OPTS);
          X = C'*U; 

          % step 2
          G = D - kron(A_est,ones(1,q)) - C*X + 1/mu*Y;
          E = sign(G).*max(abs(G)-1/mu,0);

          % step 3
          Y = Y + mu*(D - kron(A_est,ones(1,q)) - C*X - E);

          mu = min(1e6,mu*1.01);
          iter_count = iter_count + 1;

          obj = sum(sum(abs(E)));
          error = max(max(abs(D - kron(A_est,ones(1,q)) - C*X - E)));
    end
    
elseif omega == 2 
    U = [];
    for i=1:num_sub
        A_est(:,i) = mean(D(:,(i-1)*q+1:i*q),2);
        U = [U D(:,(i-1)*q+1:i*q)-A_est(:,i)*ones(1,q)];
    end
    [C D] = eigs(U*U',p,'LM',OPTS);
    X = C'*U;  
    E = U - C*X; 
end
display('*** Finding principal subspace accomplished ***')


%=========== dictionary learning ====================
num_cand = 10;            
corr = zeros(num_cand,1);
Wp = cell(num_cand,1);
Bias_p = cell(num_cand,1);
cardinality = zeros(num_cand,1);
            
W = []; 
T = kron(eye(num_sub),ones(1,q))';
Low_rankX = zeros(p,num_sub);
Pw_per = eye(p);
for j=1:p
    %Pbasis_Wrowsj = [p j]
    display(['*** Identifying the atom ' num2str(j) '/' num2str(p) ' ***'])
    theta = rand(n,1); 
    theta = theta/sum(theta);
    obj_old = 1e13; 
    obj_diff = 10;
    
    [w] = ADMMLP(X, theta, T, Pw_per, p);
    %[w] = CVXLP(X, theta, T, Pw_per, num_sub, p);
                    
    while obj_diff>1e-6
        [val ind] = sort(w'*Pw_per*X,'descend');
        for i=1:num_cand
            theta = zeros(n,1);
            theta(ind(i)) = 1;
            
            [w bias] = ADMMLP(X, theta, T, Pw_per, p);
            %[w bias] = CVXLP(X, theta, T, Pw_per, num_sub, p);
            
            Wp{i} = w;
            Bias_p{i} = bias; 
            cardinality(i) = sum(abs(w'*X-kron(bias',ones(1,q)))>1e-6);
            corr(i) = norm(Pw_per*w/norm(w));
        end
        cardinality(cardinality==0) = 1e13;
        val_card = min(cardinality); 
        index = find(cardinality==val_card);
        [vall ind_corr] = max(corr(index));
        ind = index(ind_corr); 
                    
        obj_diff = (obj_old-val_card)/obj_old;
        if obj_diff>=0
             w_new = Wp{ind};
             bias_new = Bias_p{ind};  
             obj_old = val_card;
        end
    end
    W = [W w_new];
    Low_rankX(j,:) = bias_new';
    Pw_per = eye(p)- W*inv(W'*W)*W';
end
W = W';

% updating the variables
B_est = C*inv(W);
B_scale = sqrt(sum(B_est.^2));
B_est = B_est./(ones(m,1)*B_scale);
S_est = W*X - kron(Low_rankX,ones(1,q));
S_est = diag(B_scale)*S_est; 
A_est = A_est + C*pinv(W)*Low_rankX;



%% Untils

function [w, bias] = ADMMLP(X, theta, TT, Pw_per, p)
% minimize(norm(X'*w-T*bias,1));
% s.t. w'*Pw_per*X*theta == 1;

w = zeros(p,1);
Tt = zeros(p,p-1);

b = Pw_per*X*theta;
[valb indb] = max(abs(b));
non_indb = find(1:p~=indb);
                    
Tt(non_indb,:) = eye(p-1); 
Tt(indb,:) = -b(non_indb)'/b(indb); 
                                       
T = [X'*Tt -TT];
t = X(indb,:)'/b(indb); 
total = ADMM_LP(T,t);         

w(non_indb) = total(1:p-1); 
w(indb) = (1-b(non_indb)'*total(1:p-1))/b(indb);
bias = total(p:end);



function x = ADMM_LP(A,b)
% ADMM for solving a linear program in form of
% min ||Ax + b||_1 
[m, n] = size(A);

mu = 2 *m / norm(b,1);
error = 1;
iter = 0; 
Q = pinv(A);
x = -Q*b;
z = A*x-b; 
maxIter = 5000;

while error>1e-6 && iter<=maxIter
    iter = iter + 1; 
    g = A*x+b-z/mu;
    y = sign(g).*max(abs(g)-1/mu,0);
    x = Q*(y-b+z/mu);
    z = z + mu*(y-A*x-b);
    mu = min(1e8,1.01*mu);
    error = max(y-A*x-b);
   
    if iter == maxIter
        disp('Maximum Iterations Reached');
    end
end




          
            