function [B, rank, rcond] = sparseQrSolveStep(H, Beta, R, B, tol)
%sparseQrSolveStep Solve sparse least-square problem using outputs of QR.
%
%   This function takes as parameters the local part of the Householder
%   vectors H, and scalars Beta, the distributed R as returned by 
%   sparseQrFactorizationStep, as well as the tolerance to which rank 
%   deficiency is detected.
%
%   See also sparseQrFactorizationStep

%   Copyright 2020-2022 The MathWorks, Inc.

import distributedutil.TransposeFlags

% Detect if the original matrix is rank deficient.
[rank, rcond, imageR] = iRankDeficiencyDetection(R, tol);
kernelR = ~imageR; % If it's not in the image, it's in the kernel.
% Get the corresponding local indices of the image of R.
locImageR = getLocalPart(imageR);
m = size(H, 1);
n = size(R, 2);

% Compute Q'*B using the columns of H associated to the image of R.
B = iApplyQ(H(:, locImageR), Beta(locImageR), B, TransposeFlags.CONJUGATE_TRANSPOSE);
% If the problem is strictly overdetermined we truncate B because we have
% applied the full Q and not the economic one. Note that it's not trivial to
% apply the economic Q directly as we do not form Q explicitly.
if m > n
    B = B(1:n,:);
end

% Solve with the image part of R only.
B(imageR, :) = R(imageR, imageR) \ B(imageR, :);
% Fill with zeros for the kernel part, mimics in-memory.
B(kernelR, :) = 0;
end

function [rank, rcond, imageR] = iRankDeficiencyDetection(R, tol)
% Detect if R is numerically rank deficient, up to tol. Returns the rank
% and the columns of R associated to the image.

esv = abs(diag(R)); % estimated singular values
rcond = max(esv) / min(esv);
imageR = (esv > tol);
rank = nnz(imageR);
end

function B = iApplyQ(H, Beta, B, transposeFlagQ)
% If B is replicated we are able to reduce the overall communication,
% otherwise we start by redistributing B by columns, if needed, and we
% treat each column chunk at a time as a replicated B.

if iscodistributed(B)
    % We redistribute B if it is not 1d column codistributed.
    codistB = getCodistributor(B);
    if ~(isa(codistB, 'codistributor1d') && codistB.Dimension == 2)
        codistB = codistributor1d(2, codistributor1d.unsetPartition, size(B));
        B = redistribute(B, codistB);
    end
    B = iApplyQDistributedB(H, Beta, B, transposeFlagQ);
else % B is replicated
    B = iApplyQReplicatedB(H, Beta, B, transposeFlagQ);
end

end

function B = iApplyQDistributedB(H, Beta, B, transposeFlagQ)
% Apply Q or Q' when B is distributed. H and Beta are local parts of the
% Householder vectors and scalars distributed by columns. We assume B is
% distributed by columns to reduce communication.

% We loop on the workers in the same way as iApplyQReplicatedB to
% better overlap computation and communication.
if transposeFlagQ == distributedutil.TransposeFlags.NO_TRANSPOSE
    workerIndices = spmdSize:-1:1;
else
    workerIndices = 1:spmdSize;
end

for ii = 1:numel(workerIndices)
    workerIdx = workerIndices(ii);
    chunkB = spmdBroadcast(workerIdx, getLocalPart(B));
    % iApplyQReplicated must be called by all the workers, but only one
    % needs to store the output as it is its local part of B. Thus in the 
    % end each worker stores different local parts.
    if spmdIndex == workerIdx
        localB = iApplyQReplicatedB(H, Beta, chunkB, transposeFlagQ);
    else
        iApplyQReplicatedB(H, Beta, chunkB, transposeFlagQ);
    end
end

B = codistributed.build(localB, getCodistributor(B));

end

function B = iApplyQReplicatedB(H, Beta, B, transposeFlagQ)
% Apply Q or Q' when B is replicated. H and Beta are local parts of the
% Householder vectors and scalars distributed by columns.

if transposeFlagQ == distributedutil.TransposeFlags.NO_TRANSPOSE
    workerIndices = spmdSize:-1:1;
else
    workerIndices = 1:spmdSize;
end

pctTag = 2^15 - mod(sum(mfilename),2^15);
% Compte Q*B or Q'*B
for idx = 1:numel(workerIndices)-1
    thisWorker = workerIndices(idx);
    nextWorker = workerIndices(idx+1);
    if spmdIndex == thisWorker
        B = iApplyLocalQ(H, Beta, B, transposeFlagQ);
        spmdSend(B, nextWorker, pctTag + thisWorker);
    elseif spmdIndex == nextWorker
        B = spmdReceive(thisWorker, pctTag + thisWorker);
    end
end
if spmdIndex == workerIndices(end)
    B = iApplyLocalQ(H, Beta, B, transposeFlagQ);
end

B = spmdBroadcast(workerIndices(end), B);

end

function localB = iApplyLocalQ(H, Beta, localB, transposeFlagQ)
[~, n] = size(H);

if transposeFlagQ == distributedutil.TransposeFlags.NO_TRANSPOSE
    colIndices = n:-1:1;
else
    colIndices = 1:n;
end

for j = colIndices
    h = H(:,j);
    localB = localB - h * (Beta(j) * (h' * localB));
end
end
