function [LPC, codistrC] = hMtimesReplicatedImpl(~, codistrA, LPA, transA, ...
    codistrB, LPB, transB)
%hMtimesReplicatedImpl mtimes for codistributor1d when one matrix is replicated

%   There are 4 combinations of distributions of the input matrices:
%     - rows * replicated
%     - columns * replicated
%     - replicated * rows
%     - replicated * columns
%   Each combination has its own implementation which lives in the
%   corresponding iXXXTimesXXX function.

%   Copyright 2010-2022 The MathWorks, Inc.

import distributedutil.TransposeFlags

if isempty(codistrB)  % A distributed, B replicated
    [codistrA, LPA] = pRedistribute2DMatrixToDimOneOrTwo(codistrA, LPA);
    if transA ~= TransposeFlags.NO_TRANSPOSE
        codistrA = codistrA.pTransposeCodistributor();
    end
    numRowsC = codistrA.Cached.GlobalSize(1);
    if transB == TransposeFlags.NO_TRANSPOSE
        numColsC = size(LPB, 2);
    else
        numColsC = size(LPB, 1);
    end
    sizeC = [numRowsC numColsC];
    if codistrA.Dimension == 1
        [LPC, codistrC] = iRowsTimesReplicated(sizeC, codistrA, LPA, transA, ...
            LPB, transB);
    elseif codistrA.Dimension == 2
        [LPC, codistrC] = iColumnsTimesReplicated(sizeC, codistrA, LPA, transA, ...
            LPB, transB);
    else
        ME = MException(message(...
            'parallel:codistributors:c1dMtimesHighDistDimNotSupported'));
        throwAsCaller(ME);
    end
else   % A replicated, B distributed
    [codistrB, LPB] = pRedistribute2DMatrixToDimOneOrTwo(codistrB, LPB);
    if transA == TransposeFlags.NO_TRANSPOSE
        numRowsC = size(LPA, 1);
    else
        numRowsC = size(LPA, 2);
    end
    if transB ~= TransposeFlags.NO_TRANSPOSE
        codistrB = codistrB.pTransposeCodistributor();
    end
    numColsC = codistrB.Cached.GlobalSize(2);
    sizeC = [numRowsC numColsC];
    if codistrB.Dimension == 1
        [LPC, codistrC] = iReplicatedTimesRows(sizeC, LPA, transA, ...
            codistrB, LPB, transB);
    elseif codistrB.Dimension == 2
        [LPC, codistrC] = iReplicatedTimesColumns(sizeC, LPA, transA, ...
            codistrB, LPB, transB);
    else
        ME = MException(message(...
            'parallel:codistributors:c1dMtimesHighDistDimNotSupported'));
        throwAsCaller(ME);
    end
end

end  % End of hMtimesReplicatedImpl

function [LPC, codistrC] = iRowsTimesReplicated(sizeC, codistrA, LPA, transA, ...
    LPB, transB)
% Rows * replicated -> rows.
% Completely local operation, LPC = LPA*LPB.

% Get the corresponding local mtimes operator
mtimesOp = buildMtimesTransposeOp(transA, transB);
codistrC = codistributor1d(1, codistrA.Partition, sizeC);
LPC = distributedutil.Allocator.create(codistrC.hLocalSize(), LPA([])*LPB([]));
LPC(:, :) = mtimesOp(LPA, LPB);
end

function [LPC, codistrC] = iColumnsTimesReplicated(sizeC, codistrA, LPA, transA, ...
    LPB, transB)
% Columns * replicated -> columns.
% LPC = sum over k A(:, k)*B(k, j) where B(k, j) is a submatrix of LPB
% whose rows correspond to the columns of LPA stored on k-th worker and
% columns correspond to the columns of LPA stored locally, A(:, k) is
% the entire LPA on the k-th worker. Send A(:, k)'s around a ring of workers.

% Get the corresponding local mtimes operator
mtimesOp = buildMtimesTransposeOp(transA, transB);
codistrC = codistributor1d(2, codistributor1d.unsetPartition, sizeC);
LPC = distributedutil.Allocator.create(codistrC.hLocalSize(), LPA([])*LPB([]));
i = codistrA.globalIndices(2);
j = codistrC.globalIndices(2);
% Define the slicing: we want slicing of LPB here
sliceOp = i2dSlicing(transB);
LPC(:, :) = mtimesOp(LPA, sliceOp(LPB, i, j));
to = mod(spmdIndex-2, spmdSize)+1;
from = mod(spmdIndex, spmdSize)+1;
mwTag1 = 30618;
for p = [spmdIndex+1:spmdSize 1:spmdIndex-1]
    LPA = spmdSendReceive(to, from, LPA, mwTag1);
    k = codistrA.globalIndices(2, p);
    LPC = LPC + mtimesOp(LPA, sliceOp(LPB, k, j));
end
end

function [LPC, codistrC] = iReplicatedTimesRows(sizeC, LPA, transA, ...
    codistrB, LPB, transB)
% Replicated * rows -> rows.
% LPC = sum over k A(i, k)*B(k, :) where A(i, k) is a submatrix of LPA
% whose rows correspond to the rows of LPB stored locally and
% columns correspond to the rows of LPB stored on the k-th worker, B(k, :)
% is the entire LPB on the k-th worker. Send B(k, :)'s around a ring of workers.

% Get the corresponding local mtimes operator
mtimesOp = buildMtimesTransposeOp(transA, transB);
codistrC = codistributor1d(1, codistributor1d.unsetPartition, sizeC);
LPC = distributedutil.Allocator.create(codistrC.hLocalSize(), LPA([])*LPB([]));
k = codistrB.globalIndices(1);
j = codistrC.globalIndices(1);
% Define the slicing: we want slicing of LPA here
sliceOp = i2dSlicing(transA);
LPC(:, :) = mtimesOp(sliceOp(LPA, j, k), LPB);
to = mod(spmdIndex-2, spmdSize)+1;
from = mod(spmdIndex, spmdSize)+1;
mwTag2 = 30619;
for p = [spmdIndex+1:spmdSize 1:spmdIndex-1]
    LPB = spmdSendReceive(to, from, LPB, mwTag2);
    k = codistrB.globalIndices(1, p);
    LPC = LPC + mtimesOp(sliceOp(LPA, j, k), LPB);
end
end

function [LPC, codistrC] = iReplicatedTimesColumns(sizeC, LPA, transA, ...
    codistrB, LPB, transB)
% Replicated * columns -> columns.
% Completely local operation, LPC = LPA*LPB.

% Get the corresponding local mtimes operator
mtimesOp = buildMtimesTransposeOp(transA, transB);
codistrC = codistributor1d(2, codistrB.Partition, sizeC);
LPC = distributedutil.Allocator.create(codistrC.hLocalSize(), LPA([])*LPB([]));
LPC(:, :) = mtimesOp(LPA, LPB);
end

function sliceOp = i2dSlicing(transA)
% Returns a 2d slicing function depending on transA.
% In order to call in-memory fused mtimes the slicing needs to be done
% BEFORE any local transposition, e.g., A(j, i)'*B rather than At = tranpose(A)
% and At(i, j)*B.

if transA == distributedutil.TransposeFlags.NO_TRANSPOSE
    sliceOp = @(A, i, j) A(i, j);
else
    sliceOp = @(A, i, j) A(j, i);
end
end
