function localA = replaceZerosOnDiagonalWithNaNs(localA, diagOffset)
% Helper to replace every zero on the diagonal of a local part with NaN. 
% TODO(g1986555) Remove helper and its usage once fixed.

%   Copyright 2019-2020 The MathWorks, Inc.

assert(issparse(localA), "localA must be sparse.");

sizeLocalA = size(localA);
[minSizeLocalA, dim] = min(sizeLocalA);

if isvector(localA)
    diagA = localA(diagOffset + 1);
else
    % Positive offset if dim == 1, but negative otherwise.
    diagA = diag(localA, (2*(dim == 1) - 1)*diagOffset);
end
zeroDiagInd = (diagA == 0);
if any(zeroDiagInd)
    correction = zeros(size(diagA, 1), size(diagA, 2));
    correction(zeroDiagInd) = NaN;
    % We need to shift row/col indices according to dim.
    idx = 1:minSizeLocalA;
    if dim == 1
        idxI = idx;
        idxJ = idx + diagOffset;
    else % dim == 2
        idxI = idx + diagOffset;
        idxJ = idx;
    end
    localA = localA + sparse(idxI, idxJ, correction, sizeLocalA(1), sizeLocalA(2));
end
end
