function F = interpLocalGriddedInterpolant(X, V, method, extrap, noX, szX, dimInterp)
%INTERPLOCALGRIDDEDINTERPOLANT Produce a GriddedInterpolant suitable for 
%evaluating queries on this worker.  Where necessary, coordinates and
%values are exchanged with other workers to extend local arrays enough to
%produce accurate interpolated values.  This function is used whenever
%either the grid or the values are distributed.
%   Arguments:
%   X         Cell array of grid vectors, one cell for each interpolation
%             dimension.  At most one can be codistributed.
%   V         Array of values, can be codistributed.
%   method    Interpolation method.
%   extrap    Extrapolation method.
%   noX       True if X was not specified in the call to interp.
%   szX       The size that X, Y, Z, ... must have.
%   dimInterp The dimension(s) along which interpolation will be done.
%
%   Output:
%   F         A GriddedInterpolant unless method == 'spline' in which case
%             F is a piecewise cubic Hermite polynomial.  F is empty if no
%             interpolation will be done on this worker.

%   Copyright 2021-2022 The MathWorks, Inc.

interpDimensionality = numel(dimInterp);
switch interpDimensionality
    case 1
        interpFunName = 'interp1';
    case 2
        interpFunName = 'interp2';
    case 3
        interpFunName = 'interp3';
    otherwise
        interpFunName = 'interpn';
end
distDimX = find(cellfun(@iscodistributed, X), 1); % X is a cell array with at
                                      % most one distributed column vector
if iscodistributed(V)
    codistributed.pVerifyUsing1d(interpFunName, V);
    codV = getCodistributor(V);
    distDimV = codV.Dimension;
    newCodX = codistributor1d(1, codV.Partition);
    if ~any(distDimV == dimInterp)
        % V is not distributed in an interpolation dimension.
        F = iDistVNotInterpDim(X, V, method, extrap, noX, szX);
        return
        
    elseif nnz(codV.Partition) == 1
        % V is distributed in an interpolation dimension, but all of V is
        % on one worker.
        F = iAllVOnOne(X, V, method, extrap, noX, szX);
        return
        
    else % V is distributed in an interpolation dimension with some values
         % on different workers.  Distribute X the same way.
        if noX
            X = arrayfun(@(x) (1:x)', szX, 'UniformOutput', false);
        else % find out how X is distributed
            if ~isempty(distDimX) && distDimX ~= distDimV
                X(distDimX) = {gather(X{distDimX})};
            elseif distDimX == distDimV
                X(distDimV) = {redistribute(X{distDimV}, newCodX)};
            end
        end
        if ~iscodistributed(X{distDimV})
            X(distDimV) = {codistributed(X{distDimV}, newCodX)};
        end
        distDimX = distDimV;
    end

else  % X is codistributed because this function wouldn't be called if
      % neither were codistributed.
    codX = getCodistributor(X{distDimX});
    if nnz(codX.Partition) == 1
        % All of X is on one worker and V is not distributed. That worker
        % can do all the work.
        if codX.Partition(spmdIndex)>0
            localX = X;
            localX(distDimX) = {getLocalPart(X{distDimX})};
            columnX = cellfun(@iscolumn, localX);
            localX(columnX) = cellfun(@transpose, localX(columnX), 'UniformOutput', false);
            F = griddedInterpolant(localX, V, method, extrap);
        else
            F = [];
        end
        return
    end
    
    % Distribute V to match the dimension of X that is distributed
    codV = codistributor1d(distDimX, codX.Partition, size(V));
    V = codistributed(V, codV);
end

% By this point both X and V must be codistributed, and will have matching
% distribution schemes.
assert(iscodistributed(X{distDimX}) && iscodistributed(V), ...
    "Should only get here if both X and V are distributed.");

if any(strcmpi(method, {'cubic','makima'})) && interpDimensionality > 1
    % Cubic needs three points to distinguish between non-uniform and not
    % enough points warnings.
    padSz = 3;
elseif any(strcmpi(method, {'pchip','cubic','v5cubic','makima','spline'}))
    % Cubic-based methods may need two points on either side.
    padSz = 2;
else
    % Linear or nearest-neighbor methods need only one point on either side.
    padSz = 1;
end

[localX, localV] = padLocalGrid(X, V, padSz);
isInterpOnDistDim = any(distDimX == dimInterp);
if strcmpi(method, 'spline') && interpDimensionality == 1 && isInterpOnDistDim
    F = interpSpline(X{distDimX}, localX, V, localV, distDimX);
elseif numel(localX) > 1
    X(distDimX) = {localX};
    localX = X;
    F = griddedInterpolant(localX, localV, method, extrap);
else
    F = [];
end
end

function F = iAllVOnOne(X, V, method, extrap, noX, szX)
%iAllVonOne returns a griddedInterpolant for the case that V is
%codistributed on an interpolation dimension, but all of it is on one worker.

codV = getCodistributor(V);
distDimV = codV.Dimension;
newCodX = codistributor1d(1, codV.Partition);
thisWorkerHasV = codV.Partition(spmdIndex);
distDimX = find(cellfun(@iscodistributed, X));
if ~isempty(distDimX)  % X is codistributed
    if distDimX == distDimV
        xOfDistDim = redistribute(X{distDimV}, newCodX);
    else % Not codistributed on same dimension as V
        X(distDimX) = {gather(X{distDimX})};
        xOfDistDim = codistributed(X{distDimV}, newCodX);
    end
    localX = X;
    localX(distDimV) = {getLocalPart(xOfDistDim)};
elseif thisWorkerHasV
    if noX
        localX = arrayfun(@(x) 1:x, szX, 'UniformOutput', false);
    else
        localX = X;
    end
end

if thisWorkerHasV
    localV = getLocalPart(V);
    F = griddedInterpolant(localX, localV, method, extrap);
else
    F = [];
end

end

function F = iDistVNotInterpDim(X, V, method, extrap, noX, szX)
%iDistVNotInterpDim returns a griddedInterpolant for the case where V is
%codistributed, but not an an interpolation dimension.

if noX
    X = arrayfun(@(x) (1:x)', szX, 'UniformOutput', false);
else 
    X = distributedutil.CodistParser.gatherElements(X);
end
localV = getLocalPart(V);
if isempty(localV)
    % The worker has no values, so return an empty F
    F = [];
else
    F = griddedInterpolant(X, localV, method, extrap);
end

end