function [LPC, codistrC] = hMtimesImpl(~, codistrA, LPA, transA, ...
    codistrB, LPB, transB)
%hMtimesImpl mtimes for codistributor1d when both matrices are codistributor1d

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

%   Copyright 2010-2023 The MathWorks, Inc.

import distributedutil.TransposeFlags

% If needed, transpose the codistributor
if transA ~= TransposeFlags.NO_TRANSPOSE
    codistrA = codistrA.pTransposeCodistributor();
end
if transB ~= TransposeFlags.NO_TRANSPOSE
    codistrB = codistrB.pTransposeCodistributor();
end
[codistrA, LPA] = pRedistribute2DMatrixToDimOneOrTwo(codistrA, LPA);
[codistrB, LPB] = pRedistribute2DMatrixToDimOneOrTwo(codistrB, LPB);
numRowsC = codistrA.Cached.GlobalSize(1);
numColsC = codistrB.Cached.GlobalSize(2);
sizeC = [numRowsC numColsC];
distDims = [codistrA.Dimension, codistrB.Dimension];

% These different combinations of distributions have different
% implementations
if isequal(distDims, [1, 1])
    [LPC, codistrC] = iRowsTimesRows(sizeC, codistrA, LPA, transA, ...
        codistrB, LPB, transB);
elseif isequal(distDims, [1, 2])
    [LPC, codistrC] = iRowsTimesColumns(sizeC, codistrA, LPA, transA, ...
        codistrB, LPB, transB);
elseif isequal(distDims, [2, 1])
    [LPC, codistrC] = iColumnsTimesRows(sizeC, codistrA, LPA, transA, ...
        codistrB, LPB, transB);
elseif isequal(distDims, [2, 2])
    [LPC, codistrC] = iColumnsTimesColumns(sizeC, codistrA, LPA, transA, ...
        codistrB, LPB, transB);
else
    ME = MException(message(...
        "parallel:codistributors:c1dMtimesHighDistDimNotSupported"));
    throwAsCaller(ME);
end

end  % End of hMtimesImpl

function [LPC, codistrC] = iRowsTimesRows(sizeC, codistrA, LPA, transA, ...
    codistrB, LPB, transB)
% Rows * rows -> rows.
% C(i, :) = sum over k A(i, k)*B(k, :) where A(i, k) is a submatrix
% of LPA on the k-th worker and 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, codistrA.Partition, sizeC);
LPC = zeros(codistrC.hLocalSize(), "like", LPA([])*LPB([]));%distributedutil.Allocator.create(codistrC.hLocalSize(), LPA([])*LPB([]));
k = iGetGlobalIndicesAsColonObj(codistrB, 1);
% Define the slicing: we want column slicing of LPA here
sliceOp = i1dSlicing(transA, true);
LPC(:, :) = mtimesOp(sliceOp(LPA, k), LPB);
to = mod(spmdIndex-2, spmdSize)+1;
from = mod(spmdIndex, spmdSize)+1;
mwTag1 = 31639;
for p = [spmdIndex+1:spmdSize 1:spmdIndex-1]
    LPB = spmdSendReceive(to, from, LPB, mwTag1);
    k = iGetGlobalIndicesAsColonObj(codistrB, 1, p);
    LPC = LPC + mtimesOp(sliceOp(LPA, k), LPB);
end
end

function [LPC, codistrC] = iRowsTimesColumns(sizeC, codistrA, LPA, transA, ...
    codistrB, LPB, transB)
% Rows * columns -> columns.
% No partial sums required.
% C(i, j) = A(i, :)*B(:, j) where A(i, :) and B(:, j)
% are the entire local portions on the i-th and j-th workers
% Send A(i, :)'s around a ring of workers.

% Get the corresponding local mtimes operator
mtimesOp = buildMtimesTransposeOp(transA, transB);
codistrC = codistributor1d(2, codistrB.Partition, sizeC);
LPC = zeros(codistrC.hLocalSize(), "like", LPA([])*LPB([]));
ii = iGetGlobalIndicesAsColonObj(codistrA, 1);
LPC(ii, :) = mtimesOp(LPA, LPB);
to = mod(spmdIndex-2, spmdSize)+1;
from = mod(spmdIndex, spmdSize)+1;
mwTag2 = 31640;
for p = [spmdIndex+1:spmdSize 1:spmdIndex-1]
    LPA = spmdSendReceive(to, from, LPA, mwTag2);
    ii = iGetGlobalIndicesAsColonObj(codistrA, 1, p);
    LPC(ii, :) = mtimesOp(LPA, LPB);
end
end

function [LPC, codistrC] = iColumnsTimesRows(sizeC, codistrA, LPA, transA, ...
    codistrB, LPB, transB)
