function [xTilde,stopvec,iter] = APGAlg2(F,PAR)
%APGALG2 APG method with termination criterion (Algorithm 2, arXiv:2206.01209).
%
%   [xTilde,stopvec,iter] = APGAlg2(F,PAR)
%
% Required oracle fields:
%   F.fs          smooth part f(x)
%   F.grad        gradient of f(x)
%   F.prox        prox_{lambda P}(x), called as F.prox(x,lambda)
%
% Required/typical PAR fields:
%   PAR.x0         starting point, used for x^1=z^1
%   PAR.epsilon    residual tolerance epsilon
%   PAR.gamma0     initial stepsize gamma_0
%   PAR.mu         strong convexity parameter used by Algorithm 2; must be > 0
%   PAR.alpha0     initial alpha_0 in [sqrt(mu*gamma0),1]
%   PAR.delta      backtracking contraction in (0,1)
%   PAR.M_alg2     termination-check period M from the paper
%   PAR.maxiter    maximum number of outer APG iterations
%   PAR.norm_fn    norm function
%   PAR.prod_fn    inner product function
%
% The output iter counts proximal evaluations in the main APG loop and in
% Algorithm 3 termination checks.

x = PAR.x0;                 % x^t
z = PAR.x0;                 % z^t
xTilde = x;                 % initialize in case maxiter is reached before termination

alphaPrev = PAR.alpha0;     % alpha_{t-1}, initialized by alpha_0
gammaPrev = PAR.gamma0;     % gamma_{t-1}, initialized by gamma_0

iter = 0;
t = 0;
stopcrit = inf;
stopvec = inf;

while stopcrit > PAR.epsilon && t < PAR.maxiter

    t = t + 1;

    % Algorithm 2 uses gamma_t = gamma0 * delta^{n_t}, where n_t is the
    % smallest nonnegative integer satisfying the line-search condition.
    gamma = PAR.gamma0;

    alpha = solveAlpha(alphaPrev,gammaPrev,gamma,PAR.mu);
    beta = PAR.mu*gamma/alpha;
    denom = 1 - alpha*beta;       % equals 1 - mu*gamma

    y = ((1-alpha)*x + alpha*(1-beta)*z)/denom;
    grady = F.grad(y);
    q = beta*y + (1-beta)*z;

    zNext = F.prox(q - (gamma/alpha)*grady, gamma/alpha);
    iter = iter + 1;
    xNext = (1-alpha)*x + alpha*zNext;

    NormDiff = PAR.norm_fn(y - xNext)^2;
    LinearTerm = 2*gamma*(F.fs(xNext) - F.fs(y) ...
        - PAR.prod_fn(grady,xNext-y));

    while LinearTerm > NormDiff

        gamma = gamma*PAR.delta;

        alpha = solveAlpha(alphaPrev,gammaPrev,gamma,PAR.mu);
        beta = PAR.mu*gamma/alpha;
        denom = 1 - alpha*beta;   % equals 1 - mu*gamma

        y = ((1-alpha)*x + alpha*(1-beta)*z)/denom;
        grady = F.grad(y);
        q = beta*y + (1-beta)*z;

        zNext = F.prox(q - (gamma/alpha)*grady, gamma/alpha);
        iter = iter + 1;
        xNext = (1-alpha)*x + alpha*zNext;

        NormDiff = PAR.norm_fn(y - xNext)^2;
        LinearTerm = 2*gamma*(F.fs(xNext) - F.fs(y) ...
            - PAR.prod_fn(grady,xNext-y));
    end

    x = xNext;
    z = zNext;
    alphaPrev = alpha;
    gammaPrev = gamma;
    if mod(t,PAR.M_alg2) == 0

        [xTilde,gammaTilde,iter3] = APGAlg3(F,PAR,x);

        iter = iter + iter3;
        stopvec = (x - xTilde)/gammaTilde + F.grad(xTilde) - F.grad(x);
        stopcrit = PAR.norm_fn(stopvec);
    end

end

end


function alpha = solveAlpha(alphaPrev,gammaPrev,gamma,mu)
% Solve:
%   gammaPrev*a^2 = (1-a)*alphaPrev^2*gamma + mu*a*gamma*gammaPrev.
%
% Equivalently:
%   gammaPrev*a^2 + (gamma*alphaPrev^2 - mu*gamma*gammaPrev)*a
%       - gamma*alphaPrev^2 = 0.

A = gammaPrev;
B = gamma*alphaPrev^2 - mu*gamma*gammaPrev;
C = -gamma*alphaPrev^2;

disc = B^2 - 4*A*C;
disc = max(disc,0);

alpha = (-B + sqrt(disc))/(2*A);
alpha = real(alpha);
alpha = min(max(alpha,eps),1);

end
