function C = mtimesImpl(A, B, transA, transB)
% mtimesImpl: Entrance point of mtimes implementations for codistributed
% inputs that checks inputs and calls appropriate implementation depending
% on used codistributor

%   Copyright 2016-2022 The MathWorks, Inc.

import distributedutil.TransposeFlags

narginchk(2, 4);

distributedutil.CodistParser.verifyNonCodistributedInputs({A, B});

% Set transpose flags if not provided
if nargin == 2
    transA = TransposeFlags.NO_TRANSPOSE;
    transB = TransposeFlags.NO_TRANSPOSE;
elseif nargin == 3
    transB = TransposeFlags.NO_TRANSPOSE;
end

if isscalar(A) || isscalar(B)
    A = iMaybeTranspose(A, transA);
    B = iMaybeTranspose(B, transB);
    C = times(A, B);
    return
end

[A, B] = iValidateInputArgs(A, B, transA, transB);

areBothMatricesDense = ~(issparse(A) || issparse(B));

isCodistributedA = isa(A, 'codistributed');
isCodistributedB = isa(B, 'codistributed');

% If either input is non-numeric we must go through the MATLAB
% implementation that is defined for 1D distributions (Scalapack only
% supports floating-point). 
if ~isfloat(A) || ~isfloat(B)
    A = iEnsure1DDist(A);
    B = iEnsure1DDist(B);
end

if isCodistributedA && isCodistributedB
    % Both A and B are codistributed arrays.
    codistrA = getCodistributor(A);
    codistrB = getCodistributor(B);
    if ~strcmpi(class(codistrA), class(codistrB))
        % A & B are of different distribution classes
        % Redistribute to the same class before proceeding
        [A, B] = iUnifyCodistributors(A, B, areBothMatricesDense);
    end
    if ~codistrA.hMtimesCheck(areBothMatricesDense)   % Both are distributed in same class
        A = iConvertTo1d(A);
        B = iConvertTo1d(B);
    end
    
    [localA, codistrA] = iExtractCodistrAndLP(A);
    [localB, codistrB] = iExtractCodistrAndLP(B);
    
    % Shortcut if the result is a scalar
    [isDotProd, opDotProd] = iBuildDotProdOperatorIfNeeded(A, transA, B, transB);
    if isDotProd
        % Compute replicated sum reduction to check the output class
        localC = spmdReduce(@(a,b) sum([a b]), opDotProd(localA, localB));
        % Set the local part empty, with correct class, on all the workers
        % but 1
        if spmdIndex ~= 1
            localC = zeros(1, 0, 'like', localC);
        end
        codistrC = codistributor1d(2, [1, zeros(1, spmdSize - 1)], [1, 1]);
        C = codistributed.pDoBuildFromLocalPart(localC, codistrC);
        return
    end
    
    % General mtimes implementation
    [localC, codistrC] = codistrA.hMtimesImpl(codistrA, localA, transA, ...
        codistrB, localB, transB);
else % One array is codistributed; the other is replicated
    if ~iIsValidCodistributor(A, areBothMatricesDense)
        A = iConvertTo1d(A);
    end
    if ~iIsValidCodistributor(B, areBothMatricesDense)
        B = iConvertTo1d(B);
    end
    
    [localA, codistrA]  = iExtractCodistrAndLP(A);
    [localB, codistrB]  = iExtractCodistrAndLP(B);
    if ~isempty(codistrA)
        codistr = codistrA;
    else
        codistr = codistrB;
    end
    
    [localC, codistrC] = codistr.hMtimesReplicatedImpl(codistrA, localA, transA, ...
        codistrB, localB, transB);
end
C = codistributed.pDoBuildFromLocalPart(localC, codistrC);
end % End of mtimes

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [A, B] = iValidateInputArgs(A, B, transA, transB)

% Use the in-memory code to throw some specific errors for unsupported
% types (e.g. datetime, duration*duration, int*int, etc.).
% NB: no need to "throwAsCaller" as this whole function is already in a
% try-catch.
Aproto = distributedutil.Allocator.extractTemplate(A);
Bproto = distributedutil.Allocator.extractTemplate(B);
[~] = mtimes(Aproto,Bproto);

% Logicals are converted to doubles.
if islogical(A)
    A = double(A);
end
if islogical(B)
    B = double(B);
end
% Numeric inputs must be floating-point
if (isnumeric(A) && ~isfloat(A)) ...
        || (isnumeric(B) && ~isfloat(B))
    error(message('MATLAB:mtimes:integerNotSupported'));
end

if ~ismatrix(A) || ~ismatrix(B)
    error(message('MATLAB:mtimes:inputsMustBe2D'));
end

% Determine inner dimensions
if transA
    innerDimensionA = size(A,1);
