function [res1, res2] = spdiags(varargin)
%SPDIAGS Sparse matrix formed from diagonals.
%   
%   [B,d] = SPDIAGS(A)
%   B = SPDIAGS(A,d)
%   A = SPDIAGS(B,d,A)
%   A = SPDIAGS(B,d,m,n)
%   
%   
%   Example:
%   spmd
%       n = 1e3;
%       b = codistributed.ones(n,1);
%       A = SPDIAGS([b -2*b b], -1:1, 1e3, 1e3);
%   end
%   
%   See also SPDIAGS, CODISTRIBUTED


%   Copyright 2016-2023 The MathWorks, Inc.

% Check number of in-/outputs
narginchk(1, 4);
nargoutchk(1, 2);

% Call spdiags on the host if none of the data inputs are distributed
if nargin<=2 && ~iscodistributed(varargin{1})
    [res1, res2] = spdiags(varargin{1}, gather(varargin{2}));
    return
elseif nargin==3 && ~(iscodistributed(varargin{1}) || ...
        iscodistributed(varargin{3}))
    res1 = spdiags(varargin{1}, gather(varargin{2}), varargin{3});
    return
elseif nargin==4 && ~iscodistributed(varargin{1})
    res1 = spdiags(varargin{1}, gather(varargin{2}), gather(varargin{3}), ...
        gather(varargin{4}));
    return
end

% Check distribution of input A and set up variables
checkedVariables = iCheckInputsAndSetUpVariables(varargin{:});

% For case 2 and D empty, create output of respective dimensions and
% return early
if nargin==2 && isempty(checkedVariables.D)
    res1 = codistributed(zeros(min(checkedVariables.distA.Cached.GlobalSize),0));
    res2 = zeros(0,1);
    return
end

% Collect local data
localA = getLocalPart(checkedVariables.A);
globalIndicesOnWorker = globalIndices(checkedVariables.distA, ...
    checkedVariables.distA.Dimension);
localOffset = [];
if ~isempty(globalIndicesOnWorker)
    localOffset = globalIndicesOnWorker(1)-1;
end

% Enter code path depending on number of inputs
if nargin<=2 % Extract diagonal(s) from given matrix
    [res1, res2] = iExtractDiagonals(localA, checkedVariables.distA, ...
        checkedVariables.D, localOffset, nargin==1);
else % Create/modify matrix with given diagonal(s)
    res1 = iCreateOrModifyMatrix(localA, checkedVariables.distA, ...
        checkedVariables.B, checkedVariables.D, localOffset, nargin==3);
    % res2 is not set to be conform with core math implementation
end
end

function returnVariables = iCheckInputsAndSetUpVariables(varargin)
% Helper function to do various checks with inputs and to create the main
% variables that are used later on

% Check that all inputs are numeric or char and throw otherwise.
validInputs = true;
for i=1:nargin
    validInputs = validInputs && ( ...
        isnumeric(varargin{i}) || ischar(varargin{i}) || islogical(varargin{i}));
end
if ~validInputs
    ME = MException(message('parallel:distributed:SpdiagsInvalidTypesOfInputs'));
    throwAsCaller(ME);
end

if nargin<=3
    % Check matrix A
    if nargin<=2
        A = varargin{1};
    else
        A = varargin{3};
    end
    if ~iscodistributed(A)
        sizeA = size(A);
        [~, dominantDimension] = max(sizeA);
        % Use dominantDimension as distribution dimension such that
        % pre-/appending of zeros needs to be done on only one worker.
        distA = codistributor1d(dominantDimension, ...
            codistributor1d.defaultPartition(sizeA(dominantDimension)), sizeA);
        A = codistributed(A, distA);
    else
        distA = getCodistributor(A);
        [~, dominantDimension] = max(distA.Cached.GlobalSize);
    end
    % If A is 2dbc distributed (case 1/2) or the 1d distribution
    % dimension is not correct (case 1), redistribute A.
    if ~isa(distA, 'codistributor1d') || ...
            ( nargin==1 && isa(distA, 'codistributor1d') && ...
            distA.Dimension~=(3-dominantDimension) )
        if nargin==1
            distDimensionA = 3-dominantDimension;
        else
            distDimensionA = dominantDimension;
        end
        A = redistribute(A, codistributor1d(distDimensionA));
        distA = getCodistributor(A);
    end
