function C = tensorprod(A, B, varargin)
%TENSORPROD Tensor products between two codistributed tensors.
%   C = TENSORPROD(A, B)
%   C = TENSORPROD(A, B, "ALL")
%   C = TENSORPROD(A, B, DIM)
%   C = TENSORPROD(A, B, DIMA, DIMB)
%   C = TENSORPROD(..., NUMDIMENSIONSA = NDIMSA)
%   
%   See also TENSORPROD, CODISTRIBUTED.
%   


%   Copyright 2022 The MathWorks, Inc.

% codistributed/TENSORPROD can take between 2 and 6 inputs, and can return strictly 1 output.
narginchk(2, 6)
nargoutchk(0, 1) % 0 outputs is 'ans'

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

A = isCheckIsDistributedInHighDimension(A);
B = isCheckIsDistributedInHighDimension(B);

if ~isfloat(A) || ~isfloat(B)
    error(message("MATLAB:tensorprod:InvalidDataType"));
end

if issparse(A) || issparse(B)
    error(message("MATLAB:tensorprod:InvalidInputSparse"));
end

% Gather the trailing arguments.
argList = distributedutil.CodistParser.gatherElements(varargin(1:end));
% If we are left with nothing distributed use the in-memory function
if ~iscodistributed(A) && ~iscodistributed(B)
    C = tensorprod(A, B, varargin{:});
    return;
end

% Parse the trailing input arguments.
[allFlag, nDimsFlag, dimFlag, dimAdimBFlag] = iParseFlags(argList);

% Check the N-V pair arguments.
ndimsA = ndims(A);
if ~isempty(nDimsFlag)
    ndimsA = iSetAndCheckNumDimsA(A, argList);
end

% Check the dimension arguments.
if ~isempty(dimFlag) || ~isempty(dimAdimBFlag)
    [dimA, dimB] = iSetAndCheckReductionDims(A, B, argList, dimFlag, dimAdimBFlag);
    % Check that all the elements of dimA are <= ndimsA
    if ~isempty(nDimsFlag)
        for ii = 1:numel(dimA)
            if dimA(ii) > ndimsA
                error(message("MATLAB:tensorprod:DimAGreaterThanNdimsA"))
            end
        end
    end
end

% Call the version of the codistributed function for the flags passed in.
if isempty(allFlag) && isempty(dimFlag) && isempty(dimAdimBFlag)
    if iscodistributed(A) && ~iscodistributed(B) % A codistributed and B replicated
        C = iTensorprodNoDimACodistributedBReplicated(A, B, ndimsA);
    else
        C = iTensorprodNoDim(A, B, ndimsA);
    end
elseif ~isempty(allFlag)
    C = iTensorprodAll(A, B);
elseif ~isempty(dimFlag) || ~isempty(dimAdimBFlag)
    C = iTensorprodDimADimB(A, B, ndimsA, dimA, dimB);
end
end

%%%%%%%%%%%
function A = isCheckIsDistributedInHighDimension(A)
% Remove nd-arrays on workers if distribution is in the third dimension or
% higher.
if iscodistributed(A)
    codistA = getCodistributor(A);
    if isa(getCodistributor(A), 'codistributor1d') && codistA.Dimension > ndims(A)
        A = gather(A);
    end
end
end

%%%%%%%%
function C = iTensorprodNoDim(A, B, ndimsA)
% Implements tensorprod(A, B): compute the outer product.
colA = iUnfold(A);
colB = iUnfold(B);
matAB = colA*colB.';
C = reshape(matAB, [size(A), ones(1, ndimsA-numel(size(A))), size(B)]);
end

%%%%%%%%
function C = iTensorprodNoDimACodistributedBReplicated(A, B, ndimsA)
% Implements tensorprod(A, B): compute the outer product when A is codistributed and B is in-memory.
assert(~iscodistributed(B), "B must be in-memory to use this function.")
if isa(getCodistributor(A), 'codistributor2dbc')
    A = redistribute(A, codistributor1d());
