function varargout = hIluColWise(codistrA, localA, setup)
%hIluColWise incomplete LU of a sparse matrix distributed column wise.
%
%   LU = hIluColWise(CODISTRA, LOCALA, SETUP) returns "L+U-speye(size(A))"
%
%   [L,U] = hIluColWise(CODISTRA, LOCALA, SETUP) returns unit lower
%   triangular L and upper triangular U

%   Copyright 2020-2022 The MathWorks, Inc.

partitionA = codistrA.Partition;
blockLimits = cumsum(partitionA);
blockLimits = [0 blockLimits(1:end-1); blockLimits]';
firstDiag = blockLimits(spmdIndex, 1);
milu = false;
if isfield(setup, 'milu')
    assert(strcmpi(setup.milu, "off") || strcmpi(setup.milu, "col"), ...
        "setup.milu must be 'off' or 'col' for hIluColImpl");
    if strcmp(setup.milu, 'col')
        milu = true;
    end
end

% Get the global column indices needed to perform the update
upidx = parallel.internal.codistextern.sparseNnzPerRowOrColumn(localA, 1);
upidx = upidx(1:firstDiag);
% Gather the information as a cell array of sparse vectors, the overhead is
% assumed to be acceptable because there are generally few columns needed
% to do the update
upidx = sparse(upidx);
upidx = spmdCat({upidx});

myTagSeed = iTagSeedGenerator();
[sendsToMe, receivesFromMe] = iCreateCommunicationPattern(upidx, blockLimits);
[nrow, ncol] = size(localA);
update = zeros(0, 0, 'like', localA);
% Main loop
for workerIdx = 1:spmdSize
    if spmdIndex == workerIdx
        [localL,localU] = parallel.internal.codistextern.sparseIluFactorizeLocalPart(localA, update, 'nofill', milu, firstDiag);
        % Send only the columns needed to workers on the right
        for iRecvr = workerIdx+1:spmdSize
            if receivesFromMe(iRecvr)
                idx = find(upidx{iRecvr}(blockLimits(spmdIndex,1)+1:blockLimits(spmdIndex,2)));
                updateToSend = sparse(nrow, ncol);
                % We ensured that idx ~= [], and column indexing is not
                % decreasing performance here.
                updateToSend(:, idx) = localL(:, idx); %#ok<SPRIX>
                spmdSend(updateToSend, iRecvr, iTagGenerator(myTagSeed, workerIdx));
            end
        end
    elseif spmdIndex > workerIdx && sendsToMe(workerIdx)
        workerUp = spmdReceive(workerIdx, iTagGenerator(myTagSeed, workerIdx));
        % Pre-allocating would decrease the readability for minor
        % performance improvement, if any.
        update = [update workerUp]; %#ok<AGROW>
    elseif spmdIndex > workerIdx && ~sendsToMe(workerIdx)
        % Pre-allocating would decrease the readability for minor
        % performance improvement, if any.
        update = [update sparse(nrow, blockLimits(workerIdx, 2)-blockLimits(workerIdx, 1))]; %#ok<AGROW>
    end
end

% Prepare output
if nargout == 1
    % Compact form, U already owns the diagonal entries so we discard L
    % diagonal entries
    varargout = {tril(localL, -firstDiag-1) + localU};
else
    varargout = {localL, localU};
end
end

function [sendsToMe, receivesFromMe] = iCreateCommunicationPattern(upidx, blockLimits)
% Purely local function that creates 2 boolean arrays sendsToMe and
% receivesFromMe, both return true if worker is a sender/receiver of
% spmdIndex and false otherwise.

sendsToMe = false(spmdSize, 1);
receivesFromMe = false(spmdSize, 1);
% Receivers, only on the right
for workerIdx = spmdIndex+1:spmdSize
    receivesFromMe(workerIdx) = nnz(upidx{1,workerIdx}(blockLimits(spmdIndex,1)+1:blockLimits(spmdIndex,2))) ~= 0;
end
% Senders, only on the left
for workerIdx = 1:spmdIndex-1
    sendsToMe(workerIdx) = nnz(upidx{1,spmdIndex}(blockLimits(workerIdx,1)+1:blockLimits(workerIdx,2))) ~= 0;
end
end

function myTag = iTagGenerator(myTagSeed, msgId)
% Helper function to generate tags for messages, requires a tag seed which
% is supposed to be unique to this file

myTag = myTagSeed + mod(msgId, 2^15);
end

function myTagSeed = iTagSeedGenerator()
% Helper function to create a tag seed for this file

myTagSeed = mod(sum(mfilename), 2^15);
end