else
    % Create empty A of correct dimension from gathered inputs three and
    % four and convert to double to treat chars correctly
    try
        sizeA = parallel.internal.array.checkSparseDimensions(varargin{3}, varargin{4});
    catch ME
        throwAsCaller(ME);
    end

    [~, dominantDimension] = max(sizeA);
    distA = codistributor1d(dominantDimension, ...
        codistributor1d.defaultPartition(sizeA(dominantDimension)), sizeA);
    A = sparse(sizeA(1), sizeA(2), distA, 'noCommunication');
end

if nargin>=3
    % Ensure B is not 2dbc distributed. If a row-vector then gather to
    % avoid problems when applying dimension expansion.
    B = varargin{1};
    if iscodistributed(B)
        if isrow(B)
            B = gather(B);
        else
            distB = getCodistributor(B);
            if isa(distB, 'codistributor2dbc')
                B = redistribute(B, codistributor1d(1));
            end
        end
    end
    returnVariables.B = B;
end

returnVariables.A = A;
returnVariables.distA = distA;
returnVariables.D = [];

if nargin~=1
    if iscodistributed(varargin{2})
        D = gather(varargin{2});
    else
        D = varargin{2};
    end
    % Colonize D
    D = D(:);
    % Check that D is an integer
    if any(ceil(D)~=D) || any(~isreal(D))
        ME = MException(message('parallel:array:SpdiagsInvalidTypeFor2ndInput'));
        throwAsCaller(ME);
    end
    returnVariables.D = D;
end
end

function [res1, res2] = iExtractDiagonals(localA, distA, D, ...
    localOffset, extractAll)
% Helper for one and two inputs that extract all or specific diagonals
if extractAll
    % Determine which local diagonals are not empty
    if ~isempty(localA)
        [i,j] = find(localA);
        if isrow(i)
            i = i';
            j = j';
        end
        dIndex = sort(j-i);
        localDiagonalIndex = dIndex(diff([-inf; dIndex])~=0);
        % Determine global diagonal index
        if distA.Dimension==1
            globalMap = localDiagonalIndex-localOffset;
        else
            globalMap = localDiagonalIndex+localOffset;
        end
        % Map between local and global diagonal index
        localToGlobalMap = [localDiagonalIndex, globalMap];
    else
        globalMap = [];
    end
    globalDiagonalsIndex = unique(spmdReduce(@vertcat, ...
        globalMap), 'sorted');
    numberDiagGlobal = length(globalDiagonalsIndex);
else
    numberDiagGlobal = length(D);
    localDiagData = cell(numberDiagGlobal,1);
end
% Create local data for output on non-empty workers
if ~isempty(localA)
    if extractAll
        localDiagData = iExtractLocalDataOfAllNonEmptyDiagonals(localA, ...
            distA, localToGlobalMap, globalDiagonalsIndex);
    else
        for i=1:numberDiagGlobal
            localDiagData{i} = iExtracNonEmptyLocalDataOfDiagonal(D(i), ...
                localA, localOffset, distA);
        end
    end
else % Create empty data of correct dimension on empty workers
    if extractAll
        localDiagData = zeros(0, numberDiagGlobal, 'like', localA);
    else
        for i=1:numberDiagGlobal
            localDiagData{i} = zeros(0, 1, 'like', localA);
        end
    end
end
% Create outputs
if extractAll
    globalSizeB = [min(distA.Cached.GlobalSize), numberDiagGlobal];
    partitionB = distA.Partition;
    distB = codistributor1d(1, partitionB, globalSizeB);
    res1 = codistributed.pDoBuildFromLocalPart(localDiagData, distB);
    res2 = globalDiagonalsIndex;
else
    columnsOfB = cell(numberDiagGlobal,1);
    globalSizeColumnsOfB = [min(distA.Cached.GlobalSize), 1];
    globalSizeB = [globalSizeColumnsOfB(1), numberDiagGlobal];
    % For each requested diagonal, distribute data as necessary
    for i=1:numberDiagGlobal
        partitionColumnsOfB = spmdReduce(@horzcat, size(localDiagData{i},1));
        columnsOfB{i} = codistributed.pDoBuildFromLocalPart(localDiagData{i}, ...
            codistributor1d(1, partitionColumnsOfB, globalSizeColumnsOfB));
    end
    % Redistribute columns of B to single distribution scheme
    [localRedistB, distB] = codistributed.pRedistSameSizeToSingleDist(...
        columnsOfB);
    distB = codistributor1d(1, distB.Partition, globalSizeB);
    % Concatenate all diagonals into one block of size globalSizeB
    res1 = codistributed.pDoBuildFromLocalPart([localRedistB{:}], distB);
    res2 = D;
