% This is the companion code for the following paper that appears in ICASSP
% 2018:
% Y. Yang, M. Pesavento, "A parallel best-response algorithm 
% with exact line search for nonconvex sparsity-regularized rank minimization"
% available at http://orbilu.uni.lu/handle/10993/33772

clear; clc;

% format long;

%parameters
N=1000; K=4000; I=4000; % X: N*K, D: N*I, S: I*K
rho=10  ; % rank of X; P: N*rho, Q: rho*K

% number of samples in Monte Carlo simulations
Sample=20;

% optimal objective functin value
val_opt=zeros(Sample,1);

% the following algorithms are compared:
% "_g" stands for the BCD algorithm proposed in [5]
% "_j" stands for the proposed parallel best-response algorithm
% "_a" stands for the ADMM method proposed in [4]
% maximum number of iterations
MaxIter_g=10; val_g=zeros(Sample,MaxIter_g+1);  time_g=zeros(Sample,MaxIter_g+1); %
MaxIter_j=20; val_j=zeros(Sample,MaxIter_j+1);  time_j=zeros(Sample,MaxIter_j+1); 
MaxIter_a=100; val_a=zeros(Sample,MaxIter_a+1);  time_a=zeros(Sample,MaxIter_a+1); 
for s=1:1:Sample
    disp(['sample ' num2str(s)]);
    
    % generate the data
    D=randi([0 1],N,I);
    P=sqrt(100/I)*randn(N,rho);
    Q=sqrt(100/K)*randn(rho,K);
    X0=P*Q; % perfect X
    V=0.1*randn(N,K); % noise

    S=zeros(I,K);    
    Z=randi([1 100],I,K);
    S=S+(Z>=95*ones(I,K))-(Z<=5*ones(I,K));
    clear Z;
    
    Y=X0+D*S+V; % observation

    lambda=0.1*norm(Y); %spectral norm
    mu=0.1*norm(D'*Y,inf);    
    
    % computing the optimal solution by running the proposed algorithm for
    % a sufficiently large number of iterations
%     val_opt(s)=FUN_optsol(D,Y,lambda,mu,P,Q,S);

    % initial point (common for all algorithms)
    P0=sqrt(I)/10*randn(N,rho);
    Q0=sqrt(K)/10*randn(rho,K);
    S0=zeros(I,K);
    % initial value
    val0=0.5*norm(Y-P0*Q0-D*S0,'fro')^2+0.5*lambda*(norm(P0,'fro')^2+norm(Q0,'fro')^2)+mu*norm(vec(S0),1);
    
    
    % BCD (Gauss-Seidel) algorithm
    P=P0; Q=Q0; S=S0;
    d_DtD=diag(diag(D'*D));
    val_g(s,1)=val0; 
    time_g(s,1)=0;
    for t=1:1:MaxIter_g
        tic;
        
        P=(Y-D*S)*Q'*(Q*Q'+lambda*eye(rho))^-1;
        
        Q=(P'*P+lambda*eye(rho))^-1*P'*(Y-D*S);
        
        for i=1:1:I
            q_i=-D(:,i)'*(P*Q+D*S-D(:,i)*S(i,:)-Y);
            S(i,:)=(max(q_i-mu*ones(1,K),0)-max(-q_i-mu*ones(1,K),0))/d_DtD(i,i);
            clear q_i;
        end
        
        time_g(s,t+1)=toc+time_g(s,t);
        
        val_g(s,t+1)=0.5*norm(Y-P*Q-D*S,'fro')^2+0.5*lambda*(norm(P,'fro')^2+norm(Q,'fro')^2)+mu*norm(vec(S),1);
        
        disp(['Gauss-Seidel algorithm, iteration ' num2str(t) ', time ' num2str(time_g(s,t+1)) ', value ' num2str(val_g(s,t+1))]);
    end
    
    % Jacobi algorithm
    P=P0; Q=Q0; S=S0;    
    val_j(s,1)=val0; 
    time_j(s,1)=0;
    d_DtD=diag(diag(D'*D));
    for t=1:1:MaxIter_j
        tic;
        
        Y_DS=Y-D*S;
        
        P_new=Y_DS*Q'*(Q*Q'+lambda*eye(rho))^-1;
        cP=P_new-P;
        
        Q_new=(P'*P+lambda*eye(rho))^-1*P'*Y_DS;
        cQ=Q_new-Q;
        
        G=d_DtD*S-D'*(P*Q-Y_DS);
        
        S_new=d_DtD^-1*(max(G-mu*ones(I,K),zeros(I,K))-max(-G-mu*ones(I,K),zeros(I,K)));
        cS=S_new-S;
        
        %-------------------- to calculate the stepsize by exact line search----------------
        A=cP*cQ;
        B=P*cQ+cP*Q+D*cS;
        C=P*Q+D*S-Y;
        
        a=2*sum(sum(A.^2,1));
        b=3*sum(sum(A.*B,1));
        c=sum(sum(B.^2,1))+2*sum(sum(A.*C,1))+lambda*sum(sum(cP.^2,1))+lambda*sum(sum(cQ.^2,1));
        d=sum(sum(B.*C,1))+lambda*sum(sum(cP.*P,1))+lambda*sum(sum(cQ.*Q,1))+mu*(norm(vec(S_new),1)-norm(vec(S),1));

        % calculating the stepsize by closed-form expression
        Sigma1=(-(b/3/a)^3+b*c/6/a^2-d/2/a);
        Sigma2=c/3/a-(b/3/a)^2;
        Sigma3=Sigma1^2+Sigma2^3;
        Sigma3_sqrt=sqrt(Sigma3);
        if Sigma3>=0
            gamma=nthroot(Sigma1+Sigma3_sqrt,3)...
                +nthroot(Sigma1-Sigma3_sqrt,3)...
                -b/3/a;
        else
            C1=1; C1(4)=-(Sigma1+Sigma3_sqrt);
            C2=1; C2(4)=-(Sigma1-Sigma3_sqrt);
            R=real(roots(C1)+roots(C2))-b/3/a*ones(3,1);
            gamma=min(R(R>0));
            clear C1 C2 R;
        end
        gamma=max(0,min(gamma,1));

%         % performing the exact line search the stepsize by Matlab
%         f=@(x) a/4*x^4+b/3*x^3+c/2*x^2+d*x;
%         options = optimoptions('fmincon','Display','off','Algorithm','sqp');
%         gamma = fmincon(f,0,[],[],[],[],0,1,[],options);
        
        % variable update
        P=P+gamma*cP;
        Q=Q+gamma*cQ;
        S=S+gamma*cS;
        
        time_j(s,t+1)=toc+time_j(s,t);                
        
        val_j(s,t+1)=0.5*norm(Y-P*Q-D*S,'fro')^2+0.5*lambda*(norm(P,'fro')^2+norm(Q,'fro')^2)+mu*norm(vec(S),1);
        disp(['Jacobi algorithm, iteration ' num2str(t) ', time ' num2str(time_j(s,t+1)) ', value ' num2str(val_j(s,t+1))...
            ', stepsize ' num2str(gamma)]);
    end

    % ADMM algorithm
    P=P0; Q=Q0; E=S0; F=S0; Pi=zeros(I,K); c=100;
    val_a(s,1)=val0;
    time_a(s,1)=0;
    for t=1:1:MaxIter_a,
        tic;
        Q=(P'*P+lambda*eye(rho))^-1*P'*(Y-D*E);
        F=max(E+Pi/c-mu/c*ones(I,K),zeros(I,K))-max(-E-Pi/c-mu/c*ones(I,K),zeros(I,K));
        P=(Y-D*E)*Q'*(Q*Q'+lambda*eye(rho))^-1;
        E=(D'*D+c*eye(I))^-1*(c*F-D'*(P*Q-Y)-Pi);
        Pi=Pi+c*(E-F);
        val_a(s,t+1)=0.5*norm(Y-P*Q-D*E,'fro')^2+0.5*lambda*(norm(P,'fro')^2+norm(Q,'fro')^2)+mu*norm(vec(E),1);
        time_a(s,t+1)=toc+time_a(s,t);
        disp(['ADMM algorithm, iteration ' num2str(t) ', time ' num2str(time_a(s,t+1)) ', value ' num2str(val_a(s,t+1))]);
    end;

end

mean_val_g=mean(val_g,1);
mean_val_j=mean(val_j,1);
mean_val_a=mean(val_a,1);

figure;
semilogy(0:1:MaxIter_g,mean_val_g,'k--','LineWidth',1);
hold on; box on;
semilogy(0:1:MaxIter_j,mean_val_j,'r','LineWidth',1);
legend('BCD algorithm','Parallel best-response algorithm');
xlabel('number of iterations');ylabel('objective function value');

figure;
semilogy(mean(time_g,1),mean_val_g,'k--','LineWidth',1);
hold on; box on;
semilogy(mean(time_j,1),mean_val_j,'r','LineWidth',1);
legend('Gauss-Seidel algorithm','Jacobi algorithm');
xlabel('CPU time (seconds)');ylabel('objective function value');
% axis([0 2 10^5 10^7]);