function varargout = scalaLu(A, permIsMat)
%SCALALU  ScaLAPACK LU factorization

%   Copyright 2006-2022 The MathWorks, Inc.

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

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

localA = hGetLocalPartWithCorrectComplexity(A);

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

[localA, localPerm] = parallel.internal.codistextern.denseLuFactor(localA, descA, ...
    lbgrid(1), lbgrid(2), aDist.Orientation);
A = codistributed.pDoBuildFromLocalPart(localA, aDist);

% A is the packed LU factor at this point, which we return in case of
% single output.
if nargout==1
    varargout{1} = A;
    return;
end

% Post process the permutation vector:
% We need P even in the 2 output case in order to permute L correctly.
processCol = aDist.hSpmdIndexToProcessorCol(spmdIndex);
localPerm = localPerm(1:size(localA, 1), processCol == 1);
codistr = codistributor2dbc(lbgrid, aDist.BlockSize,  ...
    aDist.Orientation, [size(A,1) 1]);
vecPerm = codistributed.pDoBuildFromLocalPart(localPerm, codistr);
vecPerm = gather(vecPerm);
vecPerm = iPermVector(vecPerm(1:min(size(A))), size(A, 1));

% Post process L and U:
% subsref is currently only supported on 1D, so we have no choice but to
% convert into 1D at this point to set up the L and U factors from the
% packed LU factor
LU = redistribute(A, codistributor1d(2));
L = tril(LU, -1) + eye(size(LU), codistributor1d(2));
if size(L, 1) < size(L, 2)
    L = subsref(L, substruct('()', {':',1:size(L,1)}));
end
U = triu(LU);
if size(U, 1) > size(U, 2)
    U = subsref(U, substruct('()', {1:size(U,2),':'}));
end

varargout{2} = U;
if nargout==2
    % 2 output case returns L(vecP', :) as first input
    [~, vecPermTranpose] = sort(vecPerm);
    varargout{1} = subsref(L, substruct('()', {vecPermTranpose, ':'}));
else % 3 output case
    varargout{1} = L;
    if permIsMat
        % Permutation matrix P = E(vecP, :) with E identity matrix of 
        % dimension size(A, 1). We cannot use eye(..., 'like', A) here 
        % since that would create a 2dbc distributed array.
        matPerm = codistributed.eye(size(A, 1), underlyingType(A));
        varargout{3} = subsref(matPerm, substruct('()', {vecPerm, ':'}));
    else
        varargout{3} = vecPerm;
    end
end
end

function vecPerm = iPermVector(scalaVecPerm, n)
% Convert ScaLAPACK output to 1d codistributed MATLAB permutation vector
vecPerm = 1:n;
for i = 1:length(scalaVecPerm)
    swap = vecPerm(scalaVecPerm(i));
    vecPerm(scalaVecPerm(i)) = vecPerm(i);
    vecPerm(i) = swap;
end
vecPerm = codistributed.pConstructFromReplicated(vecPerm, codistributor1d());
end
