function K = kron(A, B)
%KRON Kronecker tensor product.
%   K = KRON(A, B)
%   
%   See also KRON.
%   


%   Copyright 2022 The MathWorks, Inc.

% Check the arguments and throw an error if something is wrong.
arguments
    A {mustBeMatrix(A)}
    B {mustBeMatrix(B), integerTypesMustBeConsistent(A, B)}
end

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

A = iReplicateIfDistributedInHighDim(A);
B = iReplicateIfDistributedInHighDim(B);

% Kron(A,B) for different input types.
if issparse(A) || issparse(B)
    if iscodistributed(A) && ~iscodistributed(B)
        K = iKronACodistributedBReplicated(A, B);
    elseif ~iscodistributed(A) && iscodistributed(B)
        K = iKronAReplicatedBCodistributed(A, B);
    else % Both A and B are codistributed. Gather the matrix with the least number of elements.
        if numel(A) < numel(B)
            % Gather A to replicate to all workers.
            A = gather(A);
            K = iKronAReplicatedBCodistributed(A, B);
        else
            % Gather B to replicate to all workers.
            B = gather(B);
            K = iKronACodistributedBReplicated(A, B);
        end
    end
else % A and B dense.
    if iscodistributed(A) && ~iscodistributed(B)
        K = iKronACodistributedBReplicated(A, B);
    elseif ~iscodistributed(A) && iscodistributed(B)
        K = iKronAReplicatedBCodistributed(A, B);
    else % Both codistributed.
        K = iKronDenseAB(A, B);
    end
end

end

%%%%%%%%%%%%

function mustBeMatrix(A)
if ~ismatrix(A)
    error(message('MATLAB:kron:TwoDInput'));
end
end

function integerTypesMustBeConsistent(A, B)
if isinteger(A) && isinteger(B)
    % Both integer.
    if ~isequal(underlyingType(A), underlyingType(B))
        error(message("MATLAB:mixedClasses"));
    end
    return;
end
%
if isinteger(A) && ~isscalar(B)
    % One an integer, then the other must be a scalar double
    error(message("MATLAB:mixedClasses"));
elseif isinteger(B) && ~isscalar(A)
    % One an integer, then the other must be a scalar double
    error(message("MATLAB:mixedClasses"));
end
end

%%%%%%%%%%%
function A = iReplicateIfDistributedInHighDim(A)
% copies the 2D array to all workers if distributed in a dimension higher
% than 2
if iscodistributed(A)
    codistA = getCodistributor(A);
    if isa(codistA, 'codistributor1d') && codistA.Dimension > 2
        A = gather(A);
    end
end
end

%%%%%%%%%%%
function K = iKronACodistributedBReplicated(A, B)
assert(~iscodistributed(B), "B should be in-memory to use this function.")
% If A is distributed via codistributor2dbc then redistribute.
if isa(getCodistributor(A), 'codistributor2dbc')
    A = redistribute(A, codistributor1d());
end
codistA = getCodistributor(A);
localK = kron(getLocalPart(A), B);
% Build the codistributed solution K from its local parts, with the
% partition multiplied by the size of B in the partition dimension.
codistK = codistributor1d(codistA.Dimension, codistA.Partition*size(B, codistA.Dimension), size(A).*size(B));
K = codistributed.build(localK, codistK);
end

%%%%%%%%%%%
function K = iKronAReplicatedBCodistributed(A, B)
assert(~iscodistributed(A), "A should be in-memory to use this function.")
% If B is distributed via codistributor2dbc then redistribute.
if isa(getCodistributor(B), 'codistributor2dbc')
    B = redistribute(B, codistributor1d());
end
codistB = getCodistributor(B);
localK = kron(A, getLocalPart(B));
dim = codistB.Dimension;
nA = size(A, dim);
codistK = codistributor1d(dim, codistB.Partition*nA, size(A).*size(B));
K = codistributed.build(localK, codistK);
% The matrix has the rows/columns in the wrong order.
% Build a permutation vector to reorder.
% First check for empty or scalar matrices since then we don't need a shuffle.
if numel(A) <= 1 || numel(B) <= 1
    return;
end
partB = codistB.Partition;
partB = partB(partB > 0); % Non-zero entries of the partition vector of B, e.g. [2 2 1] if B has 5 rows/columns and 3 workers.
partK = codistK.Partition; % nA*partB, e.g. [4 4 2] with nA = 2.
% Create a vector of indices, one for each row/column of K in the distribution dimension.
% E.g. [ 1 2 3 4 5 6 7 8 9 10].
v = 1:1:sum(partK);
% Split the indices up into blocks to rearrange based on partB.
% E.g. {[1 2]}, {[3 4]}, {[5 6]}, {[7 8]}, {[9]}, {[10]}
vi = mat2cell(v, 1, kron(partB, ones(1, nA)));
% Create the permutation vector to select the first blocks from each
% worker, then the second, etc.
% We need UniformOutput=False because the output of arrayfun is arrays rather
% than scalars.
% E.g. [1 3 5 2 4 6]
permuteVi = cell2mat(arrayfun(@(x) (1:nA:numel(partB)*nA)+x, 0:1:(nA-1), 'UniformOutput', false));
% Permute the blocks and turn back into a row vector.
% E.g. [1 2 5 6 9 3 4 7 8 10]
permVec = cell2mat(vi(permuteVi));
if dim == 1 % Distribution by rows.
    K = subsref(K, substruct('()', {permVec, ':'})); % Shuffle the rows.
else % Distributed by columns.
    K = subsref(K, substruct('()', {':', permVec})); % Shuffle the columns.
end
end

%%%%%%%
function K = iKronDenseAB(A, B)
% Create 4D arrays to separate the columns of A and B.
% For A, the columns are transposed and split in the fourth dimension.
pA = permute(A, [3, 1, 4, 2]);
% For B, the columns are split in the third dimension.
pB = permute(B, [1, 3, 2, 4]);

% Multiply the permuted arrays element-wise to obtain all possible
% products between the elements of A and B.
% M(:, :, n, m) = B(:, n)*A(:, m)' for m = numCols(A),
% n = numCols(B).
M = pA.*pB;
K = reshape(M, size(A).*size(B));
end
