function [varargout] = getGridVectors(isMeshGrid, isNdGrid, varargin)
%GETGRIDVECTORS finds the grid vectors that represent the inputs
%   [X1,X2,X3,...Xn, V] = GETGRIDVECTORS(ISMESHGRID, ISNDGRID, X1,X2,X3,...Xn, V)
%   for a grid defined by X1,X2,X3,...Xn and corresponding grid point values V.
%   It assumes that checkGrid has been called to valid the grid.  It returns
%   the grid vectors as columns.  It works with either ndgrid or meshgrid
%   inputs, returning replicated ndgrid vectors except for the one that
%   corresponds to the distribution dimension of V, which is
%   correspondingly distributed.

%   Copyright 2021-2023 The MathWorks, Inc.

narginchk(5, inf);
assert(nargout == nargin - 2, 'Number of outputs and inputs do not match.')
n1 = nargin-3;
X = varargin(1:n1);
V = varargin{end};

if isMeshGrid
    X1 = X{1};
    X{1} = iGetPencil(X{2}, 1);  % ndgrid X vector (column)
    X{2} = iGetPencil(X1, 2).';  % ndgrid Y vector (row)
elseif isNdGrid
    X{1} = iGetPencil(X{1}, 1);  % ndgrid X vector (column)
    X{2} = iGetPencil(X{2}, 2).';  % ndgrid Y vector (row)
end
if n1 > 2
    for j = 3:n1
        X{j} = iGetPencil(X{j}, j);
        X{j} = shiftdim(X{j});  % Make it a column vector
    end
end

isCodistX = cellfun(@iscodistributed, X);
if iscodistributed(V)   % Grid vectors should be codistributed the same as V
    codV = getCodistributor(V);
    distDimV = codV.Dimension;
    codX = codistributor1d(1, codV.Partition, [size(V, distDimV), ones(1, ndims(V))]);
    for i = 1:n1
        if i == distDimV
            if iscodistributed(X{i})
                X{i} = redistribute(X{i}, codX);
            else
                X{i} = codistributed(X{i}, codX);
            end
        else
            if iscodistributed(X{i})
                X{i} = gather(X{i});
            end
        end
    end
else    % At most one grid vector should be codistributed.
    if any(isCodistX)
        ixCodistX = find(isCodistX);
        if numel(ixCodistX) > 1
            [X{ixCodistX(2:end)}] = gather(X{ixCodistX(2:end)});
        end
    end
end
varargout = [X, {V}];
end

function gridVector = iGetPencil(X, dim)
% iGetPencil gets a vector from the outside edge of a multidimensional
% array
%  Input    X - A multidimensional array, part of a meshgrid or ndgrid
%           dim The dimension of the array to retrieve the pencil
%
%  Output   gridVector - the pencil (X(1, 1, ..., dim ,... 1, 1)

subsArg = repmat({1}, 1, ndims(X));
subsArg(dim) = {':'};
subs = substruct('()', subsArg);
gridVector = subsref(X, subs);
end
