function [sz, varargout] = size(obj, dim, varargin)
%SIZE   Size of dlarray.
%   M = SIZE(D)
%   [M1,M2] = SIZE(D)
%   [M1,M2,M3,...,MN] = SIZE(D)
%   M = SIZE(D,DIM)
%   M = SIZE(D,DIM1,DIM2,...,DIMN)
%   [M1,M2,...,MN] = SIZE(D,...)
%
%   If D is formatted, SIZE(D) returns a vector of length equal to the
%   number of dimension labels, even if some of the dimension labels refer
%   to trailing singleton dimensions.
%
%   See also SIZE, DLARRAY, DLARRAY/NDIMS.

%   Copyright 2019 The MathWorks, Inc.

if nargin == 1
    fd = obj.FormattedData;
    % Find out where to map the size of the stored data into
    viewToData = obj.ViewToDataMap;
    
    if isempty(viewToData)
        sz = size(fd.Data);
        if ~isempty(fd.DimensionMetadata)
            nd = ndims(obj);  
            sz(end+1:nd) = 1;
        end
    else
        nd = ndims(obj);
        sz = size(fd.Data, viewToData(1:nd));
    end
    
    if nargout > 1
        varargout = cell(1, nargout-1);
        for ii=2:min(nargout, length(sz))
            varargout{ii-1} = sz(ii);
        end
        
        if nargout < length(sz)
            % [a, b] = size(ones(3, 4, 2)) returns a = 3, b = 8
            varargout{end} = prod(sz(nargout:end));
        elseif nargout > length(sz)
            % [a, b, c] = size(ones(3, 4)) returns a = 3, b = 4, c = 1
            varargout(length(sz):end) = {1};
        end
        
        % First scalar output returned as variable sz
        sz = sz(1);
    end
    
else
    if nargin > 2
        % size(x, 1, 3, 4): Compute vector [1 3 4] from these inputs.
        args = [{dim}, varargin];
        dim = zeros(size(args));
        for ii=1:length(args)
            arg = args{ii};
            if ~isscalar(arg)
                error(message('MATLAB:getdimarg:dimensionMustBePositiveInteger'));
            elseif isobject(arg) && isa(arg, 'dlarray')
                arg = gather(extractdata(arg));
            end
            dim(ii) = arg;
        end
    elseif isobject(dim) && isa(dim, 'dlarray')
        dim = gather(extractdata(dim));
    end
    
    if ~isnumeric(dim) && ~islogical(dim)
        error(message('MATLAB:getdimarg:dimensionMustBePositiveInteger'));
    end
    
    if ~isa(obj, 'dlarray')
        % Only (some of) the dim inputs are dlarray.
        [sz, varargout{1:nargout-1}] = size(obj, dim);
        return
    end
    
    data = obj.FormattedData.Data;
    viewToData = obj.ViewToDataMap;
    if isempty(viewToData)
        sz = size(data, dim);
    else
        sz = zeros(size(dim));
        nd = length(viewToData);
        for ii=1:length(sz)
            dimii = dim(ii);
            if dimii > nd
                sz(ii) = 1;
            else
                sz(ii) = size(data, viewToData(dimii));
            end
        end
    end
    
    if nargout > 1
        if length(dim) ~= nargout
            error(message('MATLAB:size:NumOutNotEqualNumDims'));
        end
        
        varargout = cell(1, nargout-1);
        for ii=2:nargout
            varargout{ii-1} = sz(ii);
        end
        sz = sz(1);
    end
end
end
