clear; clc;

% parameters
N=10000; K=50000;  % A: N*K
density=0.1;
MaxIter=50; % 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
F=1;
objval_j=zeros(Sample,F*MaxIter+1);
objval_i=zeros(Sample,MaxIter+1);

% 
% the required CPU time versus the number of iterations
time_j=zeros(Sample,F*MaxIter+1);
time_i=zeros(Sample,MaxIter+1);

time_accumulate_j=zeros(Sample,F*MaxIter+1);
time_accumulate_i=zeros(Sample,MaxIter+1);


for s=1:1:Sample
    disp(['Sample ' num2str(s)]);
    
    % generating the parameters
    [A,b,mu]=FUN_parameter(N,K,density);
    
    
    % the parallel STELA algorithm
    % initialization
    d_AtA=10*ones(K,1); % diagonal of AtA       
    mu_vec=mu*ones(K,1);
    mu_vec_normalized=mu_vec./d_AtA;
    x=zeros(K,1);
    Ax=A*x;
    residual=2*Ax+cos(Ax)-b;
    val_f=0.5*(residual'*residual);
    val_g=mu_vec'*abs(x);
    objval_j(s,1)=val_f+val_g;
    for t=1:1:F*MaxIter
        tic;

        Gradient_sigma=(2*ones(N,1)-sin(Ax));
        Gradient_f=((residual.*Gradient_sigma)'*A)';
        
        % compute best-response
        Bx=FUN_quad(x-Gradient_f./d_AtA,mu_vec_normalized,K);
        
        % compute stepsize
        x_dif=(Bx-x);
        Ax_dif=A*x_dif;
        val_g_new=mu_vec'*abs(Bx);

        % to compute the stepsize by successive line search
        descent=Gradient_f.*x_dif+mu_vec.*abs(Bx)-mu_vec.*abs(x);        
        alpha=0.01; beta=0.5; stepsize=1;
        while 1
            x_new=x+stepsize*x_dif;
            Ax_new=Ax+stepsize*Ax_dif;
            residual_new=2*Ax_new+cos(Ax_new)-b;
            val_f_new=0.5*(residual_new'*residual_new);
            if val_f_new-val_f+stepsize*(val_g_new-val_g)<=alpha*stepsize*sum(descent)
                val_f=val_f_new; clear val_f_new
                residual=residual_new; clear residual_new
                Ax=Ax_new; clear Ax_new
                x=x_new; clear x_new
                val_g=mu_vec'*abs(x);
                break;
            else
                stepsize=stepsize*beta;
            end
        end
        time_j(s,t+1)=toc;

        objval_j(s,t+1)=val_f+val_g;
        
        disp(['Parallel STELA: iteration ' num2str(t) ' with value ' num2str(objval_j(s,t+1))...
            ', descent ' num2str(sum(descent)) ', and time ' num2str(sum(time_j(s,1:t+1)))]);                
    end
    
    % the parallel ISTA algorithm
    % initialization
    mu_vec=mu*ones(K,1);
    x=zeros(K,1);
    Ax=A*x;
    residual=2*Ax+cos(Ax)-b;
    val_f=0.5*(residual'*residual);
    val_g=mu_vec'*abs(x);
    objval_i(s,1)=val_f+val_g;
    for t=1:1:MaxIter
        tic;

        Gradient_sigma=(2*ones(N,1)-sin(Ax));
        Gradient_f=((residual.*Gradient_sigma)'*A)';
        
        % compute best-response
        L=1; eta=2;        
        while 1
            d_AtA=L*ones(K,1);
            mu_vec_normalized=mu_vec./d_AtA;
            x_new=FUN_quad(x-Gradient_f./d_AtA,mu_vec_normalized,K);
            Ax_new=A*x_new;
            residual_new=2*Ax_new+cos(Ax_new)-b;
            val_f_new=0.5*(residual_new'*residual_new);
            val_g_new=mu_vec'*abs(x_new);

            if val_f_new+val_g_new<=objval_i(s,t)-10^-5*L/2*(x_new-x)'*(x_new-x)
                x=x_new; clear Bx
                Ax=Ax_new; clear ABx
                val_f=val_f_new; clear val_f_new
                val_g=val_g_new; clear val_g_new
                residual=residual_new; clear residual_new
                break;
            else
                L=L*eta;
            end
        end
       
        time_i(s,t+1)=toc;

        objval_i(s,t+1)=val_f+val_g;
                
        disp(['Parallel ISTA: iteration ' num2str(t) ' with value ' num2str(objval_i(s,t+1))...
            ', and time ' num2str(sum(time_i(s,1:t+1)))]);
    end
end

for t=1:1:MaxIter+1
    time_accumulate_j(:,t)=sum(time_j(:,1:t),2)';
    time_accumulate_i(:,t)=sum(time_i(:,1:t),2)';
end
clear t;

% optval=objval_j(:,end);
% 
% subplot(2,1,1);
% semilogy(0:1:MaxIter,mean(objval_j(:,1:MaxIter+1),1)-mean(optval)*ones(1,MaxIter+1),'r','Linewidth',1);
% hold on;
% semilogy(0:1:MaxIter,mean(objval_i,1)-mean(optval)*ones(1,MaxIter+1),'k--','Linewidth',1);
% legend('Parallel STELA (proposed)', 'ISTA (state-of-the-art)');xlabel('(a): number of iterations'); ylabel('error in function value');

subplot(2,1,1);
semilogy(0:1:MaxIter,mean(objval_j(:,1:MaxIter+1),1),'r','Linewidth',1);
hold on;
semilogy(0:1:MaxIter,mean(objval_i,1),'k--','Linewidth',1);
legend('Parallel STELA (proposed)', 'ISTA (state-of-the-art)');xlabel('(a): number of iterations'); ylabel('function value');

subplot(2,1,2);
semilogy(mean(time_accumulate_j(:,1:MaxIter+1),1),mean(objval_j(:,1:MaxIter+1),1),'r','Linewidth',1);
hold on;
semilogy(mean(time_accumulate_i,1),mean(objval_i,1),'k--','Linewidth',1);
legend('Parallel STELA (proposed)', 'ISTA (state-of-the-art)');xlabel('(a): CPU time (seconds)'); ylabel('function value');
