function B = subsrefTabularBraces(A, aidx, varNames, preserveComplexity)
% subsrefTabularBraces performs brace indexing for distributed tables/timetables.
%
% B = A{X1,X2}
%
% X1 can be: ':', numeric, or logical; local or codistributed
% X2 can be: ':', numeric, logical, cellstr, string, or vartype; local or codistributed
%
% varNames is an optional list of the variable names if they are
% already known (to avoid calculating them twice).
%
% preserveComplexity==true will try and preserve complexity of the output

% Copyright 2018-2022 The MathWorks, Inc.

if numel(aidx(1).subs) == 1
    error(message('MATLAB:table:LinearSubscript'));
elseif numel(aidx(1).subs) ~= 2
    error(message('MATLAB:table:NDSubscript'));
end
if nargin<3
    varNames = getTableVariableNames(A);
end
if nargin<4
    preserveComplexity = false;
end

aidx.subs = tabularInterpretVartype(A, aidx.subs);
[firstSub, secondSub] = deal(aidx.subs{:});
[~, secondSubNumeric] = distributedutil.lookupTableVars(secondSub, varNames);
% Table throws slightly different errors for bad indices.
try
    parallel.internal.indexing.validateTableIndexBounds(A, aidx);
catch ME
    throw(ME);
end

% Use smooth parens to select the region of interest.
psub = substruct('()', {firstSub secondSubNumeric});
tableOut = subsrefTabularParens(A, psub);
outCodistr = getCodistributor(tableOut);
distDim = outCodistr.Dimension;

% Get hold of the output template so we know the resulting type
outTmpl = table2array(getTableTemplate(tableOut));
preserveComplexity = preserveComplexity && isnumeric(outTmpl); % Only need to preserve complex for numeric output

% Convert from table to contiguous array, being careful about types and
% sizes.
outLP = iTypeCorrectedTable2Array(tableOut, outTmpl, distDim, preserveComplexity);

% Note that whilst the table itself is always 2D with width equal to the
% number of variables, the data inside can be ND, resulting in ND output.
% We must build a new codistributor. The sizes of each local part must
% match except in the distribution dimension.
if distDim==2
    % Work out the output size based on all columns in the template (i.e.
    % from all workers), but set the rows to size of the local part.
    globalSize = size(outTmpl);
    globalSize(1) = size(outLP,1);
    % Get the new partitioning by asking all workers their size
    part = spmdCat(size(outLP, distDim));

else
    % When distributed in first dimension, we already know the partition
    % and just need to know the new trailing dimensions.
    part = outCodistr.Partition;
    globalSize = size(outLP);
    globalSize(distDim) = sum(part);
end

outCodistr = codistributor1d(distDim, part, globalSize);
B = codistributed.pDoBuildFromLocalPart(outLP , outCodistr);
end


function outLP = iTypeCorrectedTable2Array(tableOut, outTmpl, distDim, preserveComplexity)
outLP = getLocalPart(tableOut);
if distDim==2
    % If workers have different columns we can end up with incorrect type
    % conversions. The outTmpl gives the final type so we use this to
    % enforce all local parts to the right type.

    % Resize the template to match the output height. Make sure to add
    % trailing dimensions for N-D arrays.
    sizeEmptyVar = [height(tableOut), 0, size(outTmpl, 3:ndims(outTmpl))];
    emptyVar = distributedutil.Allocator.create(sizeEmptyVar, outTmpl(:,[]));
    % Add the new var at the beginning to ensure it sets the type. ADDVARS
    % doesn't allow adding a char row vector so add char as double then
    % convert in this case. Also, if the result is a timetable we must
    % ensure all local parts have the right time vector, even if empty.
    if ischar(emptyVar) && isrow(emptyVar)
        outLP = addvars(outLP, double(emptyVar), 'before', 1);
        outLP.(1) = char(outLP.(1));
    elseif istimetable(emptyVar)
        % We know the result is going to be a timetable but since
        % timetables can be concatenated with tables and cells this worker
        % might not have the time vector. Get it now.
        emptyVar.Properties.RowTimes = iGetGlobalRowTimes(outLP);
        outLP = addvars(outLP, emptyVar, 'before', 1);
    else
        outLP = addvars(outLP, emptyVar, 'before', 1);
    end
end
% Finally just use table2array(t) or t.Variables do the conversion
if preserveComplexity
    % Array2table drops complexity so if the result is numeric we use
    % table.Variables instead.
    name = outLP.Properties.DimensionNames{2};
    outLP = outLP.(name);
else
    outLP = table2array(outLP);
end
end

function globalRowTimes = iGetGlobalRowTimes(LP)
% Get a single row-times vector to apply to all local part timetables.
rowTimes = [];
for ii=1:width(LP)
    if istimetable(LP.(ii))
        rowTimes = LP.(ii).Properties.RowTimes;
        break
    end
end
% Coordinate across workers (the time vectors will all be the same height
% or empty).
globalRowTimes = spmdCat(rowTimes,2);
if size(globalRowTimes,2)>1
    % More than one timetable found. We will get an error later if they are
    % not consistent so just keep the first.
    globalRowTimes = globalRowTimes(:,1);
end
end