function B = hSubsrefParensImpl(aDist, A, aidx, tableTmpl)
% hSubsrefParensImpl implements subsref for parens indexing providing a
% single index or multiple indices for arrays using codistributor1d scheme.
% First additional input varargin{1} contains a table template if provided.
%
% Using a single index:
% B = A(X)
% X can be: numeric, or logical vector (replicated)
% X cannot be: ':'
%
% Using multiple indices:
% B = A(X1, X2, ...)
% each Xi can be: ':', numeric, logical (replicated)

%   Copyright 2018-2023 The MathWorks, Inc.

% B = A(subs). Minimal communication between workers.
isFirstSubColonLiteral = parallel.internal.indexing.isColonLiteral(aidx.subs{1});
orgLengthSubs = length(aidx(1).subs);
[A, Aloc, aidx, partA, bDistDim, sizeb] = iPrepIndexAndCalcSizeOfB(A, aidx, aDist);

% If the result is going to be empty we can avoid most of the work for
% homogeneous arrays (i.e. not tables)
if prod(sizeb)==0 && ~(istable(A) || istimetable(A))
    codistr = codistributor1d(bDistDim, ...
        codistributor1d.defaultPartition(sizeb(bDistDim)), sizeb);
    Bloc = iMakeEmptyLocalPart(Aloc, codistr);
    B = codistributed(Bloc, codistr, "undocumented:ConstructFromLocalParts");
    return;
end

% Error check subscripts in the distribution dimension here because all
% workers are still working with the replicated index. This allows them all
% to throw an error at the same time.
subInBDistDim = aidx.subs{bDistDim};
% TODO(g1831002): Fix this here.  We need to colonize the index in the distribution
% dimension when there are multiple indices.
if nnz(size(subInBDistDim)~=1)>1
    error(message('parallel:distributed:SubsrefBadSubscriptSize'))
end

% We need to error check the index in the distribution dimension of B ourselves
% because we will manipulate it further on.  The builtin subsref will error
% check the indices in the other dimensions.
if ~parallel.internal.indexing.isValidIndexVector(subInBDistDim)
    error(message('parallel:array:SubsrefInvalidSubscript'))
end

sizea = size(A);
% Let smax store the upper bound on the permissible index value.
smax = sizea;
smax(orgLengthSubs) = prod(smax(orgLengthSubs:end));
smax(end+1:bDistDim) = 1;
% TODO(g1831002): Can this be simplified to any(subInBDistDim(:) > smax(bDistDim))?
if (orgLengthSubs == 1 && parallel.internal.indexing.getMaxIndex(subInBDistDim) > smax(1)) ...
        || (orgLengthSubs > 1 && parallel.internal.indexing.getMaxIndex(subInBDistDim) > smax(bDistDim))
    error(message('parallel:array:SubsrefSubscriptExceedsDim'))
end

% Subscript ranges
% Calculate vectors of length spmdSize storing the start and end indices of A in
% the distribution dimension that reside on each worker.
endA = cumsum(partA);
startA = endA - partA + 1;
% If the user used linear indexing, we have already colonized A.  If the user
% used > 1 index but < ndims(A) indices, we fold the last dimensions of A, so we
% have to update the start and end indices accordingly.
lengthsubs = length(aidx.subs);
% Note that at this point, lengthsubs is always > 1 as the output of
% iPrepIndexAndCalcSizeOfB.
foldLastDimensions = lengthsubs < length(sizea);
if foldLastDimensions
    strideInLastIndexDim = prod(sizea(lengthsubs:end-1));
    startA = strideInLastIndexDim*(startA-1)+1;
    endA = strideInLastIndexDim*endA;
end

% Determine partition of B so that little communication is required.

% Let the partition of B be such that each worker stores the same elements
% of A as it has.
% TODO(g1831002): Figure out how this maps into the communicating subsref.
partB = zeros(1,spmdSize);
for k = 1:spmdSize
    partB(k) = iNumIndicesInPartition(subInBDistDim, startA(k), endA(k));
end
endB = cumsum(partB);
startB = endB-partB+1;

% Not all types support ColonDescriptor. The distribution dimension is
% handled explicitly, but we need to expand out any colon descriptors in
% the other dimensions now.
if ~iSupportsColonDescriptor(Aloc)
    for ii=1:numel(aidx.subs)
        if ii~=bDistDim && parallel.internal.indexing.isColonDescriptor(aidx.subs{ii})
            aidx.subs{ii} = aidx.subs{ii}(:)'; % Convert to column vector
        end
    end
