function A = subsasgnParens(A, aidx, B)
% subsasgnParens performs paren assignment, i.e., (), for distributed inputs

% Copyright 2018-2022 The MathWorks, Inc.

if numel(aidx)>1
    if aidx(2).type ~= "."
        error(message("MATLAB:subsDotMustFollow"));
    end
    % If multiple levels of indexing, first get the variable and update it.
    Avar = subsref(A, aidx(1));
    B = subsasgn(Avar, aidx(2:end), B);
    % We are now left with just the outer level of indexing.
    aidx = aidx(1);
end

% Gather all distributed indices and find all logical indices.
for k = 1:length(aidx.subs)
    if isa(aidx.subs{k}, 'codistributed')
        aidx.subs{k} = gather(aidx.subs{k});
    end
    if islogical(aidx.subs{k})
        aidx.subs{k} = find(aidx.subs{k});
    end
end

if ~isa(A,'codistributed') && isa(B,'codistributed')
    if isempty(A)
        % Use SUBSASGN for creation:
        % notDefined(...,subs,...) = codistributed
        U(aidx.subs{:}) = gather(B);
        A = codistributed(U);
        return
    else
        error(message('MATLAB:legacy_two_part:UnableToConvert', ...
            class(B), class(A), ...
            getString(message('MATLAB:legacy_two_part:invalidConversion', ...
            class(A), class(B)))));
    end
end

% This implementation only supports codistributor1d.
codistributed.pVerifyUsing1d('subsasgn', A, B);

% A(...,s,...) = B

if length(aidx.subs) >= ndims(A)
    % Ensure that aidx and indexingSize include the distribution dimension of A.
    A = iSubsasgnParensGeneral(A, size(A), aidx, B);
    
elseif length(aidx.subs) == 1
    A = iSubsasgnParensLinearSubs(A, size(A), aidx.subs{1}, B);
    
else
    A = iSubsasgnParensTooFewSubs(A, size(A), aidx, B);
end
end

function A = iSubsasgnParensGeneral(A, indexingSize, aidx, B)
% Implementation of subsasgn when there at least as many subscripts as
% dimensions.
aDist = getCodistributor(A);
dim = aDist.Dimension;
indexingSize(end+1:dim) = 1;
aidx.subs(end+1:dim) = {1};

if ~parallel.internal.indexing.isValidIndexVector(aidx.subs{dim})
    error(message("MATLAB:badsubscript", ...
        getString(message("MATLAB:matrix:badTypeIndicesPosition", dim))));
    
end

% For a 0x0 input ("uninitialized array") all colon literals should be
% converted into 1:size(rhs,dim) unless one or more subscripts is empty.
if isequal(size(A), [0 0])
    % Check for any empty subs
    hasEmptyIdx = false;
    for ii=1:numel(aidx.subs)
        if ~parallel.internal.indexing.isColonLiteral(aidx.subs{ii}) ...
                && parallel.internal.indexing.getMaxIndex(aidx.subs{ii})==0
            hasEmptyIdx = true;
            break
        end
    end
    % Only expand colon indices to RHS size if no empties
    if ~hasEmptyIdx
        for ii=1:numel(aidx.subs)
            if parallel.internal.indexing.isColonLiteral(aidx.subs{ii})
                aidx.subs{ii} = matlab.internal.ColonDescriptor(1,1,size(B,ii));
            end
        end
    end
end

% Now convert some types we don't want to handle directly
if parallel.internal.indexing.isColonLiteral(aidx.subs{dim})
    aidx.subs{dim} = matlab.internal.ColonDescriptor(1,1,indexingSize(dim));
elseif ischar(aidx.subs{dim})
    aidx.subs{dim} = double(aidx.subs{dim});
end

s = aidx.subs{dim};
maxIdx = parallel.internal.indexing.getMaxIndex(s);
if maxIdx > size(A,dim)
    % Subscript exceeds distributed dimension, expand last partition
    if ~isscalar(B)
        error(message('parallel:distributed:SubsasgnNotYetNonScalar'))
    end
    [A, aDist] = iExpandInDistributionDimension(A, aDist, aidx, maxIdx);
end

if isscalar(B)
    
    % codistributed(...,subs,...) = scalar
    
    if isa(B,'codistributed')
        B = gather(B);
    end
    
    % Work out the indices in this partition relative to the start of
    % the partition.
    [e, f] = globalIndices(A, aDist.Dimension, spmdIndex);
    aidx.subs{dim} = s(parallel.internal.indexing.getIndicesInPartition(s, e, f)) - e + 1;
    
    A = iBuildWithExpandedGlobalSize(subsasgn(getLocalPart(A),aidx,B), ...
        getCodistributor(A));