end
end

function res1 = iCreateOrModifyMatrix(localA, distA, B, D, localOffset, replaceDiags)
% Helper for three and four inputs that either replaces certain diagonals
% or creates a matrix with certain diagonals
numberDiagGlobal = length(D);
localLengthDiag = zeros(numberDiagGlobal, 1);
sizeA = distA.Cached.GlobalSize;
resultIsLogical = islogical(localA) || islogical(B);
% Check size of matrix B (should be min(sizeA)-by-numberDiagGlobal) or
% 1-by-numberDiagGlobal is using dimension expansion).
% For backwards compatibility, only error if the code would
% previously have errored out in the indexing expression.
if isrow(B)
    rowsBOK = true;
else
    maxIndexRows = max(max(1, 1-D), min(sizeA(1), sizeA(2)-D)) + ...
        (sizeA(1)>=sizeA(2))*D;
    maxIndexRows(max(1,1-D) > min(sizeA(1), sizeA(2)-D)) = 0;
    rowsBOK = all(maxIndexRows <= size(B, 1));
end

if ~rowsBOK || numberDiagGlobal > size(B, 2)
    if replaceDiags
        error(message('MATLAB:spdiags:InvalidSizeBThreeInput'));
    else
        error(message('MATLAB:spdiags:InvalidSizeBFourInput'));
    end
end

if ~isempty(localA)
    sizeLocalA = size(localA);
    localD = D+(3-2*distA.Dimension)*localOffset;
    % Determine workers that need to ap-/prepend zeros
    prependZ = zeros(numberDiagGlobal, 1);
    appendZ = zeros(numberDiagGlobal, 1);
    localInnerLength = zeros(numberDiagGlobal, 1);
    for i=1:numberDiagGlobal
        [prependZ(i), appendZ(i)] = iWorkersToAddZeros(distA, D(i));
        [localInnerLength(i,1), localLengthDiag(i,1)] = iLocalLengthOfDiagonal(distA, ...
            sizeLocalA, localD(i), prependZ(i), appendZ(i));
        % Truncate localLengthDiag in case of prependZ to replicate core
        % MATLAB behavior that silently accepts shorter inputs if elements
        % would not be indexed
        if appendZ(i)
            localLengthDiag(i,1) = localInnerLength(i,1);
        end
    end
    % Determine the necessary size of the input for the local call to SPDIAGS
    localInputLength = min(sizeLocalA);
end
newPartitionB = spmdReduce(@horzcat, localLengthDiag);
usedLengthDiag = sum(newPartitionB,2);
usedDiags = 0;
if iscodistributed(B)
    distB = getCodistributor(B);
    localB = getLocalPart(B);
