function [time,NormResidualz,NormResidualR,z,p] = Alg5PAL(F,PAR,z,p)
%ALG5PAL First-order proximal augmented Lagrangian method, specialized to Ax=b.
%
%   [time,NormResidualz,NormResidualR,z,p] = Alg5PAL(F,PAR,z,p)
%
% This implements Algorithm 5 of arXiv:2206.01209 for the linear equality
% constrained case
%
%       minimize    f(x) + P(x)
%       subject to  A*x = b,
%
% represented by F.Residual(x) = A*x-b and F.Adjoint(y) = A'*y.
% The cone is K={0}, so K*=R^m and the multiplier projection is the identity.
%
% The proximal AL subproblem at outer iterate k is
%
%   min_x f(x) + <p_k,A*x-b> + rho_k/2 ||A*x-b||^2
%           + 1/(2*rho_k)||x-x_k||^2 + P(x),
%
% and it is solved by APGAlg2 with the paper's inputs
%
%   gamma0 = 1/rho_k,   mu = mu0 + 1/rho_k.
%

if ~isfield(PAR,'verbose')
    PAR.verbose = true;
end
NormResidualz = PAR.norm_fn(F.Residual(z));
NormResidualR = inf;
totalIter = 0;

if PAR.verbose
    fprintf('\nAlgorithm 5 Proximal AL for Ax=b\n')
    fprintf('%6s %10s %12s %14s %14s %14s %14s\n', ...
        'outer','inner','total','primal_res','inner_res','eta','rho')
end

tic
k=0;
while NormResidualz>PAR.epsilonhat || NormResidualR>PAR.epsilonhat
    outerIter = k + 1;
    rho = PAR.rho0*(PAR.zeta^k);
    eta = PAR.eta0*(PAR.sigma^k);

    zold = z;
    pold = p;

    % Smooth part f_k from equation (31), specialized to K={0} and g(x)=A*x-b.
    Fk = F;
    [fs,grad] = makeAlg5Oracle(F,PAR,pold,rho,zold);
    Fk.fs = fs;
    Fk.grad = grad;

    PAR.k=outerIter;
    PAR.x0=zold;
    PAR.epsilon=eta;
    PAR.gamma0=1/rho;
    PAR.mu = PAR.mu0 + 1/rho;

    %%%Call Alg 2 to minimize Prox AL function 
    [z,APGvec,innerIter] = APGAlg2(Fk,PAR);
    totalIter = totalIter + innerIter;

    Residualz = F.Residual(z);
    NormResidualz = PAR.norm_fn(Residualz);

    % Standard dual update
    p = pold + rho*Residualz;

    % From the subproblem residual, a stationarity residual is 
    % ||APGvecx^{k+1}-x^k/rho||.
    NormResidualR = PAR.norm_fn(APGvec-(z-zold)/(rho));
    % keyboard

    if PAR.verbose
        fprintf('%6i %10i %12i %14.6e %14.6e %14.6e %14.6e\n', ...
            outerIter, innerIter, totalIter, NormResidualz, NormResidualR, eta, rho)
    end

    k=k+1;
end

time = toc;
outerIter = k;

if PAR.verbose
    fprintf('Finished Algorithm 5: outer=%d, total inner=%d, time=%.2f seconds.\n', ...
        outerIter, totalIter, time)
end

end


%%%%Subroutine that efficiently computes function and gradient evaluations of augmented Lagrangian
%%%%function.

function [fs,grad] = makeAlg5Oracle(F,PAR,p,rho,xold)
finit = F.finit;
gradinit = F.gradinit;
Residual = F.Residual;
Adjoint = F.Adjoint;
prod_fn = PAR.prod_fn;
norm_fn = PAR.norm_fn;

last_x = [];
last_r = [];

fs = @valueOracle;
grad = @gradOracle;

    function value = valueOracle(x)
        r = residualAt(x);
        dx = x - xold;
        value = finit(x) + prod_fn(p,r) + 0.5*rho*norm_fn(r)^2 ...
            + 0.5*(1/rho)*norm_fn(dx)^2;
    end

    function g = gradOracle(x)
        r = residualAt(x);
        g = gradinit(x) + Adjoint(p + rho*r) + (1/rho)*(x - xold);
    end

    function r = residualAt(x)
        if isempty(last_x) || ~isequal(x,last_x)
            last_x = x;
            last_r = Residual(x);
        end
        r = last_r;
    end
end


