function A = subsasgnTabularDot(A, S, B)
% Dot assignment for distributed table/timetable.

% Copyright 2018-2022 The MathWorks, Inc.

Atemplate = getTableTemplate(A);
Aprops = Atemplate.Properties;
AvarNames = Aprops.VariableNames;

% Check if we're assigning properties
subscript = S(1).subs;
if ischar(subscript) && strcmp(subscript,"Properties")
    A = iSubsasgnTabularProperties(A, Atemplate, S, B);
    return;
end

% Check if we're assigning a variable
if ischar(S(1).subs)
    S(1).subs = string(S(1).subs);
end
[varIdx,varName] = iResolveDotSubscript(Atemplate, S(1).subs);

if varIdx>numel(AvarNames)
    % Grow a new variable
    A = addvars(A, B, 'NewVariableNames', varName);
else
    % Replacing/updating a variable in the table.
    A = iSubsasgnVariable(A, S, B, varIdx, varName);
end
end

function A = iSubsasgnVariable(A, S, B, varIdx, varName)
% Update one table variable
Acodistr = getCodistributor(A);

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

% Now assign the result into the table
if Acodistr.Dimension==1
    % Variable is split across workers
    globalSz = size(B);
    % If B doesn't have the right number of rows, throw the right error.
    if globalSz(1) ~= size(A,1)
        error(message("MATLAB:table:RowDimensionMismatch"));
    else
        B = redistribute(B, codistributor1d(1,Acodistr.Partition,globalSz));
    end
    A.Local.(S.subs) = B.Local;
else
    % Variable is just on one worker. Work out which one. Put in a
    % try-catch-sync since only one worker will error if there's a problem.
    subs = substruct('.', varName);
    workerWithVar = find(cumsum(Acodistr.Partition)>=varIdx, 1, "first");
    if iscodistributed(B)
        % Work out which worker needs B and gather it there
        B = gather(B, workerWithVar);
    end
    
    try
        if spmdIndex == workerWithVar
            A.Local.(subs.subs) = B;
        end
        E = [];
    catch E
    end
    % Check if the assignment errored
    allE = spmdCat(E);
    if ~isempty(allE)
        throw(allE(1));
    end
end
% Data has changed. Clear metadata.
A.Metadata.TableTemplate.reset();
end

function A = iSubsasgnTabularProperties(A, Atemplate, S, B)
% Assign the special table properties field. We might be assigning the
% whole properties structure or just one field.
if numel(S)>1
    % If assigning the special RowNames/RowTimes then make sure the size is
    % right.
    Acodistr = getCodistributor(A);
    if numel(S)==2 && strcmp(S(2).type, ".") ...
            && any(strcmp(S(2).subs, ["RowNames","RowTimes"]))
        if Acodistr.Dimension==2
            % All local parts have the time vector. Just set it
            A.Local = subsasgn(A.Local, S, gather(B));
        else
            % Each worker has part of the time vector.
            if isempty(B)
                % Clear rownames on all workers.
                A.Local = subsasgn(A.Local, S, gather(B));
            elseif (size(B,1) ~= size(A,1))
                error(message("MATLAB:table:InvalidRowNames"));
            else
                % Distributed the rownames across workers
                Bcodistr = codistributor1d(1,Acodistr.Partition,size(B));
                if iscodistributed(B)
                    B = redistribute(B, Bcodistr);
                else
                    B = codistributed(B, Bcodistr);
                end
                A.Local = subsasgn(A.Local, S, B.Local);
            end
        end
    else
        % Use the template to check the assignment is valid
        Atemplate = subsasgn(Atemplate, S, B);
        % Some properties are not supported on distributed
        unsupportedProps = distributedutil.getDistributedTimetableUnsupportedProperties();
        if numel(S)==2 && strcmp(S(2).type, ".") ...
            && any(strcmp(S(2).subs, unsupportedProps))
            error(message('parallel:distributed:UnsupportedTimetableProperty', S(2).subs));
        end
    end
    % If that worked, put the properties back
    A = applyTableTemplate(A, Atemplate);
    
else
    % If assigning the whole properties block we must take care over the
    % RowNames/Times, which may also be distributed.
    if istimetable(B)
        rowField = "RowTimes";
    else
        rowField = "RowNames";
    end
    hasRows = isfield(B,rowField) && ~isempty(B.(rowField));
    if hasRows
        rowVals = B.(rowField);
        B.(rowField) = [];
    end
    % Use the template to check the assignment is valid
    Atemplate = subsasgn(Atemplate, S, B);
    % Some properties are not supported on distributed
    unsupportedProps = distributedutil.getDistributedTimetableUnsupportedProperties();
    if any(isfield(B, unsupportedProps))
        % Get all the unsupported properties.
        propName = unsupportedProps(isfield(B, unsupportedProps));
        % Print only the first unsupported property in the error
        % message.
        error(message('parallel:distributed:UnsupportedTimetableProperty', propName(1)));
    end

    % If that worked, put the properties back
    A = applyTableTemplate(A, Atemplate);

    % Now deal with RowNames/Times
    if hasRows
        % If split across workers, ensure row names have same distribution
        % as the overall table.
        codistr = getCodistributor(A);
        if codistr.Dimension==1
            rowVals = redistribute(rowVals, codistributor1d(1, codistr.Partition));
            A.Local.Properties.(rowField) = rowVals.Local;
        else
            % Every worker needs a copy.
            rowVals = gather(rowVals);
            A.Local.Properties.(rowField) = rowVals;
        end
    end
end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Resolve a variable subscript into an index
function [varIdx,varName] = iResolveDotSubscript(template, subscript)
% For cases tt.Foo and equivalently tt.('Foo'), we allow the subscript to be
% only: scalar string, char-vector, and numeric integer scalar. We need to
% deal with growing cases where the index doesn't resolve to an existing
% variable.

% Assign into the template to get all input checking done.
template2 = subsasgn(template, substruct('.',subscript), zeros(0,1));

% Now resolve the index into the new template (we already know the
% subscript is valid).
varNames = template2.Properties.VariableNames;
[~,varIdx] = distributedutil.lookupTableVars(subscript, varNames);
assert(isscalar(varIdx), ...
    'Unexpectedly resolved dot-subscript to multiple variables.');
varName = varNames{varIdx};
end
