function Y = movfunCommon(movfcn, X, k, varargin)
%MOVFUNCOMMON common implementation of distributed moving window functions

%   Copyright 2021-2022 The MathWorks, Inc.

% Make sure all arguments other than X are gathered.
distributedutil.CodistParser.verifyNonCodistributedInputs([{X}, {k}, varargin]);
k = gather(k);
hasSamplePoints = false;
if ~isempty(varargin)
    % Partial match for "SamplePoints" will also match "shrink" if it is
    % abbreviated to 's', so we seek args beginning with 's' that don't
    % follow "EndPoints".
    isSPinVarargin = cellfun(@(x) iIsPartialMatchFor("SamplePoints", x), varargin);
    isEPinVarargin = cellfun(@(x) iIsPartialMatchFor("EndPoints", x), varargin);
    isEPinVarargin = [false isEPinVarargin(1:end-1)]; % Move Name array to align with Value array.
    hasSamplePoints = any(isSPinVarargin & ~isEPinVarargin);
end
if hasSamplePoints
    samplePtsIndex = find(isSPinVarargin, 1);
    if samplePtsIndex == length(varargin)
        error(message("MATLAB:movfun:noSamplePointsValue"))
    end
    samplePoints = varargin(samplePtsIndex:samplePtsIndex+1);
    varargin(samplePtsIndex:samplePtsIndex+1) = [];
else
    samplePoints = {};
end
argList = distributedutil.CodistParser.gatherElements(varargin);
% Ensure we have a 1D distribution.
if iscodistributed(X)
    codistr = getCodistributor(X);
    if isa(codistr, 'codistributor2dbc')
        X = redistribute(X, codistributor1d());
        codistr = getCodistributor(X);
    end
end
if ~iscodistributed(X)
    if hasSamplePoints && iscodistributed(samplePoints{2})
        samplePoints = distributedutil.CodistParser.gatherElements(samplePoints);
    end
    % If X is not distributed, just call the standard function.
    Y = movfcn(X, k, argList{:}, samplePoints{:});
    return
end

% Use the local function to check syntax. (only for function with
% SamplePoints removed from the argument list).
if hasSamplePoints && isduration(k)
    ktest = 2;
else
    ktest = k;
end 
distributedutil.validateSyntax(movfcn, [{X}, {ktest}, argList]);

% Try to extract the dimension
try
    [dim,args] = iParseArgs(movfcn, X, argList);
catch err
    throwAsCaller(err);
end

if hasSamplePoints
    if istimetable(X)
        error(message("MATLAB:movfun:SamplePointsTimeTable"))
    end
    T = samplePoints{2};
    % Check for errors in SamplePoints vector because this wasn't checked.
    if isdatetime(T) && ~isduration(k) 
        error(message("MATLAB:movfun:winsizeNotDuration", "datetime"))
    end
    if ~(isnumeric(T) || isdatetime(T) || isduration(T))
        error(message("MATLAB:movfun:SamplePointsInvalidDatatype"))
    end
    if issparse(samplePoints{2})
        error(message("MATLAB:movfun:SamplePointsSparse"))
    end
    if ~(isvector(T) || isempty(T)) || length(T) ~= size(X, dim)
        error(message("MATLAB:movfun:SamplePointsLength",size(X, dim)))
    end
    if isnumeric(T) && ~isreal(T)
        error(message("MATLAB:movfun:SamplePointsComplex"))
    end
    if ~all(isfinite(T))
        error(message("MATLAB:movfun:SamplePointsNonFinite", "SamplePoints", "NaN"))
    end
    if isrow(T)  % Also make it a column vector.
        T = T';
    end
    if iscodistributed(T)
        codistrT = getCodistributor(T);
        if ~isa(codistrT, 'codistributor1d')
            T = redistribute(T, codistributor1d());
        end
    end
    if ~issorted(T)
        error(message("MATLAB:movfun:SamplePointsSorted", "SamplePoints"))
    end
    if any(diff(T) == 0)
        error(message("MATLAB:movfun:SamplePointsDuplicate", "SamplePoints"))
    end
    samplePoints(2) = {T};
end

if codistr.Dimension == dim
    % Apply windows in the distribution dimension.
    Y = iApplyWindowsInDistDim(movfcn, X, k, dim, args, samplePoints);
else
    % Apply windows independently in each worker.
    Y = iApplyIndependentWindows(movfcn, X, k, dim, args, samplePoints);
