function x = transferPortion( obj, rangeStruct, value )
%transferPortion - attempt to transfer an array portion efficiently

% Copyright 2009-2022 The MathWorks, Inc.

% Check the range against the size of the distributed object itself.
if isfield( rangeStruct, 'start' ) && isfield( rangeStruct, 'end' ) && ...
        all( rangeStruct.start(:) > 0 ) && ...
        all( rangeStruct.end <= size( obj ).' ) && ...
        length( rangeStruct.start ) == ndims( obj ) && ...
        isequal( size( rangeStruct.start ), size( rangeStruct.end ) )
    % ok
else
    error(message('parallel:distributed:BadRangeForTransferPortion'));
end

% Timetables do not obey normal indexing rules. Convert to table for
% transfer then convert back.
if istimetable(obj) && height(obj)>0
    revertTimetable = true;
    if nargin>2
        % Get the properties from the local object, then convert to table.
        props = value.Properties;
        value = timetable2table(value);
    else
        % Get the properties from the distributed array (may require
        % communication).
        props = subsref(obj, substruct('.', 'Properties'));
    end
    timeVar = props.DimensionNames{1};
    obj = timetable2table(obj);
    % Adjust range for the extra column
    rangeStruct.end(2) = rangeStruct.end(2) + 1;
else
    revertTimetable = false;
end

% Compute the size of data corresponding to the range.
szColumnFromRange = 1 + rangeStruct.end - rangeStruct.start;
szColumnFromRange( szColumnFromRange < 0 ) = 0;
% Remove trailing ones
szColumnFromRange = distributedutil.Sizes.removeTrailingOnes( szColumnFromRange );

if nargin == 2 % No value - retrieval
    x = iCreateLocalDataForRetrieval( obj, szColumnFromRange );
    doRetrieveData = true;
    if isempty( x )
        % No data to retrieve - leave now to save a remote call.
        return;
    end
else           % Assignment
    iArgChecksForAssignment( szColumnFromRange, value );
    doRetrieveData = false;
end

% Call into SPMD to get the cell array of values, and the descriptor
[xferInfo, xferComposite] = spmd_feval_fcn( @iBuildTransferComposite, {obj, rangeStruct, doRetrieveData} );
xferInfo = xferInfo.Value;
[~, csLimsIntoX, workers] = deal( xferInfo{:} );

% Loop over workers, and either retrieve stuff via xferInfo, or plug stuff from
% value into there.
for workerIdx = 1:max( workers )
    % Only retrieve the value from the worker if we must
    regionsForThisWorker = find( workers == workerIdx );
    if ~isempty( regionsForThisWorker )
        workerValue = xferComposite{workerIdx};
    end
    
    % Loop over this worker's regions
    for region=regionsForThisWorker.'
        idxInValueCell = workerValue.regionToValueCell( region );
        
        % "csLimsIntoX" are limits into the global array range of "obj" that we need
        % either for filling out a piece of "x" (retrieval); or for sending
        % across a piece of "value". In both cases, we need to offset the
        % limits by the start of the range that we're operating on.
        currentStart = 1 + csLimsIntoX.start( :, region ) - rangeStruct.start;
        currentEnd   = 1 + csLimsIntoX.end( :, region )   - rangeStruct.start;
        
        if doRetrieveData
            dataRetrieved = workerValue.values{idxInValueCell};
            % Do in-line indexing for 2D blocks to avoid colon vectors
            % being expanded. General case covers 3D and higher.
            switch numel(currentStart)
                case 2
                    if (currentStart(1)==1 && currentEnd(1)==size(x,1))
                        % All rows - use literal colon
                        x(:,currentStart(2):currentEnd(2)) = dataRetrieved;
                    elseif (currentStart(2)==1 && currentEnd(2)==size(x,2))
                        % All cols - use literal colon
                        x(currentStart(1):currentEnd(1),:) = dataRetrieved;
                    else
                        % Sub-block
                        x(currentStart(1):currentEnd(1), currentStart(2):currentEnd(2)) = dataRetrieved;
                        assert(~istimetable(x), "Shouldn't have sub-blocks for tables or timetables")
                    end
                otherwise
                    % N-D indexing
                    thisSubStruct = iSubStructForAssignRange( currentStart, currentEnd, size(x) );
                    x = subsasgn( x, thisSubStruct, dataRetrieved );
            end
            if istable(x) && (currentStart(2)~=1 || currentEnd(2)~=size(x,2))
                % In this case, the variables of table dataRetreived is a
                % subset of the variables of table x. We need to relay this
                % information to hEnsureComplexityMatches as its action is
                % per variable.
                cols = matlab.internal.ColonDescriptor(currentStart(2), currentEnd(2));
                x = distributedutil.hEnsureComplexityMatches( x, dataRetrieved, cols );
            else
                x = distributedutil.hEnsureComplexityMatches( x, dataRetrieved );
            end
            % Also fix up any timetable rowtimes
            if istable(x) || istimetable(x)
                rows = matlab.internal.ColonDescriptor(currentStart(1), currentEnd(1));
                cols = matlab.internal.ColonDescriptor(currentStart(2), currentEnd(2));
                x = parallel.internal.indexing.copyRowTimes(x, dataRetrieved, rows, cols);
            end
        else
            % Sending
            workerValue.values{idxInValueCell} = iSelectRange( value, currentStart, currentEnd );
        end
    end
    
    if ~isempty( regionsForThisWorker ) && ~doRetrieveData
        % Send the updated value back
        xferComposite{workerIdx} = workerValue;
    end
    workerValue = [];
end

% Now do some final tidying up
if ~doRetrieveData
    x = spmd_feval_fcn( @iApplyChanges, {obj, xferComposite} );
end

% Timetables are transferred as tables. Convert back now.
if revertTimetable
    x = table2timetable(x, 'RowTimes', timeVar);
end

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function iArgChecksForAssignment( szColumn, value )
% We're doing assignment so check that the span of the range corresponds to
% the size of the value we've been given. Note that "szColumn" has trailing
% ones stripped, so we can compare directly to the size of "value".

if ~all( szColumn == size( value ).' )
    error(message('parallel:distributed:BadValueForAssignPortion'));
end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function localData = iCreateLocalDataForRetrieval( obj, szColumn )
% We're doing retrieval, extract the template for construction of x.
tmpl      = distributedutil.Allocator.extractTemplate( obj );
localData = distributedutil.Allocator.create( szColumn.', tmpl );
localData = distributedutil.hEnsureComplexityMatches( localData, tmpl );
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function obj = iApplyChanges( obj, xferVariant )

lp   = getLocalPart( obj );
dtor = getCodistributor( obj );
for ii=1:length( xferVariant.values )
    thisSubStruct = iSubStructForAssignRange( xferVariant.lpLimsCell{ii}{:}, size(lp) );
    newValues = xferVariant.values{ii};
    lp = subsasgn( lp, thisSubStruct, newValues );
    lp = distributedutil.hEnsureComplexityMatches( lp, newValues );
    %  Take special care with timetables or tables containing timetables.
    %  The time vector isn't copied by subsasgn and must be transferred
    %  explicitly.
    if istable(lp) || istimetable(lp)
        lp = parallel.internal.indexing.copyRowTimes(lp, newValues);
    end
end
% TODO: Could use setLocalPart here.
obj = codistributed.build( lp, dtor, 'noCommunication' );
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [xfer, xferComp] = iBuildTransferComposite( obj, range, doRetrieveData )
descriptor = distributedutil.CodistributedDescriptor.buildDescriptor( obj );
[csLimsIntoX, lpLims, workers] = descriptor.mapGlobalRegion( range );
% Work out which regions I own and have values for
nRegions = length( workers );
myEls    = find( spmdIndex == workers );
lp       = getLocalPart( obj );

% Build up the variant which we'll return
xferComp.regionToValueCell = zeros( 1, nRegions );
xferComp.regionToValueCell( myEls ) = 1:length( myEls );
xferComp.lpLimsCell = cell( 1, length( myEls ) );
xferComp.values  = cell( 1, length( myEls ) );
for ii = 1:length( myEls )
    if doRetrieveData
        xferComp.values{ii} = iSelectRange( lp, lpLims.start( :, myEls(ii) ), ...
            lpLims.end( :, myEls(ii) ) );
    else
        % We're pushing the data, no need to add the values.
        xferComp.lpLimsCell{ii} = { lpLims.start( :, myEls(ii) ), ...
            lpLims.end( :, myEls(ii) ) };
    end
end

xfer = distributedutil.AutoTransfer( {descriptor, csLimsIntoX, workers} );
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function x = iSelectRange( data, r_start, r_end )
subsargs = cell( 1, length( r_start ) );
for ii=1:length( r_start )
    if r_start(ii)==1 && r_end(ii)==size(data,ii)
        subsargs{ii} = ':';
    else
        subsargs{ii} = r_start( ii ) : r_end( ii );
    end
end
ss = substruct( '()', subsargs );
x = subsref( data, ss );
x = distributedutil.hEnsureComplexityMatches( x, data, [], subsargs{2} );
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function ss = iSubStructForAssignRange( r_start, r_end, size_x )
% Form a substruct suitable for use in a subsasgn expression
N = numel( r_start );
subsargs = cell( 1, N );
% Make sure dims are padded with ones if required
if numel(size_x)<N
    size_x(numel(size_x)+1 : N) = 1;
end

for ii=1:length( r_start )
    if r_start(ii)==1 && r_end(ii)==size_x(ii)
        subsargs{ii} = ':';
    else
        subsargs{ii} = r_start( ii ) : r_end( ii );
    end
end
ss = substruct( '()', subsargs );
end