function Z = pagemtimes(varargin)
%PAGEMTIMES   Page-wise matrix multiply for gpuArray.
%   Z = PAGEMTIMES(X,Y)
%   Z = PAGEMTIMES(X, transpX, Y, transpY)
%   
%   Example:
%   Using PAGEMTIMES to multiply a two batches gpuArray data:
%    
%       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 = gpuArray.rand(M, K, P1,  1, P3);
%       dB = gpuArray.rand(K, N,  1, P2, P3, P4);
%       dC = pagemtimes(dA, dB);
%       size(dC)                                % M x N x P1 x P2 x P3 x P4
%   
%   See also GPUARRAY, PAGEFUN.


%   Copyright 2020 The MathWorks, Inc.

narginchk(2, 4);

if (nargin == 3)
    error(message('MATLAB:pagemtimes:threeInputNotSupported'));
end

if (nargin == 2)
    % Ensure that only the data inputs (args 1,2) are on the GPU
    varargin = prepareSolverInputs(varargin, [1 1]);

    X = varargin{1};
    transpX = 'N';
    Y = varargin{2};
    transpY = 'N';
else
    % Ensure that only the data inputs (args 1,3) are on the GPU
    varargin = prepareSolverInputs(varargin, [1 0 1 0]);
    
    X = varargin{1};
    transpX = iGetAndCheckTransposeOption(varargin{2});
    Y = varargin{3};
    transpY = iGetAndCheckTransposeOption(varargin{4});
end

% X and Y must both be double or single, and cannot be sparse.
iCheckArrayType(X);
iCheckArrayType(Y);

% Number of dimensions of the output Z
ndZ = max(ndims(X), ndims(Y));

szX = size(X,1:ndZ);
szY = size(Y,1:ndZ);

% Error if input sizes do not match
iGetAndCheckSizes(szX, transpX, szY, transpY);

% Main function.
Z = iPageMatrixMultiply(X, transpX, Y, transpY);
end

function Z = iPageMatrixMultiply(X, transpX, Y,transpY)

X = iApplyTranspose(X, transpX);
Y = iApplyTranspose(Y, transpY);

Z = pagefun(@mtimes, X, Y);
end

function X = iApplyTranspose(X, transpX)
if (transpX == 'C')
    X = pagefun(@ctranspose, X);
elseif (transpX == 'T')
    X = pagefun(@transpose, X);
end
end

function iCheckArrayType(X)
% Check that input is a dense single or double array.
if ~isfloat(X)
    error(message('MATLAB:pagemtimes:invalidArray'));
end

if issparse(X)
    error(message('MATLAB:pagemtimes:sparseNotSupported'));
end
end

function iGetAndCheckSizes(szX, transpX, szY, transpY)
% Check that the sizes of X and Y are valid, and compute szZ based on
% the sizes of X and Y.

% If X or Y has scalar pages, the operation in each page is scalar
% expansion, not matrix multiplication. Cases where both is possible
% (e.g. 4-by-1 * 1-by1) are treated as scalar expansion.
scalarExpansion = (szX(1)==1 && szX(2)==1) || (szY(1)==1 && szY(2)==1);

% Size of first two dimensions after potential transposition
[~, colsX] = iRowsCols(szX, transpX);
[rowsY, ~] = iRowsCols(szY, transpY);

% Check if sizes are valid for matrix multiplication
if (~scalarExpansion && colsX ~= rowsY)
    error(message('MATLAB:pagemtimes:mismatchPageDimensions'));
end

% Error if there is a mismatch in batch dimensions.
if any( (szX(3:end) ~= 1) & (szY(3:end) ~= 1) & (szX(3:end) ~= szY(3:end)) )
    error(message('MATLAB:pagemtimes:mismatchImplicitExpansion'));
end
end

function [rows, cols] = iRowsCols(size, transp)
if transp == 'N'
    rows = size(1);
    cols = size(2);
else
    rows = size(2);
    cols = size(1);
end
end

% Check that option is valid, and return a char 'N', 'T' or 'C'.
function transp = iGetAndCheckTransposeOption(option)

if isempty(option) || ( ~(isstring(option) && isscalar(option)) && ~(ischar(option) && isrow(option)))
    error(message('MATLAB:string:MustBeStringScalarOrCharacterVector'))
end

transp = 'N';
if strncmpi('transpose', option, max(1,length(option)))
    transp = 'T';
elseif strncmpi('ctranspose', option, max(1,length(option)))
    transp = 'C';
elseif ~strncmpi('none', option, max(1,length(option)))
    error(message('MATLAB:pagemtimes:invalidOption'))
end
end