end
codistA = getCodistributor(A);
codistDim = codistA.Dimension;
if codistDim == ndimsA
    % Account for singleton dimensions in the last dimension after
    % distributing.
    % If a worker is empty you get a '0' instead of a '1' in the place of
    % the singleton dimension.
    localC = tensorprod(getLocalPart(A), B, NumDimensionsA = codistDim);
else
    localC = tensorprod(getLocalPart(A), B);
end
codistC = codistributor1d(codistDim, codistA.Partition, [size(A) size(B)]);
C = codistributed.build(localC, codistC);
% Allow for singleton dimensions
szC = [size(A), ones(1, ndimsA-numel(size(A))) size(B)];
if ~isempty(szC)
    if isscalar(szC)
        szC = [szC 1];
    end
    C = reshape(C, szC);
end
end

%%%%%%%%
function C = iTensorprodAll(A, B)
% Implements tensorprod(A, B, "all"): compute the inner product.
% Note that the error checking for A and B being of equal size comes for free with the matrix multiplication.
colA = iUnfold(A);
colB = iUnfold(B);
C = colA.'*colB;
end

%%%%%%%%
function C = iTensorprodDimADimB(A, B, ndimsA, dimA, dimB)
% Implements tensorprod(A, B, dimA, dimB): compute the tensor product.
% Account for trailing singleton dimensions.
maxDim = max([ndimsA, ndims(B), dimA, dimB]);
szA = size(A, 1:maxDim);
szB = size(B, 1:maxDim);
% Create logicals for the indices of the dimensions to be
% contracted.
contractedDimsA = false(1,maxDim);
contractedDimsB = false(1,maxDim);
contractedDimsA(dimA) = true;
contractedDimsB(dimB) = true;
% Permute A and B by a vector containing the dimensions not being contracted followed by the dimensions
% %that are being contracted for A (and vice versa for B).
permuteA = permute(A, [find(~contractedDimsA) dimA]);
permuteB = permute(B, [dimB find(~contractedDimsB)]);
% Reshape the permuted tensors into matrices.
% The reshape size is so that we multiply the (un)contracted dimensions
% of A by the (un)contracted dimensions of B.
matA = reshape(permuteA, prod(szA(~contractedDimsA)), prod(szA(contractedDimsA)));
matB = reshape(permuteB, prod(szB(contractedDimsB)), prod(szB(~contractedDimsB)));
% Calculate the size of the final tensor which is the size of
% the original tensors minus the contracted dimensions, using the reshaped
% matrix dimensions
szC = [szA(~contractedDimsA) szB(~contractedDimsB)];
% Calculate the tensor product.
C = matA*matB;
if ~isempty(szC)
    C = reshape(C, szC);
end
end

%%%%%%%%
function [allFlag, nDimsFlag, dimFlag, dimAdimBFlag] = iParseFlags(args)
% Parse the flags, and error accordingly if something is wrong.

% Set empty flags.
allFlag = [];
nDimsFlag = [];
dimFlag = [];
dimAdimBFlag = [];

% Early exit if no flag passed.
if isempty(args)
    return
end

% Check for "all" and "NumDims" flags.
for ii = 1:numel(args)
    if matlab.internal.datatypes.isScalarText(args{ii})
        if startsWith("all", args{ii}, "IgnoreCase", true)
            if ~isempty(allFlag)
                error(message("MATLAB:tensorprod:KeyWithoutValue"));
            end
            allFlag = true;
        elseif startsWith("NumDimensionsA", args{ii}, "IgnoreCase", true)
            nDimsFlag = true;
            % Make sure the arguments come in a pair.
            if mod(numel(args(ii:end)), 2) ~= 0
                error(message("MATLAB:tensorprod:KeyWithoutValue"));
            end
        else
            error(message("MATLAB:tensorprod:ParseFlags"))
        end
    end
end