else
    innerDimensionA = size(A,2);
end
if transB
    innerDimensionB = size(B,2);
else
    innerDimensionB = size(B,1);
end
% Check that inner dimension fits
if innerDimensionA ~= innerDimensionB
    error(message('MATLAB:innerdim'));
end
end

function [A, B] = iUnifyCodistributors(A, B, areBothMatricesDense)
codistrA = getCodistributor(A);
codistrB = getCodistributor(B);
if numel(A) > numel(B)
    if codistrA.hMtimesCheck(areBothMatricesDense)
        codistrB = codistrA.hGetNewForSize( size(B) );
        B = redistribute(B, codistrB);
    else
        if codistrB.hMtimesCheck(areBothMatricesDense)
            codistrA = codistrB.hGetNewForSize( size(A) );
            A = redistribute(A, codistrA);
        else
            A = iConvertTo1d(A);
            B = iConvertTo1d(B);
        end
    end
else
    if codistrB.hMtimesCheck(areBothMatricesDense)
        codistrA = codistrB.hGetNewForSize( size(A) );
        A = redistribute(A, codistrA);
    else
        if codistrA.hMtimesCheck(areBothMatricesDense)
            codistrB = codistrA.hGetNewForSize( size(B) );
            B = redistribute(B, codistrB);
        else
            A = iConvertTo1d(A);
            B = iConvertTo1d(B);
        end
    end
end
end

function A = iConvertTo1d(A)
% To be called if hMtimesCheck fails for the distributions of an input
[~, dim] = max(size(A));
A = redistribute(A, codistributor1d(dim));
end

function [LP, codistr] = iExtractCodistrAndLP(D)
if isa(D, 'codistributed')
    % Get codistributor and local part from codistributed array.
    LP = getLocalPart(D);
    codistr = getCodistributor(D);
else
    % Let empty codistributor indicate that local part represents a
    % replicated array.
    LP = D;
    codistr = [];
end
end

function tf = iIsValidCodistributor(X, areBothMatricesDense)
if isa(X, 'codistributed')
    tf = hMtimesReplicatedCheck( getCodistributor(X), areBothMatricesDense );
else
    % replicated array has a "valid" codistributor
    tf = true;
end
end

function A = iMaybeTranspose(A, transA)
% Perform the appropriate transposes (if any)
if transA == distributedutil.TransposeFlags.TRANSPOSE
    A = transpose(A);
elseif transA == distributedutil.TransposeFlags.CONJUGATE_TRANSPOSE
    A = ctranspose(A);
end
end

function [isDotProd, opDotProd] = iBuildDotProdOperatorIfNeeded(A, transA, B, transB)
import distributedutil.TransposeFlags
% Check for the special case of scalar output and return a flag and
% the correspnding local operator.

opDotProd = @(u,v) NaN;
isDotProd = false;
codistrA = getCodistributor(A);
codistrB = getCodistributor(B);
% At this point, A and B are of same codistributor class. We don't use any
% shortcut if they're 2dbc codistributed because it is not possible to
% represent a scalar as 2dbc codistributed.
if isa(codistrA, 'codistributor1d')
    % If the partitions and the local sizes are exactly the same we do not need to redistribute
    if isequal(codistrA.Partition, codistrB.Partition) && ...
            isequal(size(getLocalPart(A)), size(getLocalPart(B)))
        % Now we can check for the different cases returning a scalar
        if (transA == TransposeFlags.TRANSPOSE && iscolumn(A) && ...
                transB == TransposeFlags.NO_TRANSPOSE && iscolumn(B))
            isDotProd = true;
            opDotProd = @(u,v) u.'*v;
        elseif (transA == TransposeFlags.CONJUGATE_TRANSPOSE && iscolumn(A) && ...
                transB == TransposeFlags.NO_TRANSPOSE && iscolumn(B))
            isDotProd = true;
            opDotProd = @(u,v) u'*v;
        elseif (transA == TransposeFlags.NO_TRANSPOSE && isrow(A) && ...
                transB == TransposeFlags.TRANSPOSE && isrow(B))
            isDotProd = true;
            opDotProd = @(u,v) u*v.';
        elseif (transA == TransposeFlags.NO_TRANSPOSE && isrow(A) && ...
                transB == TransposeFlags.CONJUGATE_TRANSPOSE && isrow(B))
            isDotProd = true;
            opDotProd = @(u,v) u*v';
        end
    end
end
end

function X = iEnsure1DDist(X)
% Helper to ensure that an input is using a 1D distribution scheme if
% distributed. Inputs that are replicated or already 1D distributed are
% left alone. Any that are distributed but not using codistributor1d are
% redistributed.
if iscodistributed(X) && ~isa(getCodistributor(X), "codistributor1d")
    X = iConvertTo1d(X);
end
end