end

if istable(A) || istimetable(A)
    % Tables are special as we must get the right variable names too. Note
    % that all workers must get here, even if all data stays on the worker
    % so we handle this case first.
    % First, check the provided template
    assert(istable(tableTmpl) || istimetable(tableTmpl), ...
        "Provided template is not a table or timetable.")
    % Get a template for the output table
    tmpl = subsref(tableTmpl, substruct('()', {':', aidx.subs{2}}));
    if bDistDim==2
        % If workers have different variables, sub-select only those for this worker
        tmpl = subsref(tmpl, substruct('()', {':', startB(spmdIndex):endB(spmdIndex)}));
    end
    sizeBloc = sizeb;
    sizeBloc(bDistDim) = partB(spmdIndex);
    Bloc = distributedutil.Allocator.create(sizeBloc, tmpl);
    Bloc = iGetBlocWithCommunications(Aloc, Bloc, aidx, endA, startA, endB, startB, bDistDim, sizea);
    % Table preserves complexity of table variables when the first
    % subscript is a colon literal.
    if isFirstSubColonLiteral
        Bloc = distributedutil.hEnsureComplexityMatches(Bloc, tmpl);
    end
    % Might need to fix up rowtimes if first worker was empty (RowTimes are
    % taken from first worker), but no need to recurse into fields.
    if istimetable(Bloc) && bDistDim==2 && partB(1)==0
        if parallel.internal.indexing.isColonLiteral(aidx.subs{1})
            Bloc.Properties.RowTimes = Aloc.Properties.RowTimes;
        else
            Bloc.Properties.RowTimes = Aloc.Properties.RowTimes(aidx.subs{1});
        end
    end
else
    
    % Check if no communication is required
    % Determine q and r so that A(...,subInBDistDim(q),...) is on this worker and
    % B(..., q(r),...)  is on this worker.
    
    % Find all the subscripts into A in the dimension bDistDim that refer to values
    % that we have on this worker.
    q = parallel.internal.indexing.getIndicesInPartition(subInBDistDim, startA(spmdIndex), endA(spmdIndex));
    % Find those elements of q that will end up in elements of B that we store on
    % this worker.
    r = parallel.internal.indexing.getIndicesInPartition(q, startB(spmdIndex), endB(spmdIndex));
        
    if isinparfor
        % Does the indexing into my local part create all of B? We ignore
        % where it should end up and put all of it here.
        if length(r) == sum(partB)
            aidx.subs{bDistDim} = subInBDistDim(q(r))-startA(spmdIndex)+1;
            % Since we are inside a for-drange statement and are returning
            % a regular MATLAB array, we don't inspect or modify the
            % sparsity of B. Rather, we just return what regular MATLAB
            % indexing gives us.
            B = subsref(Aloc, aidx);
            return
        else
            error(message('parallel:distributed:SubsrefInparfor'))
        end
        
    elseif length(r) == partB(spmdIndex)
        % The indexing into my local part all ends up on this worker. We
        % therefore don't need any message passing.
        aidx.subs{bDistDim} = subInBDistDim(q(r))-startA(spmdIndex)+1;
        Bloc = subsref(Aloc, aidx);
        
    else
        % General paren subsref
        sizea = size(A);
        sizeBloc = sizeb;
        sizeBloc(bDistDim) = partB(spmdIndex);
        Bloc = distributedutil.Allocator.create(sizeBloc, Aloc);
        Bloc = iGetBlocWithCommunications(Aloc, Bloc, aidx, endA, startA, endB, startB, bDistDim, sizea);
    end
end

