function X = avgpool(X, poolSize, nvArgs)
%AVGPOOL   Pool data to average value
%
%   Y = AVGPOOL(X,POOLSIZE) computes the average of the input data X over
%   pooling regions defined by POOLSIZE. Input X must be a formatted
%   dlarray object, with format containing at least one 'S' or 'T' dimension 
%   label and no more than three 'S' dimension labels. To pool over the 'T' 
%   dimension of X, you must also specify the 'PoolFormat' option. All other 
%   dimensions are treated as observation dimensions. POOLSIZE is a vector 
%   or scalar representing the pool size. If POOLSIZE is a scalar, then the 
%   pooling region has the same size for all pooling dimensions. If POOLSIZE 
%   is a vector, then the region size is specified for each pooling 
%   dimension of X. Output Y is a formatted dlarray with the same dimension 
%   labels as X.
%
%   Y = AVGPOOL(X,'global') computes the average of all the values in the
%   input data X. This is equivalent to computing AVGPOOL(X,POOLSIZE),
%   where POOLSIZE is set to the size of the 'S' dimensions of the input
%   data X.
%
%   Y = AVGPOOL(__,'DataFormat',FMT) computes the average pooling
%   of an unformatted dlarray X. FMT specifies the dimension labels of
%   the input data X. Output Y is an unformatted dlarray with the same
%   dimension order as X.
%
%   Y = AVGPOOL(X,POOLSIZE,'PARAM1',VAL1,'PARAM2',VAL2,...) specifies optional
%   parameter name/value pairs:
%
%       'Stride'        - Step size for traversing the input, specified as
%                         one of the following
%                           - a scalar, in which case the same value is
%                             used for all spatial dimensions.
%                           - a vector [s1,...,sd] where si is the stride
%                             along the ith dimension.
%                         The default value is 1. If 'Stride' is less than
%                         POOLSIZE in any dimension, then the pooling
%                         regions will overlap. The 'Stride' option is not
%                         supported when POOLSIZE is set to 'global'.
%       'Padding'       - Size of padding to apply to the edges of the 
%                         input, specified as one of the following:
%                           - 'same'. Padding is set so that the output
%                             size is the same as the input size when the
%                             stride is 1. More generally, the output size
%                             is ceil(inputSize/stride), where inputSize is
%                             the size of the input along a spatial
%                             dimension.
%                           - a scalar, in which case the same padding is
%                             applied for all spatial dimensions.
%                           - a vector [p1,...,pd] where pi is the padding
%                             applied to the start and the end along the
%                             ith spatial dimension.
%                           - a matrix [low1,...,lowd;high1,...,highd] of
%                             size 2-by-d where lowi and highi are the
%                             padding along the lower and upper edges of
%                             the ith spatial dimension.
%                         The default value of 'Padding' is 0; no padding 
%                         is applied. The 'Padding' option is not supported
%                         when POOLSIZE is set to 'global'.
%       'PaddingValue'  - Value used to pad the input along the
%                         edges, specified as one of the following values:
%                           - 0 - Input is padded with zeros along each 
%                             padding dimension. Zero values are included
%                             in the average calculation.
%                           - 'mean' - Input is padded with the mean value  
%                             along each padding dimension. Each edge 
%                             pooling region is padded with the mean of the 
%                             input data values contained by that region.
%                             This effectively excludes padded values from
%                             the average calculation. 
%                         The default value is 0.
%       'DataFormat'    - Dimension labels of the input data X, specified
%                         as a char array or string. Required if X is not a
%                         formatted dlarray.
%       'PoolFormat'    - Dimension labels of the input data to pool over, 
%                         specified as a character vector or string scalar 
%                         containing 'S' and 'T' labels only, corresponding 
%                         to the dimensions of the pool size. The default 
%                         value is a string with as many 'S' labels as are
%                         present in the input data X.
%
%   Example 1:
%       % Create formatted dlarray input data and compute avgpool with
%       % poolSize of 3 and padding of 1 at the start and end of each
%       % spatial dimension.
%       X = dlarray(randn(10,10,3,5),'SSCB');
%       Z = avgpool(X,3,'Padding',[1 1 1 1]);
%       % Output: Z = 10(S) x 10(S) x 3(C) x 5(B) dlarray double
%
%   Example 2:
%       % Create formatted dlarray input data and compute avgpool with
%       % poolSize set to 'global'.
%       X = dlarray(randn(10,10,3,5),'SSCB');
%       Z = avgpool(X,'global');
%       % Output: Z = 1(S) x 1(S) x 3(C) x 5(B) dlarray double
%
%   Example 3:
%       % Create unformatted dlarray input data and compute avgpool with
%       % poolSize of 3 and padding of 1 at the start and end of each
%       % spatial dimension.
%       X = dlarray(randn(5,3,10,10));
%       Z = avgpool(X,3,'Padding',[1 1 1 1],'DataFormat','BCSS');
%       % Output: Z = 5x3x10x10 dlarray double
%
%   Example 4:
%       % Create formatted dlarray input data and compute avgpool with
%       % poolSize of 3 and padding size of 1 at the start and end of each
%       % spatial dimension. Exclude the padded values from the average by 
%       % padding with the mean value of each region.
%       X = dlarray(randn(10,10,3,5),'SSCB');
%       Z = avgpool(X,3,'Padding',[1 1 1 1],'PaddingValue','mean');
%       % Output: Z = 10(S) x 10(S) x 3(C) x 5(B) dlarray double
%
%   Example 5:
%       % Create input data as a formatted dlarray with spatial,
%       % channel, batch, and time dimensions.
%       numChannels = 3;
%       height = 9; width = 9; sequenceLength = 9;
%       batchSize = 16;
%       X = dlarray(...
%           randn(height,width,numChannels,batchSize,sequenceLength),'SSCBT');
%       % Apply avgpool on the spatial and time dimensions, by using the
%       % PoolFormat argument, with windows of size 3x3x3 and stride 2.
%       poolSize = 3;
%       Y = avgpool(X,poolSize,'PoolFormat','SST','Stride',2);
%       % Output: Y = 4(S) x 4(S) x 3(C) x 16(B) x 4(T) dlarray double
%
%   See also DLARRAY, MAXPOOL, RELU, BATCHNORM

