function [x,flag,relres,iter,resvec] = pcg(A,B,tol,maxit,M1,M2,x0,varargin)
%MATLAB Code Generation Library Function

%   Copyright 2020 The MathWorks, Inc.

%#codegen

localChkType(A);
coder.internal.assert(isa(B,'double') && ~issparse(B),...
    'Coder:toolbox:PcgRHSInvalidType');
ONE = coder.internal.indexInt(1);
ZERO = coder.internal.indexInt(0);
if isfloat(A)
    n = coder.internal.indexInt(size(A,1));
    coder.internal.assert(size(A,2) == n,'MATLAB:pcg:NonSquareMatrix');
    coder.internal.assert(iscolumn(B) && size(B,1) == n,...
        'MATLAB:pcg:RSHsizeMatchCoeffMatrix', n);
else
    coder.internal.assert(iscolumn(B),'MATLAB:pcg:RSHnotColumn');
    n = coder.internal.indexInt(size(B,1));
end
b = B(:);
if nargin < 3 || (coder.internal.isConst(isempty(tol)) && isempty(tol))
    toler = 1e-6;
else
    if isempty(tol)
        coder.internal.error('Coder:toolbox:PcgVarSizeInputEmpty');
    end
    coder.internal.assert(isnumeric(tol) && isscalar(tol) && isreal(tol),...
        'Coder:toolbox:PcgInvalidTolerance');
    toler = double(tol(1));
end

warned = false;
if toler <= eps
    coder.internal.warning('MATLAB:pcg:tooSmallTolerance');
    warned = true;
    toler = eps;
elseif toler >= 1
    coder.internal.warning('MATLAB:pcg:tooBigTolerance');
    warned = true;
    toler = 1 - eps;
end

if nargin < 4 || (coder.internal.isConst(isempty(maxit))&& isempty(maxit))
    itermax = min(n,20);
else
    if isempty(maxit)
         coder.internal.error('Coder:toolbox:PcgVarSizeInputEmpty');
    end
    coder.internal.assert(isnumeric(maxit) && isscalar(maxit) && isreal(maxit),...
        'Coder:toolbox:PcgInvalidMaxit');
    itermax = coder.internal.indexInt(maxit(1));
    itermax = max(itermax,ZERO);
end

n2b = norm(b);
if coder.internal.safeEq(n2b,0)
    x = zeros(n,1);
    flag = 0;
    relres = 0;
    iter = 0;
    resvec = 0;
    return
end

istriuM1 = false;
istrilM1  = false;
ishermitianM1 = false;

if nargin < 5 || (coder.internal.isConst(isempty(M1)) && isempty(M1))
    hasM1 = false;
else
    if isempty(M1)
        coder.internal.error('Coder:toolbox:PcgVarSizeInputEmpty');
    end
    hasM1 = true;
    localChkType(M1);
    if isfloat(M1)
        coder.internal.assert(size(M1,1) == n && size(M1,2) == n,...
            'MATLAB:pcg:WrongPrecondSize', n);
        coder.internal.assert(~issparse(M1),...
            'Coder:toolbox:PcgNonSparsePreconditioner')
        [istriuM1,istrilM1,ishermitianM1] = chkMatrixType(M1);
    end
end

istriuM2 = false;
istrilM2 = false;
ishermitianM2 = false;

if nargin < 6 || (coder.internal.isConst(isempty(M2)) && isempty(M2))
    hasM2 = false;
else
    if isempty(M2)
        coder.internal.error('Coder:toolbox:PcgVarSizeInputEmpty');
    end
    hasM2 = true;
    localChkType(M2);
    if isfloat(M2)
        coder.internal.assert(~issparse(M2),...
                     'Coder:toolbox:PcgNonSparsePreconditioner')
        coder.internal.assert(size(M2,1) == n && size(M2,2) == n,...
            'MATLAB:pcg:WrongPrecondSize', n);
        [istriuM2,istrilM2,ishermitianM2] = chkMatrixType(M2);
    end