end
% (Re)distribute each column of B such that data lies on the worker where
% it is needed
for i=1:numberDiagGlobal
    % Skip current loop pass if diagonal would be of size 0
    if usedLengthDiag(i)==0
        continue;
    end
    % Set up codistributor that reflects the correct data distribution
    distC = codistributor1d(1, newPartitionB(i,:), [usedLengthDiag(i), 1]);
    if iscodistributed(B)
        % Extract the correct column from B and crop to the desired size
        columnB = iExtractAndCropColumn(distB, localB, i, ...
            usedLengthDiag(i));
        C = redistribute(columnB, distC);
    else
        % Create the distributed column from the in-memory column
        if isrow(B)
            % Maybe expand to fill column
            C = B(i) .* codistributed.ones(usedLengthDiag(i), 1, distC);
        else
            C = codistributed(B(1:usedLengthDiag(i),i), distC);
        end
    end
    localC = getLocalPart(C);
    % Only work on workers that store data and where the diagonal will
    % actually modify data
    if ~isempty(localC) && ((localD(i)>=0 && localD(i)<sizeLocalA(2)) || ...
            (localD(i)<0 && -localD(i)<sizeLocalA(1)))
        if prependZ(i)
            localC = localC(end-localInnerLength(i)+1:end,1);
        end
        if appendZ(i)
            localC = localC(1:localInnerLength(i),1);
        end
        
        if localInnerLength(i)~=localInputLength
            % Add zeros to localC to create correct input for local
            % SPDIAGS call
            % Prepend zeros
            if sizeLocalA(1)>=sizeLocalA(2) && localD(i)>0 || ...
                    sizeLocalA(1)<sizeLocalA(2) && localD(i)<0
                localC = [zeros(localInputLength-localInnerLength(i), 1); localC]; %#ok<AGROW>
            else % Append zeros
                localC = [localC; zeros(localInputLength-localInnerLength(i), 1)]; %#ok<AGROW>
            end
        end
        % Save only parts that need to be applied and count the number
        % of effective diagonals that will be used
        usedDiags = usedDiags+1;
        localDiags{usedDiags} = localC; %#ok<AGROW>
        localUsedD{usedDiags} = localD(i);  %#ok<AGROW>
    end
end
% Only assign data that has been chosen to be effective
if usedDiags
    % Assign block of diagonals into local data
    localA = spdiags([localDiags{:}],[localUsedD{:}],localA);
end
% Ensure output is logical if A or B is logical
if resultIsLogical
    localA = logical(localA);
end
res1 = codistributed.pDoBuildFromLocalPart(localA, distA);
end

function localDiagData = iExtractLocalDataOfAllNonEmptyDiagonals(localA, ...
    distA, localToGlobalMap, globalDiagonalsIndex)
% Subroutine that extracts all non-empty diagonals on non-empty workers.
% Distribution dimension must not be dominant dimension of global matrix to
% create correct data.

% Create area for data of local diagonals of correct type
localDiagData = zeros(size(localA, distA.Dimension), ...
    length(globalDiagonalsIndex), 'like', full(localA([])));

if ~(isempty(globalDiagonalsIndex) || isempty(localToGlobalMap))
    % Get global position for local diags and apply spdiags locally
    [~,globalPositionOfLocal,~] = intersect(globalDiagonalsIndex, localToGlobalMap(:,2));
    localDiagData(:,globalPositionOfLocal) = spdiags(localA);
end
end

function localDiagData = iExtracNonEmptyLocalDataOfDiagonal(D, localA, localOffset, distA)
% Subroutine that extracts the local data of a specific diagonal D from all
% non-empty workers. Local data is pre-/appended with zeros if necessary.

% Local parameters
localD = D+(3-2*distA.Dimension)*localOffset;
sizeLocalA = size(localA);
% Determine workers that need to pre-/append zeros
[prependZeros, appendZeros] = iWorkersToAddZeros(distA, D);
% Determine necessary length of local diagonal parts to pre-/append zeros
[~, lenWithZeros] = iLocalLengthOfDiagonal(distA, sizeLocalA, ...
    localD, prependZeros, appendZeros);
% Extract localD-th diagonal from local data
if ~isvector(localA)
    localDiagData = full(diag(localA, localD));
else % DIAG cannot be applied to vectors
    % Index in row/column that represents the global diagonal entry
    % 1+localD (localA is a row)
    % 1-localD (localA is a column)
    [~, dominantLocalDimension] = max(sizeLocalA);
    localDiagIndex = 1+(2*dominantLocalDimension-3)*localD;
    if localDiagIndex<1 || localDiagIndex>length(localA)
        localDiagData = zeros(0,1, 'like', full(localA([])));
    else
        localDiagData = full(localA(localDiagIndex));
    end
end
% Pre-/append zeros to fill up localDiagData
numberOfZerosToAdd = lenWithZeros-size(localDiagData, 1);
if prependZeros
    localDiagData = [zeros(numberOfZerosToAdd, 1); localDiagData];
elseif appendZeros
    localDiagData = [localDiagData; zeros(numberOfZerosToAdd,1)];
end
end

function [prependZeros, appendZeros] = iWorkersToAddZeros(distA, D)
% Helper to determine which workers need to prepend or append zeros to the
% global super-/sub-diagonals.