%   Copyright 2019-2022 The MathWorks, Inc.

arguments
    X
    poolSize {deep.internal.dlarray.mustBeNonDlarray}
    nvArgs.Stride {deep.internal.dlarray.mustBeNonDlarray, mustBeNumeric, ...
        mustBePositive, mustBeReal, mustBeInteger, mustBeNonempty}
    nvArgs.Padding {deep.internal.dlarray.mustBeNonDlarray}
    nvArgs.PaddingValue {deep.internal.dlarray.mustBeCharStringOrNumeric, ...
        mustBeNonempty, deep.internal.dlarray.mustBeNonDlarray} = 0
    nvArgs.DataFormat {deep.internal.dlarray.mustBeDataFormatType} = ''
    nvArgs.PoolFormat {deep.internal.dlarray.mustBeDataFormatType} = ''
end

% Only data defines execution environment, not parameters
poolSize = gather(poolSize);
nvArgs.PaddingValue = gather(nvArgs.PaddingValue);

% Handle input label
[X, perm] = deep.internal.dlarray.validateDataFormatArg(X, nvArgs.DataFormat);

% Validate PoolFormat argument. 
labelsX = dims(X);
[nvArgs.PoolFormat, poolSize, numSpatialDimsInPoolFmt, poolOverTime, ...
    numSpatialDims] =  deep.internal.dlarray.validatePoolFormatArg(nvArgs.PoolFormat, poolSize, labelsX);

% Get information about dimensions to be pooled over
[poolingDimsIdx, numPoolingDims] =  deep.internal.dlarray.getPoolingDimsBasedOnPoolFmt(labelsX, ...
    numSpatialDims, numSpatialDimsInPoolFmt, poolOverTime);

% The inputSize returned in output is always ordered
inputSize = size(X);

% Validate arguments
[nvArgs, poolSize, isGlobal] = deep.internal.dlarray.validatePoolingArguments(nvArgs, poolSize, numPoolingDims, inputSize(poolingDimsIdx));

if isGlobal
    X = mean(X, poolingDimsIdx);
else
    internalArgs = {poolSize, nvArgs.Padding, nvArgs.Stride, nvArgs.IncludePadding};
    X = deep.internal.dlarray.poolData(X, numPoolingDims, numSpatialDims, poolOverTime, @internal_avgpool, internalArgs);
end

% Remove dimensions from the dlarray if it was passed as unformatted.
if ~isempty(nvArgs.DataFormat)
    X = stripdims(X);
    X = ipermute(X, perm); % Permutes back in the same order as DataFormat
end
end
