function [p, avg_gsq] = rmspropupdate(p, g, avg_gsq, lr, beta2, epsilon)
%RMSPROPUPDATE Update parameters via root mean square propagation
%
%   RMSPROPUPDATE(NET,GRAD,AVG_SQG) changed to use a default squared
%   gradient decay factor of 0.9 in R2020a. To use the previous default
%   decay factor use RMSPROPUPDATE(NET,GRAD,AVG_SQG,0.001,0.999).
%   
%   [NET,AVG_SQG] = RMSPROPUPDATE(NET,GRAD,AVG_SQG) updates the learnable
%   parameters of the dlnetwork NET using the RMSProp gradient descent
%   algorithm, with a default global learning rate of 0.001 and a squared
%   gradient decay factor of 0.9. Use RMSPROPUPDATE to iteratively update
%   learnable parameters during network training.
%
%   Input GRAD contains the gradients of the loss with respect to each of
%   the network parameters. Input AVG_SQG contains the moving average of
%   the element-wise squares of the parameter gradients from the previous
%   call to RMSPROPUPDATE. GRAD and AVG_SQG must be tables with the same
%   structure as NET.Learnables, with a Value variable containing a cell
%   array of parameter gradients or average squared gradients. Input GRAD
%   can be obtained using the dlgradient and dlfeval functions. The global
%   learning rate is multiplied by the corresponding learning rate factor
%   for each parameter in each layer of the dlnetwork.
%
%   If input AVG_SQG is empty, the function assumes no previous gradients
%   and executes like the first update in a series of iterations. 
%
%   Outputs NET and AVG_SQG are the updated dlnetwork and average squared
%   gradients, respectively.
%
%   [PARAMS,AVG_SQG] = RMSPROPUPDATE(PARAMS,GRAD,AVG_SQG) updates the deep
%   learning parameters in PARAMS using the RMSProp gradient descent
%   algorithm, with a default learning rate of 0.001 and a squared gradient
%   decay factor of 0.9. Input PARAMS can be a dlarray, a numeric array, a
%   cell array, a structure, or a table with a Value variable containing
%   the learnable parameters of the network. GRAD and AVG_SQG must have the
%   same datatype and ordering as PARAMS. Input GRAD can be obtained using
%   the dlgradient and dlfeval functions. All parameter values are updated
%   using the global learning rate.
%
%   Outputs PARAMS and AVG_SQG are the updated parameters and average
%   squared gradients, respectively.
%
%   [___] = RMSPROPUPDATE(___,LEARNRATE,SQG_DECAY) also specifies values to
%   use for the global learning rate and squared gradient decay factor.
%   LEARNRATE must be a positive scalar. SQG_DECAY must be a scalar between
%   0 and 1.
%
%   [___] = RMSPROPUPDATE(___,LEARNRATE,SQG_DECAY,EPSILON) specifies a
%   small constant used to prevent division by zero in the update equation.
%   The default value of EPSILON is 1e-8.
%
%   Example 1:
%      % Perform an RMSProp update step with learning rate of 0.01 and
%      % squared gradient decay factor of 0.95.
%      p = rand(3,3,4); 
%      avg_sqg = ones(3,3,4); 
%      g = ones(3,3,4); 
%      [p,avg_sqg] = rmspropupdate(p,g,avg_sqg,0.01,0.95);
%
%   Example 2:
%      % Perform a single epoch of training on a network that classifies 
%      % handwritten digits.
%
%      % Load training data and encode.
%      [XTrain,YTrain] = digitTrain4DArrayData;
%      XTrain = single(XTrain);
%      YEncoded = onehotencode(YTrain,2)';
%      
%      % Define the network and specify the average image using the 'Mean' option 
%      % in the image input layer.
%      layers = [
%          imageInputLayer([28 28 1],'Name','input','Mean',mean(XTrain,4))
%          convolution2dLayer(5,20,'Name','conv1')
%          reluLayer('Name','relu1')
%          convolution2dLayer(3,20,'Padding',1,'Name','conv2')
%          reluLayer('Name','relu2')
%          fullyConnectedLayer(size(YEncoded,1),'Name','fc')];
%      net = dlnetwork(layerGraph(layers));
%      
%      % Initialize average parameter gradients.
%      squaredGradientsAvg = [];
%      
%      % Loop over one epoch of mini-batches.
%      miniBatchSize = 128;
%      for i = 1:miniBatchSize:numel(YTrain)
%          % Read mini-batch of data and convert the labels to dummy variables.
%          idx = i:min(i+miniBatchSize-1,numel(YTrain));
%          
%          % Convert mini-batch of data to dlarray.
%          X = dlarray(XTrain(:,:,:,idx),'SSCB');
%          
%          % Evaluate the model gradients using dlfeval and the
%          % modelGradients function defined below.
%          gradients = dlfeval(@modelGradients,net,X,YEncoded(:,idx));
%          
%          % Update the network parameters using the RMSProp algorithm.
%          [net.Learnables,squaredGradientsAvg] = rmspropupdate(net.Learnables,gradients,squaredGradientsAvg);
%      end
%
%      function gradients = modelGradients(net,X,Y)
%          loss = crossentropy(softmax(forward(net,X)),Y);
%          gradients = dlgradient(loss,net.Learnables);
%      end
%
%   See also ADAMUPDATE, DLFEVAL, DLGRADIENT, DLNETWORK, DLUPDATE, SGDMUPDATE

%   Copyright 2019-2020 The MathWorks, Inc.

if nargin<6
    epsilon = 1e-8;
else
    validateattributes(epsilon, "numeric", {'scalar','real','finite','positive'}, ...
        'rmspropupdate', 'EPSILON', 6);
end

if nargin<5
    beta2 = 0.9;
else
    validateattributes(beta2, "numeric", {'scalar','real','finite','>=',0,'<',1}, ...
        'rmspropupdate', 'GSQ_DECAY', 5);
end

if nargin<4
    lr = 0.001;
else
    validateattributes(lr, "numeric", {'scalar', 'real', 'finite', 'nonnegative'}, ...
        'rmspropupdate', 'LR', 4);
end

if isempty(avg_gsq)
    % Execute a first-step update with g_sq_av set to 0. The step will create
    % an array for vel that is the correct size
    func = deep.internal.LearnableUpdateFunction( ...
        @(p, g) iSingleStepValue(p, g, 0, lr, beta2, epsilon), ...
        @(p, g) iSingleStepParameter(p, g, 0, lr, beta2, epsilon) );
    
    [p, avg_gsq] = deep.internal.networkContainerFun(func, p, g);
else
    % Execute the normal update
    func = deep.internal.LearnableUpdateFunction( ...
        @(p, g, avg_gsq) iSingleStepValue(p, g, avg_gsq, lr, beta2, epsilon), ...
        @(p, g, avg_gsq) iSingleStepParameter(p, g, avg_gsq, lr, beta2, epsilon) );
    
    [p, avg_gsq] = deep.internal.networkContainerFun(func, p, g, avg_gsq);
end

end

function [p, avg_gsq] = iSingleStepParameter(p, g, avg_gsq, lr, beta2, epsilon)
% Apply per-parameter learn-rate factor
lr = lr .* p.LearnRateFactor;

[step, avg_gsq] = nnet.internal.cnn.solver.rmspropstep(g, avg_gsq, lr, beta2, epsilon);
p.Value = p.Value + step;
end

function [p, avg_gsq] = iSingleStepValue(p, g, avg_gsq, lr, beta2, epsilon)
[step, avg_gsq] = nnet.internal.cnn.solver.rmspropstep(g, avg_gsq, lr, beta2, epsilon);
p = p + step;
end
