function meanX = mean(varargin)
%MEAN Average or mean value
%   S = MEAN(X)
%   S = MEAN(X,DIM)
%   S = MEAN(...,TYPE)
%   S = MEAN(...,MISSING)
%
%   See also MEAN, TALL.

%   Copyright 2015-2018 The MathWorks, Inc.

narginchk(1,4);
tall.checkNotTall(upper(mfilename), 1, varargin{2:end});

% Use the in-memory version to check the arguments
tall.validateSyntax(@mean,varargin,'DefaultType','double');

[args, flags] = splitArgsAndFlags(varargin{:});

% We cannot support datetime/mean since we don't have the building blocks, but
% we allow it through the validateType so that we can throw a more specific
% error.
assert(istall(args{1})); % we have checked that args 2:end are not tall.
x = args{1};
x = tall.validateType(x, mfilename, ...
                      {'numeric', 'logical', 'duration', 'datetime', 'char'}, ...
                      1);
adaptor = args{1}.Adaptor;
[nanFlagCell, precisionFlagCell] = adaptor.interpretReductionFlags(upper(mfilename), flags);

assert(~isempty(nanFlagCell)); % We need the nan flag at the client to switch implementations.
nanFlag = nanFlagCell{1};

% If no dimension specified, try to deduce it (will be empty if we can't).
if numel(args) == 1
    dim = matlab.bigdata.internal.util.deduceReductionDimension(adaptor);
else
    dim = args{2};
    % DIM is allowed to be a column or row vector. We want a row vector.
    if iscolumn(dim)
        dim = dim';
    end
end

if strcmp(adaptor.Class, "datetime")
    meanX = iDatetimeMean(x, dim, nanFlag);
else
    % General version
    meanX = iGeneralMean(x, dim, nanFlag, precisionFlagCell);
end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function meanX = iGeneralMean(x, dim, nanFlag, precisionFlagCell)
if isempty(dim)
    % Default dimension

    % Compute adaptors for results from reduceInDefaultDim
    sumXAdaptor = matlab.bigdata.internal.adaptors.getAdaptorForType(...
        computeSumResultType(x.Adaptor.Class, precisionFlagCell));
    szInRedDimAdaptor = matlab.bigdata.internal.adaptors.getScalarDoubleAdaptor();

    if ismember(nanFlag, ["includenan", "includenat"])
        % Mean including all elements
        [result, resolvedDimension] = ...
            reduceInDefaultDim(@(x, dim) sum(x, dim, 'includenan', precisionFlagCell{:}), x);
        result = tall(result, sumXAdaptor);
        % Note: we don't care about the type of resolvedDimension ...
        resolvedDimension = tall(resolvedDimension);
        szTv = size(x);
        
        % ... but we do care about the type of sizeInReductionDimension as we need it
        % for the type of 'result' to propagate through the RDIVIDE call.
        sizeInReductionDimension = clientfun(@iGetSizeinDim, szTv, resolvedDimension);
        sizeInReductionDimension.Adaptor = szInRedDimAdaptor;
        meanX = result ./ sizeInReductionDimension;
    else
        % Mean excluding NaNs
        result = reduceInDefaultDim(@(x, dim) sum(x, dim, 'omitnan', precisionFlagCell{:}), x);
        result = tall(result, sumXAdaptor);
        meanX = result ./ sum(~isnan(x));
    end
else
    % Specified dimension - note that here we rely on SUM and RDIVIDE to propagate
    % adaptors correctly where necessary to deal with duration.
    if ismember(nanFlag, ["includenan", "includenat"])
        % Mean including all elements
        meanX = sum(x, dim, 'includenan', precisionFlagCell{:}) ./ elementsPerResult(x, dim);
    else
        % Mean excluding NaNs
        meanX = sum(x, dim, 'omitnan', precisionFlagCell{:}) ./ sum(~isnan(x), dim);
    end
end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function meanX = iDatetimeMean(x, dim, nanFlag)
% Special version of MEAN for datetimes since SUM and RDIVIDE are not
% defined for datetimes. Instead we use MEAN on each block and combine
% using ratios.

% NB: Coding pattern copied from reduceInDefaultDim, but able to deal with
% reduction that returns two outputs instead of 1. We should try and fold
% back into reduceInDefaultDim.
if isempty(dim)
    % Default dimension. We calculate both reduction and slice-wise
    % versions in a single reduction then choose between them.
    [meanXReduced, ~, meanXSliced] = aggregatefun( ...
            @(y) iPerBlockDatetimeMeanUnknownDim(y, nanFlag), ...
            @(m,n,p) iCombineDatetimeMeansUnknownDim(m, n, p, nanFlag), ...
            x);
     meanXReduced.Adaptor = resetSizeInformation(x.Adaptor);
     meanXSliced.Adaptor = resetSizeInformation(x.Adaptor);
     
    meanX = clientfun(@iPickWhichOne, size(x), meanXReduced, meanXSliced);
    % In this case we can't tell the output size in advance. All we know is
    % that the result has tall size 1 either way.
    meanX.Adaptor = resetSizeInformation(x.Adaptor);
    meanX.Adaptor = setSizeInDim(meanX.Adaptor, 1, 1);
else
    % Dim is known
    if isnumeric(dim) && ~ismember(1,dim)
        % Not reducing tall dimension. Just call slicewise
        meanX = slicefun(@(y) mean(y, dim, nanFlag), x); 
    else
        % Tall reduction
        [meanX, ~] = aggregatefun( ...
            @(y) iPerBlockDatetimeMean(y, dim, nanFlag), ...
            @(m,n) iCombineDatetimeMeans(m, n, nanFlag), ...
            x);
    end
    meanX.Adaptor = x.Adaptor;
    meanX.Adaptor = computeReducedSize(meanX.Adaptor, x.Adaptor, dim, false);
end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [meanX, countX] = iPerBlockDatetimeMean(x, dim, nanFlag)
% If we have more than one row, we need to interpolate based on the number
% of controbuting elements.
meanX = mean(x, dim, nanFlag);
if ismember(nanFlag, ["omitnan", "omitnat"])
    countX = sum(~ismissing(x), dim);
else
    countX = elementsPerResult(x, dim);
    % Expand the count to apply to all elements
    countX = repmat(countX, size(meanX));
end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [meanX, countX] = iCombineDatetimeMeans(meanX, countX, nanFlag)
% If we have more than one row, we need to interpolate based on the number
% of controbuting elements.
import matlab.bigdata.internal.util.indexSlices;

if size(meanX,1)<=1
    return;
end

if all(countX(1,:)==countX(2:end,:), "all")
    % Special case where all blocks had the same number of elements. We
    % can just take the mean
    meanX = mean(meanX, 1, nanFlag);
    countX = sum(countX, 1);
else
    % Blocks had different number of elements. For each row in turn, work
    % out the ratio for each output element then interpolate the results.
    assert(isequal(size(meanX), size(countX)), "Expected means and counts to be same size")
    mean1 = indexSlices(meanX, 1);
    count1 = indexSlices(countX, 1);
    
    for slice = 2:size(meanX,1)
        mean2 = indexSlices(meanX, slice);
        count2 = indexSlices(countX, slice);
        ratio = count2 ./ (count2 + count1);
        
        % METHOD 1: Add scaled (duration) difference
        % newMean = mean1 + ratio.*(mean2 - mean1);
        
        % METHOD 2: Interlace the rows and use INTERP1 (which *is* supported
        % for datetime) to calculate the intermediate value between each
        % pair of elements. This is more complex than method 1, but turns
        % out to be significantly more accurate than going via duration.
        interlaced = reshape([mean1;mean2], 1, [])';
        interlacedRatio = reshape(ratio, 1, []);
        query = interlacedRatio + ((1:numel(mean1))*2 - 1);
        newMean = interp1(1:numel(interlaced), interlaced, query);
        newMean = reshape(newMean, size(mean1));
        
        % To avoid NaN propagation, fix up values that didn't need interpolating
        newMean(interlacedRatio==0) = mean1(interlacedRatio==0);
        newMean(interlacedRatio==1) = mean2(interlacedRatio==1);
        
        % Finally, replace the original slice with the new mean
        mean1 = newMean;
        count1 = count1 + count2;
    end
    
    % The accumulated result is in mean1, count1
    meanX = mean1;
    countX = count1;
end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [meanXTall, countXTall, meanXSliced] = iPerBlockDatetimeMeanUnknownDim(x, nanFlag)
% If we have more than one row, we need to interpolate based on the number
% of contributing elements.

% Compute mean in tall dimension
[meanXTall, countXTall] = iPerBlockDatetimeMean(x, 1, nanFlag);

% Also compute mean as though tall dim is 1
if size(x, 1)==0
    % If empty, the tall version will have expanded the array for us into
    % all NaT
    x = meanXTall;
elseif size(x, 1)>1
    x = matlab.bigdata.internal.util.indexSlices(x, 1);
end
meanXSliced = mean(x, nanFlag);
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [meanXTall, countXTall, meanXSliced] = iCombineDatetimeMeansUnknownDim(meanXTall, countXTall, meanXSliced, nanFlag)
% If we have more than one row, we need to interpolate based on the number
% of controbuting elements.

[meanXTall, countXTall] = iCombineDatetimeMeans(meanXTall, countXTall, nanFlag);

% For the slice-wise version, always keep just one
if size(meanXSliced, 1)>1
    meanXSliced = matlab.bigdata.internal.util.indexSlices(meanXSliced, 1);
end

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [result, dim] = iPickWhichOne(sz, reducedResult, slicedResult)
% A helper to use in clientfun that can establish which of the results is
% correct, and return that along with the resolved reduction dimension.

nonSingletonDims = find(sz ~= 1);
if isequal(sz, [0 0])
    % Special case (as per g1361570) for [] empty input, MATLAB doesn't
    % apply the first-non-singleton dimension rule. We get the correct
    % result in slicedResult, so we set dim = 3 to force that.
    dim = 3;
elseif isempty(nonSingletonDims)
    % No non-singleton dimensions - scalar case
    dim = 1;
else
    dim = nonSingletonDims(1);
end

% Pick the result depending on the computed full size of the tall variable.
if dim == 1
    result = reducedResult;
else
    result = slicedResult;
end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function sz = iGetSizeinDim(szVec, dim)
% Another special case as per g1364892. mean([]) returns NaN, which does not
% match any return that you get by specifying a dimension. Here we're working
% around yet another special case which is that sum([]) needs to return 0, and
% it does that by pretending that the reduction dimension is 3.
if isequal(szVec, [0 0])
    sz = 0;
else
    sz = szVec(dim);
end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Helper to work out how many input elements contributed to a given result
% element. Note that the output will be a tall scalar.
function n = elementsPerResult(data, dim)
if isnumeric(dim)
    % We need to take the product of the reduced dimension lengths, but
    % since we don't know what ndims is yet, we need to call SIZE
    % repeatedly in case we have a trailing dimension. These calls to size
    % will be folded into a single reduction.
    n = size(data, dim(1));
    for ii=2:numel(dim)
        n = n .* size(data, dim(ii));
    end
else
    % all
    n = numel(data);
end
end