globalIndicesOnWorker = globalIndices(distA, distA.Dimension);
if distA.Cached.GlobalSize(1)>=distA.Cached.GlobalSize(2)
    % Dominant matrix dimension==1 --> add zeros on the top or bottom
    prependZeros = D>0 ... % upper diagonals only
        ... %segment with first matrix row
        && ( (distA.Dimension==1 && any(globalIndicesOnWorker==1))...
        ...%diagonal precedes segment
        || (distA.Dimension==2 && any(globalIndicesOnWorker<=D)) );
    appendZeros = D<0 && ... % lower diagonals only
        ... % segment with last matrix row
        ( (distA.Dimension==1 && any(globalIndicesOnWorker==distA.Cached.GlobalSize(1)) ...
        && (distA.Cached.GlobalSize(2)-D)>distA.Cached.GlobalSize(1)) ...
        ... % diagonal exceeds segment
        || (distA.Dimension==2 && any(globalIndicesOnWorker>distA.Cached.GlobalSize(1)+D)) );
else % Dominant matrix dimension==2 --> add zeros on the left or right
    prependZeros = D<0 && ...% lower diagonals only
        ...%diagonal precedes segment
        ( (distA.Dimension==1 && any(globalIndicesOnWorker<=-D)) ...
        || (distA.Dimension==2 && any(globalIndicesOnWorker==1)) );
    appendZeros = D>0 && ... % upper diagonals only
        ...% segment with last column
        ( (distA.Dimension==2 && any(globalIndicesOnWorker==distA.Cached.GlobalSize(2)) ...
        && (distA.Cached.GlobalSize(1)+D)>distA.Cached.GlobalSize(2)) ...
        ...% diagonal exceeds segment
        || (distA.Dimension==1 && any(globalIndicesOnWorker>distA.Cached.GlobalSize(2)-D)) );
end
end

function [len, lenWithZeros] = iLocalLengthOfDiagonal(distA, sizeLocalA, ...
    localD, prependZ, appendZ)
% Helper that determines the length of the local diagonal and the length
% after adding  zeros

% One worker cannot pre- and append zeros.
assert(~(prependZ && appendZ));
% Standard case where no zeros are added
len = max(0, min(sizeLocalA(1), sizeLocalA(2)-localD) - max(1, 1-localD) + 1);
lenWithZeros = len;
% Calculate dominant dimension to determine if zeros are add top/bottom or
% left/right
[~, dominantDimension] = max(distA.Cached.GlobalSize);
if prependZ
    if dominantDimension~=distA.Dimension
        lenWithZeros = sizeLocalA(distA.Dimension);
    else
        lenWithZeros = min(sizeLocalA(distA.Dimension)+abs(localD), ...
            sizeLocalA(3-distA.Dimension));
    end
end
if appendZ
    if dominantDimension~=distA.Dimension
        lenWithZeros = sizeLocalA(distA.Dimension);
    else
        lenWithZeros = min(sizeLocalA(3-distA.Dimension)+...
            (2*distA.Dimension-3)*localD, sizeLocalA(3-distA.Dimension));
    end
end
end

function croppedColumn = iExtractAndCropColumn(distCol, localB, i, length)
% Helper that extracts a certain column from a codistributed matrix, crops
% it to the desired length and forms a new codistributed column with no
% data communication

% Create new cropped partition and crop local parts
localCol = distCol.hGetColumn(localB, i);
if distCol.Dimension==1
    if sum(distCol.Partition(1:spmdIndex-1))>=length
        localDim = 0;
        localCol = zeros(0,1);
    elseif sum(distCol.Partition(1:spmdIndex))<=length
        localDim = distCol.Partition(spmdIndex);
    else
        localDim = length-sum(distCol.Partition(1:spmdIndex-1));
        localCol = localCol(1:localDim,1);
    end
else
    if ~isempty(localCol)
        localDim = 1;
        localCol = localCol(1:length);
    else
        localDim = 0;
        localCol = zeros(length, 0);
    end
end
% Communicate new partition and create cropped codistributed array
croppedPartition = spmdReduce(@horzcat, localDim);
distCroppedCol = codistributor1d(distCol.Dimension, croppedPartition, ...
    [length, 1]);
croppedColumn = codistributed.pDoBuildFromLocalPart(localCol, ...
    distCroppedCol);
end
