function [X,arg2,condest] = sqrtm(A) %#codegen
%MATLAB Code Generation Library Function

%   Copyright 1984-2017 The MathWorks, Inc.
%#codegen

coder.columnMajor;
eml_invariant(nargin > 0, ...
    'MATLAB:minrhs');
eml_invariant(isa(A,'float'), ...
    'Coder:toolbox:unsupportedClass','sqrtm',class(A));
eml_invariant(ismatrix(A) && size(A,1) == size(A,2), ...
    'Coder:MATLAB:square');
[Q,T] = schur(A,'complex');  % T is complex Schur form.
n = cast(size(A,1),coder.internal.indexIntClass);
R = eml_expand(coder.internal.scalarEg(T),size(T));
if isUTmatD(T) % Check if T is diagonal.
    % Square root always exists.
    for j = 1:n
        R(j,j) = sqrt(T(j,j));
    end
else
    % Compute upper triangular square root R of T, a column at a time.
    for j = 1:n
        R(j,j) = sqrt(T(j,j));
        % Changing branches in some cases improves accuracy and
        % reliability, but the resulting square root would not be the
        % principal one.
        % if imag(R(j,j)) < 0 && real(T(j,j)) < 0
        %     R(j,j) = -R(j,j);
        % end
        for i = coder.internal.indexMinus(j,1):-1:1
            s = coder.internal.scalarEg(R);
            for k = coder.internal.indexPlus(i,1):coder.internal.indexMinus(j,1)
                s = s + R(i,k)*R(k,j);
            end
            R(i,j) = (T(i,j) - s) / (R(i,i) + R(j,j));
        end
    end
end
X = Q*R*Q';
if nargout ~= 2
    % Calculate: nzeig = any(diag(T)==0);
    nzeig = false;
    for k = 1:n
        if T(k,k) == 0
            nzeig = true;
            if ~coder.setwarningflag
                coder.internal.warning('MATLAB:sqrtm:SingularMatrix');
            end
            break
        end
    end
    check_for_cancellation(R);
end
if nargout == 1
elseif nargout == 2
    arg2 = norm(X*X-A,'fro') / norm(A,'fro');
elseif nargout == 3
    arg2 = norm(R,'fro')^2 / norm(T,'fro');
    if nzeig
        condest = coder.internal.inf(class(A));
    else
        % Power method to get condition number estimate.
        tol = 1e-2;
        x = complex(ones(n*n,1,class(A)));    % Starting vector.
        cnt = ones(coder.internal.indexIntClass);
        e = ones(class(A));
        e0 = zeros(class(A));
        while abs(e-e0) > tol*e && cnt <= 6
            x = x/norm(x);
            x0 = x;
            e0 = e;
            Sx = tksolve(R,x);
            x = tksolveT(R,Sx);
            e = sqrt(real(dot(x0,x)));  % sqrt of Rayleigh quotient.
            cnt = coder.internal.indexPlus(cnt,1);
        end
        condest = e*norm(A,'fro')/norm(X,'fro');
    end
end
% As in FUNM:
if isreal(A) && norm(imag(X),1) <= 10*cast(n,class(X))*eps*norm(X,1)
    % Zero the imaginary parts.
    for j = 1:n
        for i = 1:n
            X(i,j) = real(X(i,j));
        end
    end
end

%--------------------------------------------------------------------------

function x = tksolve(R,b)
%TKSOLVE     Solves block triangular Kronecker system.
%            x = TKSOLVE(R, b) solves
%                  A*x = b
%            where A = KRON(EYE,R) + KRON(TRANSPOSE(R),EYE).
n = cast(size(R,1),coder.internal.indexIntClass);
x = coder.nullcopy(eml_expand(coder.internal.scalarEg(R,b),[coder.internal.indexTimes(n,n),1]));
temp = coder.nullcopy(eml_expand(coder.internal.scalarEg(R,b),[n,1]));
% Forward substitution.
for i = 1:n
    i0 = coder.internal.indexTimes(coder.internal.indexMinus(i,1),n);
    for j = 1:n
        temp(j) = b(coder.internal.indexPlus(i0,j));
    end
    for j = 1:coder.internal.indexMinus(i,1)
        j0 = coder.internal.indexTimes(coder.internal.indexMinus(j,1),n);
        for k = 1:n
            temp(k) = temp(k) - R(j,i)*x(coder.internal.indexPlus(j0,k));
        end
    end
    Rtemp = R;
    for j = 1:n
        Rtemp(j,j) = Rtemp(j,j) + R(i,i);
    end
    temp = linsolve(Rtemp,temp,struct('UT',true));
    for j = 1:n
        x(coder.internal.indexPlus(i0,j)) = temp(j);
    end
end

%--------------------------------------------------------------------------

function x = tksolveT(R,b)
%TKSOLVE     Solves block triangular Kronecker system.
%            x = TKSOLVET(R,b) solves
%                 A'*x = b
%            where A = KRON(EYE,R) + KRON(TRANSPOSE(R),EYE).
n = cast(size(R,1),coder.internal.indexIntClass);
x = coder.nullcopy(eml_expand(coder.internal.scalarEg(R,b),[coder.internal.indexTimes(n,n),1]));
temp = coder.nullcopy(eml_expand(coder.internal.scalarEg(R,b),[n,1]));
% Back substitution.
for i = n:-1:1
    i0 = coder.internal.indexTimes(coder.internal.indexMinus(i,1),n);
    for j = 1:n
        temp(j) = b(coder.internal.indexPlus(i0,j));
    end
    for j = coder.internal.indexPlus(i,1):n
        j0 = coder.internal.indexTimes(coder.internal.indexMinus(j,1),n);
        for k = 1:n
            temp(k) = temp(k) - conj(R(i,j))*x(coder.internal.indexPlus(j0,k));
        end
    end
    Rtemp = R';
    for j = 1:n
        Rtemp(j,j) = Rtemp(j,j) + conj(R(i,i));
    end
    temp = linsolve(Rtemp,temp,struct('LT',true));
    for j = 1:n
        x(coder.internal.indexPlus(i0,j)) = temp(j);
    end
end

%--------------------------------------------------------------------------

function p = isUTmatD(T)
% Returns true if upper triangular matrix T is diagonal, false otherwise.
n = cast(size(T,2),coder.internal.indexIntClass);
for j = 1:n
    for i = 1:coder.internal.indexMinus(j,1)
        if T(i,j) ~= 0
            p = false;
            return
        end
    end
end
p = true;

%--------------------------------------------------------------------------

function y = abs1(x)
y = abs(real(x)) + abs(imag(x));

%--------------------------------------------------------------------------

function check_for_cancellation(R)
% Warn if there is substantial cancellation of principal square roots of
% eigenvalues.
n = cast(size(R,2),coder.internal.indexIntClass);
for j = 1:coder.internal.indexMinus(n,1)
    for i = coder.internal.indexPlus(j,1):n
        if abs1(R(i,i)+R(j,j)) <=  ...
                10*sqrt(eps(class(R)))*(abs1(R(i,i))+abs1(R(j,j)))
            if ~coder.setwarningflag
                coder.internal.warning('Coder:toolbox:sqrtm_rootEigenCancel');
            end
            return
        end
    end
end

%--------------------------------------------------------------------------
