function [tf, step] = isuniform(x)
%ISUNIFORM Check if codistributed array data is uniformly spaced.
%   TF = ISUNIFORM(A)
%   [TF,STEP] = ISUNIFORM(A)
%   
%   See also ISUNIFORM, CODISTRIBUTED.


%   Copyright 2023 The MathWorks, Inc.

% ISUNIFORM requires the input to be a real numeric-ish vector.
if ~isreal(x) || ~(isnumeric(x) || islogical(x))
    error(message('MATLAB:isuniform:MustBeReal'));
end

if isinteger(x) || islogical(x)
    step = NaN;
else
    step = full(NaN('like',getLocalPart(x)));
end

if isscalar(x) || isempty(x)
    tf = false;
    return
end

if ~isvector(x)
    error(message('MATLAB:isuniform:MustBeVector'));
end

% Force to a column vector in 1D distribution.
x = colonize(x);

% Calculate the expected step size using the end-points
firstVal = gather(subsref(x, substruct('()', {1})));
firstValF = iMustBeFullRealFloat(firstVal);
lastVal = gather(subsref(x, substruct('()', {numel(x)})));
lastValF = iMustBeFullRealFloat(lastVal);
numSteps = numel(x)-1;
meanStep = (lastValF - firstValF) ./ numSteps;
% Check if we overflowed
if ~isfinite(meanStep) && isfinite(firstValF) && isfinite(lastValF)
    meanStep = lastValF./numSteps - firstValF./numSteps; % Divide first to avoid overflow
end
tol = 4*eps(max(abs(firstValF), abs(lastValF))); % Based on value used for in-memory

% If only two elements, will be uniform unless  not finite
if numel(x)==2
    tf = gather(allfinite(x));
    if nargout>1 && tf
        iThrowIfStepTooBig(x, meanStep, firstVal, lastVal);
        step = meanStep;
    end
    return
end

% Step must be finite
if ~isfinite(meanStep)
    tf = false;
    return;
end

% We can efficiently calculate the result with minimal communication using
% a binary tree reduction. For each partition call "isuniform" and store
% the result in a struct along with the first value, last value, step size,
% and number of values.
S = iMakeStructForLocalPart(getLocalPart(x), meanStep, tol);

% Combine the structs from each partition to get the overall result.
S = spmdReduce(@(a,b) iCombineStructs(a,b,meanStep,tol), S);

% In common with other IS* functions the result is on the client.
tf = gather(S.CouldBeUniform);
if nargout>1 && tf
    iThrowIfStepTooBig(x, meanStep, firstVal, lastVal);
    step = meanStep;
end

end

function s = iMakeStructForLocalPart(lp, meanStep, tol)
% Build an information struct describing one local part.
if isempty(lp)
    s = struct(...
        "Numel", 0,...
        "CouldBeUniform", true,...
        "FirstVal", NaN, "LastVal", NaN ...
        );
elseif isscalar(lp)
    % For a scalar partition we just need to store the value, although if
    % not finite then it cannot be uniform anyway.
    lp = iMustBeFullRealFloat(lp);
    s = struct(...
        "Numel", 1,...
        "CouldBeUniform", isfinite(lp), ...
        "FirstVal", lp, "LastVal", lp ...
        );
else
    % We would like to use ISUNIFORM on the local part here for all
    % datatypes but we must use the expected step and tolerance from the
    % entire range. Therefore only use if integer (no tolerance).
    if isinteger(lp)
        tf = isuniform(lp);
    else
        d = diff(lp);
        tf = all(abs(d-meanStep) < tol);
    end
    s = struct(...
        "Numel", numel(lp),...
        "CouldBeUniform", tf,...
        "FirstVal", iMustBeFullRealFloat(lp(1)), ...
        "LastVal", iMustBeFullRealFloat(lp(end)) ...
        );
end
end

function val = iMustBeFullRealFloat(val)
% Helper to ensure an input is full, real, and floating-point
if ~isfloat(val)
    val = double(val);
end
val = full(real(val));
end

function S = iCombineStructs(S1, S2, meanStep, tol)
% Combine a pair of structs from neighbouring partitions.

% If either input was empty, return the other.
if S2.Numel == 0
    S = S1;
    return
elseif S1.Numel == 0
    S = S2;
    return
end

% Make a default result.
S = struct(...
    "Numel", S1.Numel + S2.Numel,...
    "CouldBeUniform", S1.CouldBeUniform && S2.CouldBeUniform,...
    "FirstVal", S1.FirstVal, "LastVal", S2.LastVal ...
    );

% If either is definitively non-uniform then the result is non-uniform.
if ~S.CouldBeUniform
    return;
end

% We now just need to check that the partitions are separated by the step.
S.CouldBeUniform = abs(S1.LastVal + meanStep - S2.FirstVal) <= tol;
end

function iThrowIfStepTooBig(x, step, firstVal, lastVal)
% Check that the step size can be represented in double without loss of
% precision.

% We need to take care for 64-bit integers - it is possible to have an
% integer step that cannot be accurately stored.
if isUnderlyingType(x,"int64") || isUnderlyingType(x,"uint64")
    % For unsigned it is important to avoid negative step!
    if firstVal > lastVal
        intStep = (firstVal - lastVal) ./ (numel(x)-1);
    else
        intStep = (lastVal - firstVal) ./ (numel(x)-1);
    end
    if intStep > flintmax(underlyingType(step))
        error(message("MATLAB:isuniform:StepTooLarge"));
    end
elseif isfloat(x)
    % For floating-point types just check it fits
    if step > realmax(underlyingType(step))
        error(message("MATLAB:isuniform:StepTooLarge"));
    end
end
end