% Count the number of numeric elements to determine whether the argument is 'dim' or
% 'dimA, dimB' syntax.
numElems = args(:, cellfun(@isnumeric, args(1, :)));
countNumElems = numel(numElems);
numExistingFlags = ~isempty(nDimsFlag);
if isnumeric(args{1})  % Allow for repeated NV pairs.
    switch countNumElems
        case numExistingFlags
            % No additional flags
        case numExistingFlags+1
            dimFlag = true;
        case numExistingFlags+2
            dimAdimBFlag = true;
    end
end
end

%%%%%%%%
function [dimA, dimB] = iSetAndCheckReductionDims(A, B, args, dimFlag, dimAdimBFlag)
% Set and check the 'dimA', 'dimB', 'dim' inputs
if ~isempty(dimFlag)
    dimA = args{1};
    dimB = args{1};
elseif ~isempty(dimAdimBFlag)
    dimA = args{1};
    dimB = args{2};
end

% If dimA or dimB are empty, then they must be 0x0.
if isempty(dimA) || isempty(dimB)
    if ~(ismatrix(dimA) || size(dimA, 1) == 0 && size(dimA, 2) == 0)
        error(message("MATLAB:tensorprod:InvalidDIM"));
    elseif ~(ismatrix(dimB) || size(dimB, 1) == 0 && size(dimB, 2) == 0)
        error(message("MATLAB:tensorprod:InvalidDIM"));
    else
        % Return early for 0x0 empty dim
        return;
    end
end

% Make sure that dimA and dimB are numeric, real-valued 2D vectors
if any(~isnumeric(dimA)) || ~isvector(dimA) || any(~isreal(dimA))
    error(message("MATLAB:tensorprod:InvalidDIM"))
end
if any(~isnumeric(dimB)) || ~isvector(dimB) || any(~isreal(dimB))
    error(message("MATLAB:tensorprod:InvalidDIM"))
end

% Make sure that the elements of dimA and dimB are finite integers >=1
if any(~isfinite(dimA)) || any(dimA < 1) || any(round(dimA) ~= dimA)
    error(message("MATLAB:tensorprod:InvalidDIM"))
end
if any(~isfinite(dimB)) || any(dimB < 1) || any(round(dimB) ~= dimB)
    error(message("MATLAB:tensorprod:InvalidDIM"))
end

% Make sure that the elements are unique
if numel(unique(dimA)) ~= numel(dimA) || numel(unique(dimB)) ~= numel(dimB)
    error(message("MATLAB:tensorprod:InvalidDIM"))
end

% Make sure that the sizes of the reduction dimensions match.
if length(dimA)~=length(dimB)
    error(message("MATLAB:tensorprod:DimLengthsMismatch"))
end

% Make sure that the sizes of the reduced tensors match.
if size(A, dimA) ~= size(B, dimB)
    error(message("MATLAB:tensorprod:InvalidDIMSize"))
end
end

%%%%%%%%
function ndimsA = iSetAndCheckNumDimsA(A, args)
ndimsA = args{end}; % ndimsA is always the last argument in varargin.

% Make sure that ndimsA numeric, real-valued scalar
if ~isnumeric(ndimsA) || ~isscalar(ndimsA) || ~isreal(ndimsA)
    error(message("MATLAB:tensorprod:InvalidNumDimsA"))
end

% Check ndimsA is a real, finite,integer >=1
if  ~isfinite(ndimsA) || ndimsA < 1 || round(ndimsA) ~= ndimsA
    error(message("MATLAB:tensorprod:InvalidNumDimsA"))
end

% Check that ndimsA >=ndims(A) except when A is a column or row vector, when ndims(A)=2
% but ndims is allowed to be 1
if ~isvector(A) && ndimsA < ndims(A)
    error(message("MATLAB:tensorprod:InvalidNumDimsA"))
end
end

%%%%%%%%
function colT = iUnfold(T)
% Return the unfolded version of a tensor T depending on whether or not it
% is codistributed.
if ~iscodistributed(T)
    colT = T(:);
else
    colT = colonize(T);
end
end