end

end


function Y = iApplyWindowsInDistDim(movfcn, X, k, dim, args, samplePoints)
% Apply windows in the distribution dimension.

% We implement an efficient communication scheme to transfer just the
% window elements at the borders between workers. K tells us how many we
% need from the previous worker or the next (either a two-element vector
% [numBefore numAfter] or a scalar where numBefore = floor(k/2) and
% numAfter = k-numBefore-1).

if isscalar(k)
    szPad(1) = floor(k/2);        % numBefore
    if isnumeric(k)
        szPad(2) = ceil(k - szPad(1) - 1);  % numAfter
    else
        szPad(2) = ceil(k - szPad(1));  % duration after
    end
else
    szPad = [floor(k(1)) ceil(k(2))];
end
if ~isempty(samplePoints)
    T = samplePoints{2};  % SamplePoints vector.
    [localX, szLeftPad, localT] = padArraysForSamples(X, szPad, T);
    samplePoints(2) = {localT};
else
    [localX, szLeftPad] = padLocalX(X, szPad);
end
codistr = getCodistributor(X);

% Process our local part.
lpOut = movfcn(localX, k, args{:}, samplePoints{:});

% Remove padding if it hasn't been discarded
if ~any(cellfun(@(x) iIsPartialMatchFor("Discard", x), args))
    szInDistDim = codistr.Partition(spmdIndex);
    lpOut = extractSlices(lpOut, (1:szInDistDim) + szLeftPad, dim);   
end

% Build a new codistributed from the local parts. These are guaranteed to
% have the same size as the input including in the windowed direction. The
% partitioning is therefore the same too.
szOut = size(X);
if dim>numel(szOut)
    % Pad with ones
    szOut = [szOut, ones(1, dim-numel(szOut))];
end
szLpOut = size(lpOut, dim);
newPartition = spmdCat(szLpOut);
szOut(dim) = sum(newPartition);
codistOut = codistributor1d(dim, newPartition, szOut);
Y = codistributed.pDoBuildFromLocalPart(lpOut, codistOut);
end

function Y = iApplyIndependentWindows(movfcn, X, k, dim, args, samplePoints)
% Apply windows independently in each worker.
codistr = getCodistributor(X);
assert(codistr.Dimension ~= dim, ...
    "Can't apply windows independently in the distribution dimension");

if ~isempty(samplePoints) && iscodistributed(samplePoints{2})
    samplePoints(2) = {gather(samplePoints{2})};
end

% Process our local part.
lpOut = movfcn(getLocalPart(X), k, args{:}, samplePoints{:});

% Build a new codistributed from the local parts. These are guaranteed to
% have the same size as the input except in the windowed direction. The
% partitioning is therefore the same too.
szOut = size(X);
if dim>numel(szOut)
    % Pad with ones
    szOut = [szOut, ones(1, dim-numel(szOut))];
end
szOut(dim) = size(lpOut, dim);
codistOut = codistributor1d(codistr.Dimension, codistr.Partition, szOut);
Y = codistributed.pDoBuildFromLocalPart(lpOut, codistOut);
end


function [dim,args] = iParseArgs(fcn, X, args)
% Look at the optional arguments to extract the dimension. The dimension
% comes immediately after the window (first trailing argument) unless there
% is a weight, in which case it comes second.
needsWeight = ismember(func2str(fcn), ["movstd","movvar"]);
if needsWeight
    dimArg = 2;
else
    dimArg = 1;
end

if numel(args)>=dimArg && isnumeric(args{dimArg}) && isscalar(args{dimArg})
    % Get the user-specified dimension
    dim = args{dimArg};
else
    dim = distributedutil.Sizes.firstNonSingletonDimension(size(X));
    % We now need to insert the dimension into the argument list to ensure
    % local parts don't choose to operate in different dimensions.
    if needsWeight
        % Check if a weight was provided, otherwise insert []
        if ~isempty(args) && isnumeric(args{1})
            % We have a weight
            args = [args(1), {dim}, args(2:end)];
        else
            % Insert default weight
            args = [{[]}, {dim}, args];
        end
    else
        args = [{dim},args];
    end
end
end

function tf = iIsPartialMatchFor(str, arg)
tf = matlab.internal.datatypes.isScalarText(arg) ...
    && startsWith(str, arg, "IgnoreCase", true);
end
