function varargout = scalaSvd(A, isvec)
% SCALASVD  ScaLAPACK singular value decomposition

%   Copyright 2006-2021 The MathWorks, Inc.

nargoutchk(0, 3);

% When there are no output arguments, the function should still return
% the singular values.
numOutputs = max(1, nargout);

sizeA = size(A);
minMN = min(size(A));

% TODO(g2636380): remove this once MKL ScaLAPACK is fixed.
% ScaLAPACK may return wrong answers for single on Windows (see
% g2607606). Thus we might cast to double and then we'll 
% cast back to single the outputs. The general case is we don't though.
needCastToDouble = false;

% Exit early for special inputs.
if minMN == 0
    % Empty input.

    if isvec
        sizeS = [0,1];
    else
        sizeS = [0,0];
    end
    % To match the "in memory" behaviour, the outputs should be real,
    % even if the input is complex. Note that A is guarenteed to be
    % empty so the call to REAL below should not be expensive.
    outputPrototype = real(A);
    if numOutputs == 1
        varargout{1} = zeros(sizeS, 'like', outputPrototype);
    else
        varargout{1} = eye(sizeA(1), 0, 'like', outputPrototype);
        varargout{2} = zeros(sizeS, 'like', outputPrototype);
        if numOutputs > 2
            varargout{3} = eye(sizeA(2), 0, 'like', outputPrototype);
        end
    end
    return;
elseif isequal(size(A), [1, 2]) && nargout == 3
    % ScaLAPACK stalls for 1x2 input matrix (see g2517708) and 3 outputs.
    % At this point we only compute 'econ' form for rectangular matrices.
    % So it is trivial to do the computation directly.
    varargout{1} = ones(1, 1, 'like', A);
    normA = norm(A);
    varargout{2} = normA;
    varargout{3} = A' ./ normA;
    return;
elseif isUnderlyingType(A, "single") && isreal(A) && ispc()
    % TODO(g2636380): remove this once MKL ScaLAPACK is fixed.
    % This is the case where we will be casting to double and then casting
    % back to single - real, single, and on Windows.
    A = double(A);
    needCastToDouble = true;
end

% Redistribute input according to the default 2dbc distribution scheme.
if ~isa(getCodistributor(A), 'codistributor2dbc')
    A = redistribute(A, codistributor2dbc());
end

aDist = getCodistributor(A);
lbgrid = aDist.LabGrid;


codistr1d = codistributor1d();

% Create array descriptor required by ScaLAPACK.
descA = arraydescriptor(aDist);

localA = hGetLocalPartWithCorrectComplexity(A);

% Output from ScalaPack is ordered by: singular values, left singular vectors, right singular vectors.
[svdOut{1:numOutputs}] = parallel.internal.codistextern.denseSvd(localA, ...
    descA, lbgrid(1), lbgrid(2), ...
    aDist.Orientation, isreal(A));

if needCastToDouble
    % TODO(g2636380): remove this once MKL ScaLAPACK is fixed.
    % We casted to double and we now need to cast back to
    % single.
    svdOut = cellfun(@single, svdOut, "UniformOutput", false);
end

% For singular values, the ScaLAPACK output is replicated, but we want a
% codistributed output. ScaLAPACK doesn't always produce exactly replicated
% results so use the values on worker 1.
S = codistributed.pConstructFromReplicated(svdOut{1}, 1, codistributor1d(1));
% Transform S into diagonal matrix if matrix form requested.
if ~isvec
    S = diag(S);
end

if (numOutputs == 1)
    varargout{1} = S;
    return;
end

% When there are multiple output arguments, they need to be rearranged.  The
% ScaLAPACK ordering is [S, U, VT] while MATLAB expects [U, S, VT'].
codistrU = codistributor2dbc(lbgrid, aDist.BlockSize, ...
    aDist.Orientation, [size(A,1) minMN]);
dU = codistributed.pDoBuildFromLocalPart(svdOut{2}, codistrU);
varargout{1} = redistribute(dU, codistr1d);
varargout{2} = S;

if (numOutputs == 3)
    codistrV = codistributor2dbc(lbgrid, aDist.BlockSize, ...
        aDist.Orientation, [minMN size(A,2)]);
    dVT = codistributed.pDoBuildFromLocalPart(svdOut{3}, codistrV);
    varargout{3} = redistribute(dVT, codistr1d)';
end
