function Z = pagemtimes(varargin)
%PAGEMTIMES   Page-wise matrix multiply for codistributed array.
%   Z = PAGEMTIMES(X,Y)
%   Z = PAGEMTIMES(X, transpX, Y, transpY)
%   
%   Example:
%   Using PAGEMTIMES to multiply a two batches codistributed array data:
%   spmd
%       M = 3;                                  % output number of rows
%       K = 6;                                  % matrix multiply inner dimension
%       N = 2;                                  % output number of columns
%       P1 = 10;                                % size of first page dimension
%       P2 = 17;                                % size of second page dimension
%       P3 = 4;                                 % size of third page dimension
%       P4 = 12;                                % size of fourth page dimension
%       dA = codistributed.rand(M, K, P1,  1, P3);
%       dB = codistributed.rand(K, N,  1, P2, P3, P4);
%       dC = pagemtimes(dA, dB);
%       size(dC)                                % M x N x P1 x P2 x P3 x P4
%   end
%   
%   See also CODISTRIBUTED, PAGEFUN.


%   Copyright 2020-2022 The MathWorks, Inc.

narginchk(2, 4);

[inputMatrices, transA, transB] = iParseArgs(varargin{:});


% Calculate the number of pages in dimensions 3 and up, referred as 
% trailing output dimensions. Will error if the corresponding page 
% dimensions of the inputs are not consistent.
trailingOutDims = parallel.internal.array.pagefunTrailingDimensions(inputMatrices, ...
    "MATLAB:pagemtimes:mismatchImplicitExpansion");

% Get the codistributor of the output.
outCodistr = pagefunGetOutputsCodistr(inputMatrices, trailingOutDims);
% Redistribute the intputs if needed.
inputMatrices = pagefunRedistributeInputsIfNeeded(inputMatrices, outCodistr);

% Extract the local parts.
localParts = pagefunGetLocalParts(inputMatrices);
forceSymmetry = iOutputPagesMustBeSymmetric(inputMatrices{1}, transA, inputMatrices{2}, transB);
% Call in-memory pagemtimes.
try
    lpZ = pagemtimes(localParts{1}, transA, localParts{2-forceSymmetry}, transB);
    err = [];
catch err
end
% Check for errors
err = spmdCat(err);
if ~isempty(err)
    throw(err(1));
end

% Build the output.
Z = pagefunBuildOutputs({lpZ}, outCodistr, trailingOutDims);
end

function [inputMatrices, transA, transB] = iParseArgs(varargin)
% Parse and check the input arguments
if nargin == 3
    throwAsCaller(iMException('MATLAB:pagemtimes:threeInputNotSupported'));
end
if any(cellfun(@issparse, varargin))
    throwAsCaller(iMException('MATLAB:pagemtimes:sparseNotSupported'));
end
if nargin == 2
    inputMatrices = varargin(:);
    transA = 'n';
    transB = 'n';
else % nargin == 4
    inputMatrices = varargin([1 3]);
    % Transpose flag text will be check by in-memory function on all workers.
    % Here just make sure they are strings.
    transA = varargin{2};
    transB = varargin{4};
    if ~matlab.internal.datatypes.isScalarText(transA) ...
            || ~matlab.internal.datatypes.isScalarText(transB)
        throwAsCaller(iMException("MATLAB:pagemtimes:invalidOption"));
    end
end
% Both matrices must be floating point
if ~isfloat(inputMatrices{1}) || ~isfloat(inputMatrices{2})
    throwAsCaller(iMException("MATLAB:pagemtimes:invalidArray"));
end
end

function tf = iOutputPagesMustBeSymmetric(A, transA, B, transB)
% Check for special cases that must produce symmetric output. In
% particular: A'*A and A*A' for real A, or A.'*A and A*A.' for complex A.

pageSzsA = size(A, 1, 2);
pageSzsB = size(B, 1, 2);
isVectorA = any(pageSzsA == 1);
isVectorB = any(pageSzsB == 1);

if ~isequal(transA, transB) && ...                     % different transpose flags
        ~isempty(A) && ~isempty(B) && ...              % no empty input
        isequal(pageSzsA, pageSzsB) && ...             % same page size
        ~isVectorA && ~isVectorB && ...                % no vector page
        isequal(underlyingType(A), underlyingType(B))  % same numeric types
    % Finally, result must be symmetric if both inputs are equal, this is
    % expensive but unavoidable at this point.
    tf = isequaln(A, B);
else
    tf = false;
end
end

function me = iMException(id)
% Create an MException from the specified error ID
me = MException(message(id));
end