end

ISREAL = isreal(b) &&...
         (isreal(A) || isa(A,'function_handle') && isreal(A(b,varargin{:}))) &&...
         (~hasM1 || isreal(M1) || isa(M1,'function_handle') && isreal(M1(b,varargin{:}))) &&...
         (~hasM2 || isreal(M2) || isa(M2,'function_handle') && isreal(M2(b,varargin{:})));

if nargin < 7 ||(coder.internal.isConst(isempty(x0)) && isempty(x0))
    if ISREAL
        x =  zeros(n,1);
    else        
        x = complex(zeros(n,1));
    end
else
    if isempty(x0)
        coder.internal.error('Coder:toolbox:PcgVarSizeInputEmpty');
    end
    coder.internal.assert(isa(x0,'double') && ~issparse(x0),...
        'Coder:toolbox:PcgInitGuessInvalidType');
    coder.internal.assert(iscolumn(x0) && size(x0,1) == n, ...
        'MATLAB:pcg:WrongInitGuessSize', n);
    
    if ISREAL && isreal(x0)
        x = x0(:);      
    else
        x = complex(x0(:));
    end
end

% arguments after x0 are the additional inputs (if any) required to
% evaluate the function handles
coder.internal.errorIf(nargin > 7 && isfloat(A) &&...
            isfloat(M1) && isfloat(M2),'MATLAB:pcg:TooManyInputs')


flag = 1;
xmin = x;
imin = ZERO;
tolb = toler*n2b;
r =  b - localOperation('mtimes',A,x,varargin{:});
normr = norm(r);
normr_act = normr;
iter = 0;

if normr <= tolb
    flag = 0;
    relres = normr/n2b;
    resvec = normr;
    return
end
resvec = zeros(itermax+1,1);         % Preallocate vector for norm of residuals
resvec(1,:) = normr;                 % resvec(1) = norm(b-A*x0)
normrmin = normr;                    % Norm of minimum residual
rho = ones('like',x);
stag = ZERO;                          % stagnation of the method
moresteps = ZERO;
maxmsteps = min([coder.internal.indexDivide(n,coder.internal.indexInt(50)),...
           coder.internal.indexInt(5), n - itermax]);
maxstagsteps = coder.internal.indexInt(3);

ii = coder.internal.indexInt(1);
p = coder.nullcopy(r);

while(ii <= itermax)
    if hasM1
        if isa(M1,'float')
            y = localOperation('mldivide',M1,r,istriuM1,istrilM1,ishermitianM1);
        else
            y = localOperation('mldivide',M1,r,varargin{:});
        end
        if coder.internal.anyNonFinite(y)
            flag = 2;
            break
        end
    else % no preconditioner
        y = r;
    end

    if hasM2
        if isa(M2,'float')
            z = localOperation('mldivide',M2,y,istriuM2,istrilM2,ishermitianM2);
        else
            z = localOperation('mldivide',M2,y,varargin{:});
        end
        if coder.internal.anyNonFinite(z)
            flag = 2;
            break
        end
    else % no preconditioner
        z = y;
    end
    rho1 = rho;
    rho = r'*z;
    if rho == 0 || isinf(rho)
        flag = 4;
        break
    end

    if ii == ONE
        p = z;
    else
        beta = rho/rho1;
        if beta == 0 || isinf(beta)
            flag = 4;
            break
        end
          p = z + beta*p;
    end

    q = localOperation('mtimes',A,p,varargin{:});
    pq = p'*q;
    if pq <= 0 || isinf(pq)
        flag = 4;
        break
    else
        alpha = rho/pq;
    end

    if isinf(alpha)
        flag = 4;
        break
    end

    % Check for stagnation of the method
    if (norm(p)*abs(alpha) < eps*norm(x))
        stag = stag + ONE;
    else
        stag = ZERO;
    end

  %  x = x + alpha * p;             % form new iterate
    x = coder.internal.blas.xaxpy(n,alpha,p,ONE,ONE,x,ONE,ONE);
  %  r = r - alpha * q;
    r = coder.internal.blas.xaxpy(n,-alpha,q,ONE,ONE,r,ONE,ONE);

    normr = norm(r);
    normr_act = normr;
    resvec(ii+1,1) = normr;

    % check for convergence
    if (normr <= tolb || stag >= maxstagsteps || moresteps)
        % r = b - A*x;
        r =  b - localOperation('mtimes',A,x,varargin{:});
        normr_act = norm(r);
        resvec(ii+1,1) = normr_act;
        if (normr_act <= tolb)
            flag = 0;
            iter = double(ii);
            break
        else
            if stag >= maxstagsteps && moresteps == ZERO
                stag = ZERO;
            end
            moresteps = moresteps + ONE;
            if moresteps >= maxmsteps
                if ~warned
                    coder.internal.warning('MATLAB:pcg:tooSmallTolerance');
                end
                flag = 3;
                iter = double(ii);
                break;
            end
        end
    end

    if (normr_act < normrmin)      % update minimal norm quantities
        normrmin = normr_act;
        xmin = x;
        imin = ii;
    end

    if stag >= maxstagsteps
        flag = 3;
        break;
    end
    ii = ii + ONE;
