clear; clc;

%parameters
N=10000; K=50000;  % A: N*K
density=0.1;
MaxIter=100; % maximum number of iterations
P=5; Kp=K/P; % number of block variables and the corresponding dimension

Sample=10; % number of repeatitions in the Monte Carlo simulations

% the achieved objective value versus the number of iterations
val_g=zeros(Sample,MaxIter+1);
val_j=zeros(Sample,MaxIter+1);
val_b=zeros(Sample,MaxIter+1);

% the required CPU time versus the number of iterations
time_g=zeros(Sample,MaxIter+1);
time_j=zeros(Sample,MaxIter+1);
time_b=zeros(Sample,MaxIter+1);
time_accumulate_g=zeros(Sample,MaxIter+1);
time_accumulate_j=zeros(Sample,MaxIter+1);
time_accumulate_b=zeros(Sample,MaxIter+1);

% the achieved error versus the number of iterations
error_g=zeros(Sample,MaxIter+1);
error_j=zeros(Sample,MaxIter+1);
error_b=zeros(Sample,MaxIter+1);

for s=1:1:Sample
    disp(['Sample ' num2str(s)]);
    
    % generating the parameters
    [A,b,mu]=FUN_parameter(N,K,density);
    
    d_AtA=sum(A.^2,1)'; % diagonal of AtA
    % d_AtA_InvSR=d_AtA.^-0.5; % inverse square root
    
    mu_vec=mu*ones(K,1);
    mu_vec_normalized=mu_vec./d_AtA;
    
    % the parallel STELA algorithm
    % initialization
    x=zeros(K,1);
    residual= A*x-b;
    Gradient=(residual'*A)';
    error_j(s,1)=max(abs(Gradient-min(max(Gradient-x,-mu_vec),mu_vec))); % error in optimality condition
    for t=1:1:MaxIter
        tic;
        % compute best-response
        x_hat=FUN_quad(x-Gradient./d_AtA,mu_vec_normalized,K);
        
        % compute stepsize
        x_dif=(x_hat-x);
        Ax_dif=A*x_dif;
        
        stepsize=min((-residual'*Ax_dif-mu_vec'*((abs(x_hat)-abs(x))))/(Ax_dif'*Ax_dif),1); % sum(abs(x)) is faster than norm(x,1)
        
        % update variable
        x=x+stepsize*x_dif; clear x_hat x_dif;
        
        % calculate intermediate variables for the next iteration (recursive update)
        residual=residual+stepsize*Ax_dif; clear Ax_dif;
        Gradient=(residual'*A)';
        
        time_j(s,t+1)=toc;
        
        % calculate variable precision
        error_j(s,t+1)=max(abs(Gradient-min(max(Gradient-x,-mu_vec),mu_vec))); % error in optimality condition
        
        %     calculate objective value
        val_j(s,t+1)=1/2*(residual'*residual)+mu_vec'*abs(x);
        
        disp(['Parallel STELA: iteration ' num2str(t) ' with value ' num2str(val_j(s,t+1))...
            ' and precision ' num2str(error_j(s,t+1)) ' and time ' num2str(sum(time_j(s,1:t+1)))]);
    end
    clear Gradient residual x;
    
    % the hybrid STELA algorithm
    % initialization
    mu_vec_p=mu*ones(Kp,1);
    x=zeros(K,1);
    residual_p=-b; %residual_p=A*x-b;
    Gradient=(residual_p'*A)';
    error_g(s,1)=max(abs(Gradient-min(max(Gradient-x,-mu_vec),mu_vec))); % error in optimality condition
    for t=1:1:MaxIter
        tic;
        for p=1:1:P
            % to extract the p-th block of A and x
            A_p=A(:,(p-1)*Kp+1:p*Kp);
            x_p=x((p-1)*Kp+1:p*Kp);
            d_AtA_p=d_AtA((p-1)*Kp+1:p*Kp);
            
            % to calculate the best-response
            Gradient_p=(residual_p'*A_p)';            
            x_p_hat=FUN_quad(x_p-Gradient_p./d_AtA_p,mu_vec_p./d_AtA_p,Kp); clear Gradient_p d_AtA_p;
            
            % to compute stepsize
            x_p_dif=(x_p_hat-x_p);
            Ax_p_dif=A_p*x_p_dif; clear A_p;            
            stepsize=min((-residual_p'*Ax_p_dif-mu_vec_p'*((abs(x_p_hat)-abs(x_p))))/(Ax_p_dif'*Ax_p_dif),1); % sum(abs(x)) is faster than norm(x,1)            
            x((p-1)*Kp+1:p*Kp)=x_p+stepsize*x_p_dif; clear x_p x_p_hat x_p_dif;
            
            % calculate intermediate variables for the next iteration (recursive update)
            residual_p=residual_p+stepsize*Ax_p_dif; clear stepsize Ax_p_dif;
        end
        clear p;
        
        time_g(s,t+1)=toc;
        
        val_g(s,t+1)=1/2*(residual_p'*residual_p)+mu_vec'*abs(x);
        
        Gradient=(residual_p'*A)';
        error_g(s,t+1)=max(abs(Gradient-min(max(Gradient-x,-mu_vec),mu_vec))); % error in optimality condition
        
        disp(['Hybrid STELA: iteration ' num2str(t) ' with value ' num2str(val_g(s,t+1))...
            ' and precision ' num2str(error_g(s,t+1)) ' and time ' num2str(sum(time_g(s,1:t+1)))]);
    end
    clear Gradient residual_p x;
    
    % the block successive convex approximation algorithm
    % initialization
    mu_vec_p=mu*ones(Kp,1);
    x=zeros(K,1);
    residual_p=-b; %residual_p=A*x-b;
    Gradient=(residual_p'*A)';
    error_b(s,1)=max(abs(Gradient-min(max(Gradient-x,-mu_vec),mu_vec))); % error in optimality condition
    for t=1:1:MaxIter
        tic;
        for p=1:1:P
            A_p=A(:,(p-1)*Kp+1:p*Kp);
            x_p=x((p-1)*Kp+1:p*Kp);
            d_AtA_p=d_AtA((p-1)*Kp+1:p*Kp);
%             residual_p=-b_Ax;
            Gradient_p=(residual_p'*A_p)';
            
            x_p_hat=FUN_quad(x_p-Gradient_p./d_AtA_p,mu_vec_p./d_AtA_p,Kp); clear Gradient_p d_AtA_p;
            
            % decreasing stepsize
            x_p_dif=(x_p_hat-x_p);
            Ax_p_dif=A_p*x_p_dif; clear A_p;            
            stepsize=1/t;            
            x((p-1)*Kp+1:p*Kp)=x_p+stepsize*x_p_dif; clear x_p x_p_hat x_p_dif;
            
            % calculate intermediate variables for the next iteration (recursive update)
            residual_p=residual_p+stepsize*Ax_p_dif; clear stepsize Ax_p_dif;
        end
        clear p;
        
        time_b(s,t+1)=toc;
        
        val_b(s,t+1)=1/2*(residual_p'*residual_p)+mu_vec'*abs(x);
        
        Gradient=(residual_p'*A)';
        error_b(s,t+1)=max(abs(Gradient-min(max(Gradient-x,-mu_vec),mu_vec))); % error in optimality condition
        
        disp(['Block successive convex approximation: iteration ' num2str(t) ' with value ' num2str(val_b(s,t+1))...
            ' and precision ' num2str(error_b(s,t+1)) ' and time ' num2str(sum(time_b(s,1:t+1)))]);
    end
    clear Gradient residual_p x;
    
    clear A b mu mu_vec mu_vec_normalized mu_vec_p d_AtA;
end
clear s;

for t=1:1:MaxIter+1
    time_accumulate_j(:,t)=sum(time_j(:,1:t),2)';
    time_accumulate_g(:,t)=sum(time_g(:,1:t),2)';
    time_accumulate_b(:,t)=sum(time_b(:,1:t),2)';
end
clear t;

figure;
subplot(2,1,1);
semilogy(0:1:MaxIter,mean(error_j,1),'k','Linewidth',1);
hold on;
semilogy(0:1:MaxIter,mean(error_g,1),'r--','Linewidth',1);
semilogy(0:1:MaxIter,mean(error_b,1),'b-.','Linewidth',1);
legend('Parallel STELA (proposed)', 'Hybrid STELA (proposed)','BSCA (state-of-the-art)');
xlabel('(a): number of iterations'); ylabel('error e(x^t)');

subplot(2,1,2);
semilogy(mean(time_accumulate_j,1),mean(error_j,1),'k','Linewidth',1);
hold on;
semilogy(mean(time_accumulate_g,1),mean(error_g,1),'r--','Linewidth',1);
semilogy(mean(time_accumulate_b,1),mean(error_b,1),'b-.','Linewidth',1);
legend('Parallel STELA (proposed)', 'Hybrid STELA (proposed)','BSCA (state-of-the-art)');
xlabel('(b): CPU time (seconds)'); ylabel('error e(x^t)');
