function varargout = scalaQr(A, varargin)
%SCALAQR   ScaLAPACK QR factorization

%   Copyright 2010-2023 The MathWorks, Inc.

distA = getCodistributor(A);
use1DForResult = isa(distA, "codistributor1d");

% Redistribute input with non-default 2dbc LabGrid such that ncol==1 (g2642926).
if ~isa(distA, "codistributor2dbc") || ...
        (isa(distA, "codistributor2dbc") && ~isequal(distA.LabGrid, [spmdSize 1]))
    A = redistribute(A, codistributor2dbc([spmdSize 1]));
end

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

% Ensure complexity matches on all workers
localA = hGetLocalPartWithCorrectComplexity(A);

% Set up the array descriptor required by ScaLAPACK
descA = arraydescriptor(distA);

[mA, nA] = size(A);

% Set flag for "economy" factorization
if mA > nA
    wantEconQR = ismember('econ', varargin);
else
    wantEconQR = false;
end

% Set flag for pivot output
wantPivots = nargout == 3;

if wantPivots
    [localQ, localA, localE] = parallel.internal.codistextern.denseQrWithPivoting(localA, ...
        descA, lbgrid(1), lbgrid(2), ...
        distA.Orientation, wantEconQR);
    E = generatePermVector(distA, localE);
    if ismember('matrix', varargin)
        % E is by construction 1d codistributed
        E = full(sparse(E, 1:length(E), 1));

        % The permutation matrix needs to be cast to the correct output
        % type.  (If the permutation output is a vector, it is always
        % double by design so no need to cast). See g849944.
        E = cast(E, underlyingType(A));
    end
    varargout{3} = E;
else
    [localQ, localA] = parallel.internal.codistextern.denseQr(localA, descA, ...
        lbgrid(1), lbgrid(2), ...
        distA.Orientation, ...
        wantEconQR);
end

if wantEconQR  % economy factorization
    [Q, A] = iConstructFactorsFromLPs(localQ, distA, localA, nA);
else
    [A, Q] = iConstructFactorsFromLPs(localA, distA, localQ, mA);
end

A = triu(A);  % Get the upper triangular factor

if use1DForResult
    Q = redistribute(Q, codistributor1d());
    A = redistribute(A, codistributor1d());
elseif wantPivots % Redistribute E to default 2dbc codistribution
    varargout{3} = redistribute(varargout{3}, codistributor2dbc());
end

varargout{1} = Q;
varargout{2} = A;
end

%%%%%%%%%%%%%%%%%%%%%%%%%
function [A, B] = iConstructFactorsFromLPs(localA, codistrA, localB, szB)
% First array reuses codistrA because it is the same size as the original
% input.  The second array will be a square array that has a codistributor
% with the properties of codistrA but global size szB x szB

A = codistributed.pDoBuildFromLocalPart(localA, codistrA);

% New codistributor for square factor
codistr = codistributor2dbc(codistrA.LabGrid, codistrA.BlockSize, ...
    codistrA.Orientation, [szB szB]);

% Dispose of any workspace remaining in the upper rows and columns
localSz = codistr.hLocalSize();
localB = localB(1:localSz(1), 1:localSz(2));
B = codistributed.pDoBuildFromLocalPart(localB, codistr);
end
