function localL = hIcholColWise(codistrA, localA, setup, isMatrixKnownSymmetric)
%hIcholColWise incomplete Cholesky of a sparse matrix distributed column wise.
%
%   L = hIcholColWise(CODISTRA, LOCALA, SETUP)

%   Copyright 2020-2022 The MathWorks, Inc.

partitionA = codistrA.Partition;
blockLimits = cumsum(partitionA);
blockLimits = [0 blockLimits(1:end-1); blockLimits]';
firstDiag = blockLimits(spmdIndex, 1);
michol = strcmpi(setup.michol, "on");

% Check that localA is not structurally singular, i.e., if a diagonal value
% is exactly 0. Meanwhile, we apply the diagonal compensation if requested.
isStructurallySingular = false;
if iscolumn(localA)
    if localA(firstDiag+1) == 0
        isStructurallySingular = true;
    elseif isfield(setup, "diagcomp")
        localA(firstDiag+1) = localA(firstDiag+1) * (1 + setup.diagcomp);
    end
elseif ~isempty(localA)
    diagA = diag(localA, -firstDiag);
    if nnz(diagA) < size(localA, 2)
        isStructurallySingular = true;
    elseif isfield(setup, "diagcomp")
        localA(blockLimits(spmdIndex, 1)+1:blockLimits(spmdIndex, 2), :) = ...
            localA(blockLimits(spmdIndex, 1)+1:blockLimits(spmdIndex, 2), :) + diag(setup.diagcomp*diagA);
    end
end
% Collectively throw the error if A is structurally singular
isStructurallySingular = spmdReduce(@or, isStructurallySingular);
if isStructurallySingular
    error(message("MATLAB:ichol:DiagZeroColZeroBreakdown"));
end

% For each row of localA we are going to need to know if there is a
% non-zero entry.
upidx = parallel.internal.codistextern.sparseNnzPerRowOrColumn(localA, 1);
% Now we can compute the full communication pattern, as A is symmetric we
% can do this without communicating at all with the other workers.
[sendsToMe, receivesFromMe] = iCreateCommunicationPattern(upidx, blockLimits);
% We make sure the communication pattern is consistent. It must be if the
% input is a symmetric matrix, but the user may have given an unsymmetric
% matrix... And if we don't do this check, we would have a dead-lock during
% the computation of the factorisation. Of course, we filter-out the case
% where the matrix is known to be symmetric.
if ~isMatrixKnownSymmetric
    % We concatenate the sendsToMe vectors by columns, each vector is of
    % size spmdSize.
    allSendsToMe = spmdCat(sendsToMe);
    % allSendsToMe(spmdIndex, :) indicates the workers that send to spmdIndex.
    mismatch = ~isequal(receivesFromMe.', allSendsToMe(spmdIndex, :));
    mismatch = spmdReduce(@or, mismatch);
    if mismatch
        error(message("parallel:distributed:IcholNotSymmetric"));
    end
end

[nrow, ncol] = size(localA);
update = zeros(0, 0, 'like', localA);

pctTag = 2^15 - mod(sum(mfilename),2^15);

% Main loop
for workerIdx = 1:spmdSize
    if spmdIndex == workerIdx
        try
            [localL,localU] = parallel.internal.codistextern.sparseIluFactorizeLocalPart(localA, update, 'nofill', michol, firstDiag);
        catch ME
            % If there was a zero pivot we send the ICHOL warning.
            if strcmp(ME.identifier, "MATLAB:ilu:ZeroPivot")
                error(message("MATLAB:ichol:Breakdown"));
            else
                throw(ME);
            end
        end
        % Send only the columns needed to workers on the right
        for iRecvr = workerIdx+1:spmdSize
            if receivesFromMe(iRecvr)
                % The computation is morally the same than that of ILU but
                % as A is symmetric we can compute the column indices
                % needed by iRecvr using localA directly.
                idx = find(parallel.internal.codistextern.sparseNnzPerRowOrColumn(localA(blockLimits(iRecvr,1)+1:blockLimits(iRecvr,2),:),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, pctTag);
            end
        end
    elseif spmdIndex > workerIdx && sendsToMe(workerIdx)
        workerUp = spmdReceive(workerIdx, pctTag);
        % 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

% Rescale the output accordingly.
if iscolumn(localU)
    d = localU(firstDiag + 1);
else
    d = diag(localU, -firstDiag);
end
if ~isfield(setup, "shape") || (isfield(setup, "shape") && strcmpi(setup.shape, "lower"))
    localL = localL * diag(sqrt(d));
else % setup.shape == "upper"
    d = spmdCat(d, 1);
    localL = localU ./ sqrt(d);
end
end

function [sendsToMe, receivesFromMe] = iCreateCommunicationPattern(upidx, blockLimits)
% Function to create 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(blockLimits(workerIdx,1)+1:blockLimits(workerIdx,2))) ~= 0;
end
% Senders, only on the left
for workerIdx = 1:spmdIndex-1
    sendsToMe(workerIdx) = nnz(upidx(blockLimits(workerIdx,1)+1:blockLimits(workerIdx,2))) ~= 0;
end
end