end

if itermax == ZERO
    ii = ZERO;
end

if ii > itermax
    ii = itermax;
end

if (flag == 0)
    relres = normr_act / n2b;
else
    % r_comp = b - A*xmin;
     r_comp = b - localOperation('mtimes',A,xmin,varargin{:});
    if norm(r_comp) <= normr_act
        x = xmin;
        iter = double(imin);
        relres = norm(r_comp) / n2b;
    else
        iter = double(ii);
        relres = normr_act / n2b;
    end
end

% truncate the zeros from resvec
if ((flag <= 1) || (flag == 3))
    resvec = resvec(1:ii+1,:);
else
    resvec = resvec(1:ii,:);
end
end


function localChkType(a)
coder.inline('always');
coder.internal.assert(isa(a,'double') || isa(a,'function_handle'),...
    'Coder:toolbox:PcgInvalidInput');
end

function y = localOperation(op,a,x,varargin)
coder.inline('always');
if isa(a,'function_handle')
    y = a(x,varargin{:});
    coder.internal.assert(iscolumn(y),'MATLAB:iterapp:MustReturnColumn',...
        'function',func2str(a));
else
    if strcmpi(op,'mtimes')
        y = a*x;
    elseif strcmpi(op,'mldivide')
        ISTRIU = varargin{1};
        ISTRIL = varargin{2};
        ISHERMITIAN = varargin{3};
        if ISTRIU
            y = linsolve(a,x,struct('UT',true));
        elseif ISTRIL
            y = linsolve(a,x,struct('LT',true));
        else
            if ISHERMITIAN
                n = coder.internal.indexInt(size(a,1));
                ONE = coder.internal.indexInt(1);
                [L,info] = coder.internal.lapack.xpotrf('L',n,a,ONE,n);
                if info == cast(0,'like',coder.internal.lapack.info_t)
                    % L*L'*y = x; Note L here is tril(L);
                    % y1 = L'*y;
                    % we need only the lower triangular part of L
                    y1 = linsolve(L,x,struct('LT',true));
                    y =  linsolve(L,y1,struct('LT',true,'TRANSA',true));
                else
                    y = linsolve(a,x);
                end
            else
                y = linsolve(a,x);
            end
        end
    else
        coder.internal.assert(false,'MATLAB:iterapp:InvalidOp');
    end
end
end

function[istriuA,istrilA,ishermitianA] = chkMatrixType(A)
coder.inline('always');
istrilA = false;
ishermitianA = false;
istriuA = istriu(A);
if ~istriuA
    istrilA = istril(A);
end
if ~(istriuA || istrilA)
    ishermitianA = ishermitian(A);
end

end

% LocalWords:  Hsize Hnot Precond resvec preconditioner pq xmin iterapp TRANSA