else
    distAt = iGetCodistributorOfDestination(A, aidx);
    
    if prod(distAt.Cached.GlobalSize) ~= numel(B)
        error(message('parallel:distributed:SubsasgnSizeMismatch'));
    end
    
    assert(isa(A,'codistributed'),'At this point, A must be codistributed.');
    if ~isa(B,'codistributed')
        B = reshape(B, distAt.Cached.GlobalSize);
        % Use the codistributed build to make sure each worker only has the
        % elements of B it needs (we don't really need a codistributed here
        % but since codistributed knows how to get the data for each worker
        % we make one instead of duplicating the logic).
        % NB: for use within DRANGE it is imperative that this is done with
        % no communication, so bypass the codistributed constructor.
        [Blp, Bcodist] = distAt.hBuildFromReplicatedImpl(0, B);
        B = codistributed.build(Blp, Bcodist, 'noCommunication');
    else % codistributed(B)
        if ~isequal(distAt.Cached.GlobalSize, size(B))
            B = reshape(B, distAt.Cached.GlobalSize);
        end
        B = redistribute(B, distAt);
    end
    
    % codistributed(...,subs,...) = codistributed
    
    bDist = getCodistributor(B);
    j = globalIndices(B, bDist.Dimension, spmdIndex);
    [e,f] = globalIndices(A, aDist.Dimension, spmdIndex);
    if all(s(j) >= e & s(j) <= f)
        
        % Important special case, no communication required.
        
        aidx.subs{dim} = s(j)-e+1;
        Aloc = subsasgn(getLocalPart(A),aidx,getLocalPart(B));
        
    else
        
        % General case.
        
        Aloc = getLocalPart(A);
        bidx = substruct('()',repmat({':'},1,ndims(B)));
        mwTag = 31898;
        for p = 1:spmdSize
            j = globalIndices(B, bDist.Dimension, p);
            for q = 1:spmdSize
                [e,f] = globalIndices(A, aDist.Dimension, q);
                k = find(s(j) >= e & s(j) <= f);
                if ~isempty(k) && ~isempty(j)
                    aidx.subs{dim} = s(j(k))-e+1;
                    bidx.subs{dim} = j(k)-j(1)+1;
                    if p == q && p == spmdIndex
                        Aloc = subsasgn(Aloc,aidx,subsref(getLocalPart(B),bidx));
                    elseif p == spmdIndex
                        spmdSend(subsref(getLocalPart(B),bidx),q,mwTag);
                    elseif q == spmdIndex
                        Aloc = subsasgn(Aloc,aidx,spmdReceive(p,mwTag));
                    end
                end
            end
        end
        
    end
    A = iBuildWithExpandedGlobalSize(Aloc, getCodistributor(A));
end
end

function A = iSubsasgnParensTooFewSubs(A, indexingSize, aidx, B)
% Fewer subscripts than dimensions. Use reshape to index in the reduced
% dimensionality.

% Work out the size as the subscripts want it ("fold" the trailing dimensions)
numSubs = numel(aidx.subs);
origSize = indexingSize;
indexingSize = [indexingSize(1:numSubs-1), prod(indexingSize(numSubs:end))];

% Check for growing in the folded dimension (':' never grows)
if ~parallel.internal.indexing.isColonLiteral(aidx.subs{end})
    maxIdx = parallel.internal.indexing.getMaxIndex(aidx.subs{end});
    if maxIdx > indexingSize(end)
        iThrowAmbiguousGrowingMessage();
    end
end

% We need to compute the new size for the reshape-back operation.
newSize = origSize;
for idx = 1:(numSubs-1)
    % These dimensions might grow - but not if they're a literal colon.
    thisSub = aidx.subs{idx};
    if ~parallel.internal.indexing.isColonLiteral(thisSub)
        newSize(idx) = max(indexingSize(idx), parallel.internal.indexing.getMaxIndex(thisSub));
    end
end

% Reshape to the indexing size. This is communication free if the
% distribution dimension is not one of the reshaped dimensions but
% otherwise might be expensive.
% TODO: make the reshape notional, not actual, when we can!
A = reshape( ...
    iSubsasgnParensGeneral(reshape(A, indexingSize), indexingSize, aidx, B), ...
    newSize);