% Result is a codistributed array
codistr = codistributor1d(bDistDim, partB, sizeb);
B = codistributed(Bloc, codistr, "undocumented:ConstructFromLocalParts");
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function Bloc = iMakeEmptyLocalPart(Aloc, codistr)
% We can't use 'like' for non-numeric types (we would lose extra data like
% date formats and locales), so instead use empty subsref and reshape to
% get an empty of the right size.
bsize = hLocalSize(codistr);
proto = Aloc([]);
Bloc = reshape(proto, bsize);
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [A, Aloc, aidx, partA, bDistDim, sizeb] = iPrepIndexAndCalcSizeOfB(A, aidx, aDist)
% Pre-process aidx and calculate the size of B as well as its distribution
% dimension.
% aidx is returned in a form such that aidx.subs is of length >= bDistDim.
lengthsubs = length(aidx.subs);
nda = ndims(A);
if lengthsubs > 1
    % More than one index used, e.g. A(X, Y), A(:, [1:10; 21:30], 5), etc.
    % 1) The number of indices may exceed ndims(A), but the extra indices better be
    % 1 or ':'.
    % 2) The number of indices may be less than ndims(A), in which case we fold the
    % last dimensions of A.
    % 3) The shape of the indices is ignored, and they are treated as the colonized
    % versions.
    %
    % We plan on letting B be distributed either across its last dimension, or
    % across aDist.Dimension. The number of dimensions of B will equal the
    % number of indices used into A.
    sizea = size(A);
    if lengthsubs >= nda
        % There are at least as many indices as ndims(A). The last "real" index
        % (i.e. those that aren't required to be equivalent to 1) is in the last
        % dimension of A.
        bDistDim = aDist.Dimension;
        % Ensure that aidx and sizea include the distribution dimension of A and B.
        sizea(end+1:aDist.Dimension) = 1;
        aidx.subs(end+1:aDist.Dimension) = {1};
        lengthsubs = length(aidx.subs);
        % Expand out literal colon into a full ColonDescriptor
        if parallel.internal.indexing.isColonLiteral(aidx.subs{bDistDim})
            aidx.subs{bDistDim} = matlab.internal.ColonDescriptor(1, 1, sizea(bDistDim));
        end
    else
        if aDist.Dimension ~= nda
            % Redistribute A so that the local parts are storing contiguous blocks of the
            % linear index of A.
            aDist = codistributor('1d',nda, codistributor1d.defaultPartition(sizea(nda)));
            A = redistribute(A, aDist);
            aDist = getCodistributor(A);
        end
        % At this point, A is distributed across its last dimension, but the last index
        % is in a dimension strictly less than that.
        bDistDim = lengthsubs;
        % Expand out literal colon into a full ColonDescriptor
        if parallel.internal.indexing.isColonLiteral(aidx.subs{bDistDim})
            aidx.subs{bDistDim} = matlab.internal.ColonDescriptor(1, 1, prod(sizea(bDistDim:end)));
        end
    end
    % The shape of the indices does not matter since more than one index was
    % specified.  We therefore convert all indices into column vectors.
    for k = 1:lengthsubs
        if ~parallel.internal.indexing.isColonLiteral(aidx.subs{k}) ...
                && ~parallel.internal.indexing.isColonDescriptor(aidx.subs{k})
            aidx.subs{k} = aidx.subs{k}(:);
        end
    end
    
    partA = aDist.Partition;
    Aloc = getLocalPart(A);
    % Figure out the global size of B.
    sizeb = zeros(1,lengthsubs);
    for k = 1:length(sizeb)
        if parallel.internal.indexing.isColonLiteral(aidx.subs{k})
            sizeb(k) = size(A,k);
        elseif parallel.internal.indexing.isColonDescriptor(aidx.subs{k})
            sizeb(k) = prod(size(aidx.subs{k})); %#ok<PSIZE> Seems to be only way to get number of elements from ColonDescriptor
        else
            sizeb(k) = length(aidx.subs{k});
        end
    end
else
    % Only one index used, i.e. linear indexing, such as A(1),
    % A([1:10; 21:30]), etc. Note that A(:) has already been taken care of.
    % When performing indexing with one input, i.e. A(X),, the output has the same
    % shape as the index matrix X, except when both A and X are vectors.  When
    % that is true, A(X) has the same shape as A, not X.
    sizea = size(A);
    sizeidx = size(aidx.subs{1});
    isAVectorNotScalar = iIsNDVector(A) && ~isscalar(A);
    isIdxVector = isvector(aidx.subs{1});
    
    % Abort early for empty (but non-vector) index
    if ~isIdxVector && prod(sizeidx)==0
        sizeb = sizeidx;
        Aloc = [];
        partA = [];
        bDistDim = [];
        return;
    end
    
    % Convert A into a distributed column vector.  Note that this may be extremely
    % inefficient: Colonize can result in an all-to-all communication, and this
    % we go this code path to retrieve a single element of A, such as, say, A(1).
    A = colonize(A);
    aDist = getCodistributor(A);
    Aloc = getLocalPart(A);
    partA = aDist.Partition;
    
    % The shape of B depends on whether A is a vector as well as how it is being
    % indexed.  If A and idx are both vectors, the shape of A is preserved.
    % If not, the shape of the index vector is preserved.
    preserveAShape = (isAVectorNotScalar && isIdxVector);
    
    if preserveAShape
        % Work out the output size. This should be the same as the size of A
        % but with the non-unity dimension replaced.
        bDistDim = iFirstNonSingletonDim(sizea);
        sizeb = ones(1,numel(sizea));
        sizeb(bDistDim) = prod(sizeidx);
    else
        % Shape comes from index vector. Scalar should be distributed along
        % second dimension to match default scheme, otherwise distribute
        % along the non-unity dimension.
        if isscalar(aidx.subs{1})
            bDistDim = 2;
        else
            bDistDim = iFirstNonSingletonDim(sizeidx);
        end
        sizeb = sizeidx;
    end
    
    
    % Now create a new subscript array with one subscript per output
    % dimension
    bsubs = aidx.subs{1};
    aidx.subs = repmat({1}, 1, numel(sizea));
    aidx.subs{bDistDim} = bsubs;
    
    % Reshape local part to the same shape as the output
    sizealoc = ones(1,numel(sizeb));
    sizealoc(bDistDim) = numel(Aloc);
    Aloc = reshape(Aloc, sizealoc);
end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function Bloc = iGetBlocWithCommunications(Aloc, Bloc, aidx, endA, startA, endB, startB, bDistDim, sizea)

subInBDistDim = aidx.subs{bDistDim};
% Subscript structure for use with builtin subsref and subsasgn

bidx = substruct('()',repmat({':'}, 1, length(sizea)));

% Double loop over workers.  Determine q and r so that A(...,subInBDistDim(q(r)),...)
% is on worker srcidx and B(...,q(r),...) is on worker destIdx.

mwTag = 32006;
for srcIdx = 1:spmdSize
    q = parallel.internal.indexing.getIndicesInPartition(subInBDistDim, startA(srcIdx), endA(srcIdx));
    if ~isempty(q)
        for destIdx = 1:spmdSize
            r = parallel.internal.indexing.getIndicesInPartition(q, startB(destIdx), endB(destIdx));
            if ~isempty(r)
                aidx.subs{bDistDim} = subInBDistDim(q(r)) - startA(srcIdx) + 1;
                bidx.subs{bDistDim} = q(r) - startB(destIdx) + 1;
                if destIdx == srcIdx && srcIdx == spmdIndex
                    Bloc = iMySubsasgn(Bloc, bidx, subsref(Aloc, aidx));
                elseif srcIdx == spmdIndex
                    spmdSend(subsref(Aloc, aidx), destIdx, mwTag)
                elseif destIdx == spmdIndex
                    Bloc = iMySubsasgn(Bloc, bidx, spmdReceive(srcIdx, mwTag));
                end
            end
        end
    end
end

end

function out = iMySubsasgn(out, idx, value)
% Helper to call subsasgn and preserve timetable RowTimes if necessary
out = subsasgn(out, idx, value);
if istimetable(out)
    out = parallel.internal.indexing.copyRowTimes(out, value, idx.subs{1}, idx.subs{2});
end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function isvec = iIsNDVector(data)
% Helper that returns true for any N-D vector. That is, any array where at
% most one dimension is not unity.
isvec = (nnz(size(data)~=1) <= 1);
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function dimIdx = iFirstNonSingletonDim(dims)
% Return the index of the first non-singleton dimension, or 1 if all unity.
dimIdx = find(dims~=1,1,'first');
if isempty(dimIdx)
    dimIdx = 1;
end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function tf = iSupportsColonDescriptor(arg)
% True if the input argument is a type that supports indexing using
% ColonDescriptor; false otherwise. All built-in numeric types support
% this, but most objects (e.g. table, datetime, etc.) don't.

if isobject(arg)
    tf = isa(arg, 'matlab.mixin.internal.indexing.LazyColon');
else
    % Built-in type, so should work
    tf = true;
end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function N = iNumIndicesInPartition(idx,partStart,partEnd)
% Determine how many indices are within one partition. The input index can
% be a vector of linear indices or a ColonDescriptor
if parallel.internal.indexing.isColonDescriptor(idx)
    % Work out the colon elements in the range
    newColon = parallel.internal.indexing.limitColonDescriptorToRange(idx,partStart,partEnd);
    N = prod(size(newColon)); %#ok<PSIZE>
else
    assert(isnumeric(idx));
    N = nnz((partStart <= idx) & (idx <= partEnd));
end
end