% Columns * rows
% Two sub-cases:
% numCols < numRows -> distribute result by rows.
% numCols >= numRows -> distribute result by columns.
% Reduction over distributed dimensions.
% Make sure size(LPA, 2) == size(LPB, 1), then
% compute T = LPA*LPB, LP = codistributed(spmdPlus(T))
% without actually using spmdPlus and codistributed(...).
% Message pattern is all-to-all to distribute and sum T's.

% Get the corresponding local mtimes operator
mtimesOp = buildMtimesTransposeOp(transA, transB);
Apart = codistrA.Partition;
Bpart = codistrB.Partition;
if ~isequal(Apart, Bpart)
    origCodistrB = codistrB;
    codistrB = codistributor1d(1, Apart, codistrB.Cached.GlobalSize);
    % Revert implicit transpose to redistribute correctly
    if transB ~= distributedutil.TransposeFlags.NO_TRANSPOSE
        origCodistrB = origCodistrB.pTransposeCodistributor();
        codistrB = codistrB.pTransposeCodistributor();
    end
    LPB = distributedutil.Redistributor.redistribute(origCodistrB, LPB, codistrB);
    % We do not transpose back codistrB here because it is not used anymore
    % afterwards.
end
mwTag3 = 31641;
if sizeC(2) < sizeC(1)
    % Result distributed by rows
    Cpart = codistributor1d.defaultPartition(sizeC(1));
    codistrC = codistributor1d(1, Cpart, sizeC);
    LPC = zeros(codistrC.hLocalSize(), "like", LPA([])*LPB([]));
    ii = iGetGlobalIndicesAsColonObj(codistrC, 1);
    % Define the slicing: we want row slicing of LPA here
    sliceOp = i1dSlicing(transA, false);
    LPC(:, :) = mtimesOp(sliceOp(LPA, ii), LPB);
    for p = 1:spmdSize-1
        to = mod(spmdIndex+p-1, spmdSize)+1;
        from = mod(spmdIndex-p-1, spmdSize)+1;
        ii = iGetGlobalIndicesAsColonObj(codistrC, 1, to);
        myLPC = mtimesOp(sliceOp(LPA, ii), LPB);
        LPC = LPC + spmdSendReceive(to, from, myLPC, mwTag3);
    end
else
    % Result distributed by columns
    Cpart = codistributor1d.defaultPartition(sizeC(2));
    codistrC = codistributor1d(2, Cpart, sizeC);
    LPC = zeros(codistrC.hLocalSize(), "like", LPA([])*LPB([]));
    jj = iGetGlobalIndicesAsColonObj(codistrC, 2);
    % Define the slicing: we want column slicing of LPB here
    sliceOp = i1dSlicing(transB, true);
    LPC(:, :) = mtimesOp(LPA, sliceOp(LPB, jj));
    for p = 1:spmdSize-1
        to = mod(spmdIndex+p-1, spmdSize)+1;
        from = mod(spmdIndex-p-1, spmdSize)+1;
        jj = iGetGlobalIndicesAsColonObj(codistrC, 2, to);
        myLPC = mtimesOp(LPA, sliceOp(LPB, jj));
        LPC = LPC + spmdSendReceive(to, from, myLPC, mwTag3);
    end
end
end

function [LPC, codistrC] = iColumnsTimesColumns(sizeC, codistrA, LPA, transA, ...
    codistrB, LPB, transB)
% Columns * columns -> columns
% (Includes default matrix distributions.)
% C(:, j) = sum over k A(:, k)*B(k, j) where A(:, k) is the
% entire local portion of A on the k-th worker and B(k, j) is
% a submatrix of the local portion of B 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, codistrB.Partition, sizeC);
LPC = zeros(codistrC.hLocalSize(), "like", LPA([])*LPB([]));
k = iGetGlobalIndicesAsColonObj(codistrA, 2);
% Define the slicing: we want row slicing of LPB here
sliceOp = i1dSlicing(transB, false);
LPC(:, :) = mtimesOp(LPA, sliceOp(LPB, k));
to = mod(spmdIndex-2, spmdSize)+1;
from = mod(spmdIndex, spmdSize)+1;
mwTag4 = 31642;
for p = [spmdIndex+1:spmdSize 1:spmdIndex-1]
    LPA = spmdSendReceive(to, from, LPA, mwTag4);
    k = iGetGlobalIndicesAsColonObj(codistrA, 2, p);
    LPC = LPC + mtimesOp(LPA, sliceOp(LPB, k));
end
end

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

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

end

function colonObj = iGetGlobalIndicesAsColonObj(codist, dim, varargin)
% Returns corresponding codist.globalIndices as ColonDescriptor, thus avoiding
% explicitly creating any array.

[first, last] = globalIndices(codist, dim, varargin{:});
colonObj = matlab.internal.ColonDescriptor(first, last);
end
