function varargout = 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)
%   
%   First input sparse gpuArray is not supported.
%   
%   Example:
%    
%       n = 1e3;
%       b = gpuArray.ones(n,1);
%       A = SPDIAGS([b -2*b b], -1:1, 1e3, 1e3);
%   
%   See also SPDIAGS, GPUARRAY


%   Copyright 2016-2017 The MathWorks, Inc.

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

% Call spdiags on the host if none of the data inputs are gpuArray
if ~isa(varargin{1},'gpuArray') && ...
        ( nargin~=3 || (nargin==3 && ~isa(varargin{3},'gpuArray')) )
    [varargin{:}] = gather(varargin{:});
    nout = max(1,nargout);
    [varargout{1:nout}] = spdiags(varargin{:});
    return
end

if issparse(varargin{1}) && isa(varargin{1}, 'gpuArray')
    error(message('parallel:gpu:array:SpdiagsSparseFirstInput'));
end

if nargin <= 2
    % Extract diagonals
    A = varargin{1};
    if nargin == 1
        % Find all nonzero diagonals
        [i, j] = find(A);
        % Compute d = unique(d) without extra function call
        d = gather(sort(j-i));
        d = d(diff([-inf; d(:)])~=0);
        d = d(:);
    else
        % Diagonals are specified
        d = gather(varargin{2});
        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
    end
    [m, n] = size(A);
    numberOfDiags = length(d);
    % Create output data
    B = zeros(min(m, n), numberOfDiags, 'like', A); % A is always full at this point
    for k = 1:numberOfDiags
        if m >= n
            i = max(1, 1+d(k)) : min(n, m+d(k));
        else
            i = max(1, 1-d(k)) : min(m, n-d(k));
        end
        B = iShortSubsasgnAPI(B, iDiagK(A, d(k)), i, k);
    end
    varargout{1} = B;
    if nargout>1
        varargout{2} = d;
    end
else
    % If B lives on the host, keep it there in order to support sparse data
    B = varargin{1};
    d = gather(varargin{2});
    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
    numberOfDiags = length(d);
    if nargin == 3 % Replace specified diagonals
        A = varargin{3};
        % Send A to remote if necessary
        if ~isa(A, 'gpuArray')
            A = gpuArray(A);
        end
        [m, n] = size(A);
    else % Create new matrix with specified diagonals
        m = gather(varargin{3});
        n = gather(varargin{4});
        % Check that requested dimension are valid (non-negative, finite, real,
        % integer), otherwise throw the same error that core MATLAB throws
        if any(~isfinite([m n])) || any(~isreal([m n])) || any([m n]<0) ...
                || any(ceil([m n])~=[m n]) || size([m n],2)~=2
            ME = MException(message('MATLAB:sparsfcn:negSizeOfSparse'));
            throwAsCaller(ME);
        end
    end
    
    % Error if output B would be a sparse logical gpuArray
    if islogical(B) || (exist('A','var') && islogical(A))
        error(message('parallel:gpu:array:SpdiagsLogicalThreeOrFourInputs'));
    end
    
    % Check size of matrix B (should be min(m,n)-by-p)
    % For backwards compatibility, only error if the code would
    % previously have errored out in the indexing expression.
    maxIndexRows = max(max(1, 1-d), min(m, n-d)) + (m>=n)*d;
    maxIndexRows(max(1, 1-d) > min(m, n-d)) = 0;
    if any(maxIndexRows > size(B, 1)) || numberOfDiags > size(B, 2)
        if nargin == 3
            error(message('MATLAB:spdiags:InvalidSizeBThreeInput'));
        else
            error(message('MATLAB:spdiags:InvalidSizeBFourInput'));
        end
    end
    
    % Compute indices and values of sparse matrix with given diagonals
    % Compute lengths of diagonals:
    len = max(0, min(m, n-d) - max(1, 1-d) + 1);
    len = [0; cumsum(len)];
    
    a = gpuArray.zeros(len(numberOfDiags+1), 3);
    for k = 1:numberOfDiags
        % Insert new d(k)-th diagonal to compact form
        i = (max(1, 1-d(k)):min(m, n-d(k)))';
        a = iShortSubsasgnAPI(a, ...
            [i i+d(k) iShortSubsrefAPI(B, i+(m>=n)*d(k), k )], ...
            (len(k)+1):len(k+1), ':');
    end
    
    % Append unmodified data from original A if present
    if exist('A','var') && nnz(A) ~= 0
        % Process A in compact form
        [i, j, aold] = find(A);
        i = i(:);
        j = j(:);
        aold = aold(:);
        
        idx = ~any((j-i) == d', 2);
        i = iShortSubsrefAPI(i, idx, ':');
        j = iShortSubsrefAPI(j, idx, ':');
        aold = iShortSubsrefAPI(aold, idx, ':');
        
        % Combine new diagonals and non-diagonal entries of original matrix
        a = [a; [i j aold]];
    end
    % Create sparse output
    varargout{1} = sparse(iShortSubsrefAPI(a, ':', 1), iShortSubsrefAPI(a, ':', 2), ...
        iShortSubsrefAPI(a, ':', 3), m, n);
end
end

function variable = iShortSubsasgnAPI(variable, data, idx1, idx2)
% Helper to shorten subsasgn API that works for gpuArray and host variables
if iIsValidSubsasgnIndex(idx1) && iIsValidSubsasgnIndex(idx2)
    % If data is sparse data on the host, make it full in order to be used
    % with SUBSASGN
    variable = subsasgn(variable, substruct('()', {idx1, idx2}), full(data));
end
end

function data = iShortSubsrefAPI(data, idx1, idx2)
% Helper to shorten subsref API that works for gpuArray and host variables
if ~isempty(idx1) && ~isempty(idx2)
    data = subsref(data, substruct('()', {idx1, idx2}));
else
    data = zeros(0, 1, 'like', data);
end
end

function tf = iIsValidSubsasgnIndex(idx)
tf = ~isempty(idx) && all(idx>=0);
end

function D = iDiagK(X, k)
% Helper to extract the k-th diagonal of X, even if X is a vector.
if ~isvector(X)
    D = diag(X, k);
    D = D(:);  %Ensure column vector is returned for empty X.
else
    if ~isempty(X) && 0 <= k && 1+k <= size(X,2) % Row
        D = iShortSubsrefAPI(X, 1, 1+k);
    elseif ~isempty(X) && k < 0 && 1-k <= size(X,1) % Column
        D = iShortSubsrefAPI(X, 1-k, 1);
    else
        D = zeros(0, 1, 'like', X);
    end
end
end