end

function A = iSubsasgnParensLinearSubs(A, sizea, subs, B)
% subsasgnParensLinear performs paren assignment using a single linear
% index.

if ~parallel.internal.indexing.isValidIndexVector(subs)
    error(message('MATLAB:badsubscript', ...
        getString(message('MATLAB:badsubscriptTextRange'))));
end
if parallel.internal.indexing.isColonLiteral(subs)
    subs = matlab.internal.ColonDescriptor(1,1,prod(sizea));
end
maxIdx = parallel.internal.indexing.getMaxIndex(subs);
% Only consider growing if the maximum index exceeds the number of elements we
% already have.
if maxIdx > prod(sizea)
    newSz = parallel.internal.indexing.linearIndexAssignmentResultSize(sizea, maxIdx);
    A = iExpandForLinearIndexing(A, sizea, newSz, maxIdx);
    sizea = newSz;
end
numelInPriorDims = prod(sizea(1:end-1));

if isscalar(B)
    
    % codistributed(s) = scalar
    
    if isa(B,'codistributed')
        B = gather(B);
    end
    
    % This algorithm needs A to be distributed by last dimension.
    % Not efficient for A that is not distributed by last dimension.
    oldDistA = getCodistributor(A);
    A = redistribute(A, codistributor1d(ndims(A)));
    
    % Work out the indices in this partition relative to the start of
    % the partition.
    [firstIdx,lastIdx] = globalIndices(A, ndims(A), spmdIndex);
    firstIdx = numelInPriorDims*(firstIdx-1)+1;
    lastIdx = numelInPriorDims*lastIdx;
    aidx = substruct("()", {subs(parallel.internal.indexing.getIndicesInPartition(subs, firstIdx, lastIdx))-firstIdx+1});
    
    A = iBuildWithExpandedGlobalSize(subsasgn(getLocalPart(A), aidx, B), ...
        getCodistributor(A));
    A = redistribute(A, oldDistA);
    
else
    
    % codistributed(s) = array
    
    if numel(subs) ~= numel(B)
        error(message('parallel:distributed:SubsasgnSizeMismatch'));
    end
    
    if isa(B,'codistributed')
        B = gather(B);
    end
    
    % This algorithm needs A to be distributed by last dimension.
    % Not efficient for A that is not distributed by last dimension.
    oldDistA = getCodistributor(A);
    A = redistribute(A, codistributor1d(ndims(A)));
    
    % Work out the indices in this partition relative to the start of
    % the partition.
    [firstIdx,lastIdx] = globalIndices(A, ndims(A), spmdIndex);
    firstIdx = numelInPriorDims*(firstIdx-1)+1;
    lastIdx = numelInPriorDims*lastIdx;
    j = parallel.internal.indexing.getIndicesInPartition(subs, firstIdx, lastIdx);
    aidx = substruct("()", {subs(j)-firstIdx+1});
    
    A = iBuildWithExpandedGlobalSize(subsasgn(getLocalPart(A), aidx, B(j)), ...
        getCodistributor(A));
    A = redistribute(A, oldDistA);
    
end

end

function A = iExpandForLinearIndexing(A, oldSz, newSz, maxIdx)
% Expansion for linear indexing cases.
% Need to resize A to newSz. For valid single-subscript growing, we must
% either be expanding an empty to a vector, or a pre-existing vector.

oldDistA = getCodistributor(A);
oldDistDim = oldDistA.Dimension;
if any(oldSz == 0)
    % A was empty - let's just build a new A from scratch.
    A = distributedutil.Allocator.createCodistributed(newSz, getLocalPart(A));
    newDistA = getCodistributor(A);
    newDistDim = newDistA.Dimension;
else
    % expand (generalised) vector case - simply augment the local part on the last
    % worker, after first ensuring A is distributed in the vector dimension.
    [~, vecDim] = max(oldSz);
    newDistDim = vecDim;
    A = redistribute(A, codistributor1d(newDistDim));
    % We need an as-if substruct for iExpandInDistributionDimension.
    aidxForExpansion = substruct('()', num2cell(newSz));
    A = iExpandInDistributionDimension(A, getCodistributor(A), aidxForExpansion, maxIdx);
end
% Ensure we revert the distribution dimension
if oldDistDim ~= newDistDim
    A = redistribute(A, codistributor1d(oldDistDim));
end
end

function [A, aDist] = iExpandInDistributionDimension(A, aDist, aidx, newSize)
% Expand the size of A in the distribution dimension
dim = aDist.Dimension;
expand = newSize - size(A,dim);
Aloc = getLocalPart(A);
p = aDist.Partition;
p(end) = p(end) + expand;
s = aidx.subs{dim};
if spmdIndex == spmdSize
    [e,f] = partitionIndices(p);
    f = f + expand;
    lastidx = aidx;
    lastidx.subs{dim} = s(parallel.internal.indexing.getIndicesInPartition(s, e, f))-e+1;
    zeroElem = distributedutil.Allocator.create([1,1], Aloc);
    Aloc = subsasgn(Aloc, lastidx, zeroElem);
else
    % Assignment on the last lab might have caused expansion in
    % many dimensions - not just the distributed dimension.
    % Therefore on labs other than the last one we have to check if
    % expansion happens and if so we need to grow the local part
    % accordingly.
    [e,f] = partitionIndices(p);
    if e <= f % if there is any local data
        
        % check all non-distributed indices if they are being expanded
        
        expandF = false;
        for i = 1:length(aidx.subs)
            if i == dim, continue; end % skip distributed dimension
            
            ss = aidx.subs{i};
            if ~parallel.internal.indexing.isColonLiteral(ss)
                thisMaxIndex = parallel.internal.indexing.getMaxIndex(ss);
                if thisMaxIndex > size(A, i)
                    expandF = true;
                    break;
                end
            end
        end
        if expandF
            lastidx = aidx;
            lastidx.subs{dim} = 1;
            zeroElem = distributedutil.Allocator.create([1,1], Aloc);
            Aloc = subsasgn(Aloc,lastidx, zeroElem);
        end
    end
end

% Expand along the dimension of distribution if necessary.
globalSize = aDist.Cached.GlobalSize;
globalSize(dim) = globalSize(dim) + expand;
aDist = codistributor('1d', dim, p, globalSize);

[A, aDist] = iBuildWithExpandedGlobalSize(Aloc, aDist);
end

% Helper function that expands the codistributor global size property to
% fit the local part.
function [D, codist] = iBuildWithExpandedGlobalSize(L, codist)

globalSize = codist.Cached.GlobalSize;

localSize = size(L);
numDims = max(numel(localSize), numel(globalSize));
localSize(end+1:numDims) = 1;

globalSize(end+1:numDims) = 1;
globalSize = max(globalSize, localSize);
codist = codistributor1d(codist.Dimension, codist.Partition, globalSize);

D = codistributed.build(L, codist, 'noCommunication');
end

function [distB] = iGetCodistributorOfDestination(A, aidx)
% Helper that determines the codistributor of B to match the SUBSASGN
% destination in A to perform A(aidx) = B efficiently. In detail, the
% codistributor for B is used to redistribute/reshape B such that the
% location of its local data matches the location of the local data of A,
% into where it is assigned to. This helper avoids the use of
% SUBSREF(A, aidx) to determine this information.

lengthsubs = length(aidx.subs);
nda = ndims(A);
aDist = getCodistributor(A);
% Set distribution dimension of B accordingly
if lengthsubs >= nda
    bDistDim = aDist.Dimension;
else
    bDistDim = lengthsubs;
end

% Determine global size of B
globalSizeB = zeros(1,lengthsubs);
for k = 1:length(globalSizeB)
    if parallel.internal.indexing.isColonLiteral(aidx.subs{k})
        globalSizeB(k) = size(A,k);
    else
        globalSizeB(k) = numel(aidx.subs{k});
    end
end

% Partition of B needs to be created without any communication such that
% SUBSASGN can possibly be used inside a for-loop over distributed range
partitionB = zeros(1, spmdSize);
for i=1:spmdSize
    [partStart, partEnd] = partitionIndices(aDist.Partition, i);
    partitionB(i) = numel(parallel.internal.indexing.getIndicesInPartition(aidx.subs{bDistDim}, ...
        partStart, partEnd));
end
% Create codistributor
distB = codistributor1d(bDistDim, partitionB, globalSizeB);
end

function iThrowAmbiguousGrowingMessage()
% MATLAB uses the ID indexed_matrix_cannot_be_resized but overrides the message.
msg = message("MATLAB:matrix:ambiguousArrayGrowth");
err = MException("MATLAB:indexed_matrix_cannot_be_resized", getString(msg));
throwAsCaller(err);
end
