classdef(Abstract) TrainingOptionsMiniBatch < nnet.cnn.TrainingOptions
    % TrainingOptionsMiniBatch   Base training options for mini-batch solvers
    %
    %   This class holds the training options common to mini-batch solvers
    %
    %   MiniBatchTrainingOptions properties:
    %       LearnRateSchedule            - Method for lowering the global 
    %                                      learning rate during training
    %       LearnRateDropFactor          - A multiplicative factor applied 
    %                                      to the learning rate every time
    %                                      a certain number of epochs passed 
    %                                      if 'LearnRateSchedule' is set to 
    %                                      'piecewise'
    %       LearnRateDropPeriod          - The learning rate drop factor is 
    %                                      applied to the global learning 
    %                                      rate every time this number of 
    %                                      epochs is passed if  
    %                                      'LearnRateSchedule' is set to 
    %                                      'piecewise'
    %       Shuffle                      - This controls if the training
    %                                      data is shuffled
    %       WorkerLoad                   - Specify compute and prefetch
    %                                      workers and their relative load
    %                                      in a parallel pool
    %       CheckpointFrequencyUnit      - Unit of CheckpointFrequency
    %       DispatchInBackground         - Flag to use prefetch queueing
    %                                      when reading data
    %       PreprocessingEnvironment     - Preprocessing environment for
    %                                      training

    % Copyright 2023 The MathWorks, Inc.

    properties
        % MaxEpochs   Maximum number of epochs
        %   The maximum number of epochs that will be used for training.
        %   Training will stop once this number of epochs has passed.
        MaxEpochs {iValidatePropertyValue(MaxEpochs, 'MaxEpochs')} = iGetDefaultValue('MaxEpochs')

        % LearnRateSchedule 
        LearnRateSchedule {iValidatePropertyValue(LearnRateSchedule,'LearnRateSchedule')} = iGetDefaultValue('LearnRateSchedule')

        % LearnRateDropFactor
        LearnRateDropFactor {iValidatePropertyValue(LearnRateDropFactor,'LearnRateDropFactor')} = iGetDefaultValue('LearnRateDropFactor')

        % LearnRateDropPeriod        
        LearnRateDropPeriod {iValidatePropertyValue(LearnRateDropPeriod,'LearnRateDropPeriod')} = iGetDefaultValue('LearnRateDropPeriod')        

        % MiniBatchSize   The size of a mini-batch for training
        %   The size of the mini-batch used for each training iteration.
        MiniBatchSize {iValidatePropertyValue(MiniBatchSize,'MiniBatchSize')} = iGetDefaultValue('MiniBatchSize')

        % Shuffle   This controls when the training data is shuffled. It
        % can either be 'once' to shuffle data once before training,
        % 'every-epoch' to shuffle before every training epoch, or 'never'
        % in order not to shuffle the data.
        Shuffle {iValidatePropertyValue(Shuffle,'Shuffle')} = iGetDefaultValue('Shuffle')

        % CheckpointFrequency   Frequency for saving checkpoint networks
        CheckpointFrequency
 
        % CheckpointFrequencyUnit   Frequency unit for saving checkpoint
        % networks
        CheckpointFrequencyUnit {iValidatePropertyValue(CheckpointFrequencyUnit,'CheckpointFrequencyUnit')} = iGetDefaultValue('CheckpointFrequencyUnit')

        % SequenceLength   Determines the strategy used to create
        % mini-batches of sequence data
        SequenceLength

        % PreprocessingEnvironment  Accelerate reading, augmentation, and
        % queuing of data.
        PreprocessingEnvironment {iValidatePropertyValue(PreprocessingEnvironment,'PreprocessingEnvironment')} = iGetDefaultValue('PreprocessingEnvironment')
    end

    properties(Hidden)
        % WorkerLoad   Relative division of load between parallel workers
        % on different hardware
        WorkerLoad {iValidatePropertyValue(WorkerLoad,'WorkerLoad')} = iGetDefaultValue('WorkerLoad')
    end

    properties(Abstract)
        InitialLearnRate
    end

    properties(Hidden, Dependent)
        % DispatchInBackground  Alias for PreprocessingEnvironment
        DispatchInBackground
    end

    properties(Hidden=true,Dependent,SetAccess=private)
        % LearnRateScheduleSettings   Settings for the learning rate schedule
        %   The learning rate schedule settings. This summarizes the
        %   options for the chosen learning rate schedule. The field Method
        %   gives the name of the method for adjusting the learning rate.
        %   This can either be 'none', in which case the learning rate is
        %   not altered, or 'piecewise', in which case there will be two
        %   additional fields. These fields are DropFactor, which is a
        %   multiplicative factor for dropping the learning rate, and
        %   DropPeriod, which determines how many epochs should pass
        %   before dropping the learning rate.
        LearnRateScheduleSettings
    end

    methods (Access=protected)
        function fn = getPublicProperties(this)
            % Retrieve all the visible properties
            fn = getPublicProperties@nnet.cnn.TrainingOptions(this);
            % Also add all the hidden properties
            fn{end+1} = 'WorkerLoad';
        end
    end

    methods(Access = public)
        function this = TrainingOptionsMiniBatch(varargin)
            this = this@nnet.cnn.TrainingOptions(varargin{:});
            this = manageDispatchInBackground(this, varargin);
        end

        function out = saveobj(this)
            out = saveobj@nnet.cnn.TrainingOptions(this);
            out.MaxEpochs = this.MaxEpochs;
            out.LearnRateSchedule     = this.LearnRateSchedule;
            out.LearnRateDropFactor   = this.LearnRateDropFactor;
            out.LearnRateDropPeriod   = this.LearnRateDropPeriod;
            out.MiniBatchSize = this.MiniBatchSize;
            out.Shuffle = this.Shuffle;
            out.WorkerLoad = this.WorkerLoad;
            out.InitialLearnRate = this.InitialLearnRate;
            out.CheckpointFrequencyUnit = this.CheckpointFrequencyUnit;
            out.DispatchInBackground = this.DispatchInBackground;
            out.PreprocessingEnvironment = this.PreprocessingEnvironment;
        end
    end

    % set methods for properties requiring nontrivial canonicalization
    methods
        function obj = set.WorkerLoad(obj,newValue)
            obj.WorkerLoad = iCanonicalizePropertyValue(newValue,'WorkerLoad');
        end

        function obj = set.CheckpointFrequency(obj,newValue)
            iValidatePropertyValue(newValue,'CheckpointFrequency');
            obj.CheckpointFrequency = iCanonicalizePropertyValue(newValue,'CheckpointFrequency');
        end

        function obj = set.LearnRateDropFactor(obj,newValue)
            obj.LearnRateDropFactor = iCanonicalizePropertyValue(newValue,'LearnRateDropFactor');
        end

        function obj = set.LearnRateDropPeriod(obj,newValue)
            obj.LearnRateDropPeriod = iCanonicalizePropertyValue(newValue,'LearnRateDropPeriod');
        end

        function obj = set.MaxEpochs(obj,newValue)
            obj.MaxEpochs = iCanonicalizePropertyValue(newValue,'MaxEpochs');
        end

        function obj = set.MiniBatchSize(obj,newValue)
            obj.MiniBatchSize = iCanonicalizePropertyValue(newValue,'MiniBatchSize');
        end

        function obj = set.Shuffle(obj,newValue)
            obj.Shuffle = iCanonicalizePropertyValue(newValue,'Shuffle');
        end

        function obj = set.LearnRateSchedule(obj,newValue)
            obj.LearnRateSchedule = iCanonicalizePropertyValue(newValue,'LearnRateSchedule');
        end

        function obj = set.SequenceLength(obj,newValue)
            iValidatePropertyValue(newValue,'SequenceLength');
            obj.SequenceLength = iCanonicalizePropertyValue(newValue,'SequenceLength');
        end

        function obj = set.CheckpointFrequencyUnit(obj,newValue)
            obj.CheckpointFrequencyUnit = iCanonicalizePropertyValue(newValue,'CheckpointFrequencyUnit');
        end

        % this set method is only used when loading from releases 2016a ||
        % 2016b, which would throw an error if they don't find it
        function obj = set.LearnRateScheduleSettings(obj,newValue)
            % the Method field is the only one which has different name
            % for NVP and how it is stored in LRS struct
            obj.LearnRateSchedule = newValue.Method;
            fn = fieldnames(newValue);
            for ii=1:numel(fn)
               if ~strcmp(fn{ii},'Method')
                   obj.(fn{ii}) = newValue.(fn{ii});
               end
            end
        end

        % get method for Dependent variable LearnRateScheduleSettings
        function LRScheduleSettings = get.LearnRateScheduleSettings(this)
            LRScheduleSettings = this.createLearnRateScheduleSettings( ...
                this.LearnRateSchedule, ...
                this.LearnRateDropFactor, ...
                this.LearnRateDropPeriod);
        end

        function obj = set.DispatchInBackground(obj, dispatchInBackground)
            % Setting DispatchInBackground effectively sets
            % PreprocessingEnvironment
            nnet.internal.cnn.options.OptionsValidator.assertValidDispatchInBackground(dispatchInBackground);
            if dispatchInBackground
                obj.PreprocessingEnvironment = 'parallel';
            else
                obj.PreprocessingEnvironment = 'serial';
            end
        end

        function dispatchInBackground = get.DispatchInBackground(obj)
            dispatchInBackground = obj.PreprocessingEnvironment ~= "serial";
        end
    end

    methods(Hidden)
        function validateOptionsForTrainNetwork(this)
            validateOptionsForTrainNetwork@nnet.cnn.TrainingOptions(this)
            if this.PreprocessingEnvironment == "background"
                % Prevent backgroundPool usage with trainNetwork
                error(message('nnet_cnn:trainNetwork:UnsupportedPreprocessingEnvironment', 'background'));
            end
        end
    end

    methods(Access=private)
        function scheduleSettings = createLearnRateScheduleSettings(~, ...
                learnRateSchedule, learnRateDropFactor, learnRateDropPeriod)

            % method used to build the LearnRateScheduleSettings struct

            scheduleSettings = struct;
            switch learnRateSchedule
                case 'none'
                    scheduleSettings.Method = 'none';
                case 'piecewise'
                    scheduleSettings.Method = 'piecewise';
                    scheduleSettings.DropRateFactor = learnRateDropFactor;
                    scheduleSettings.DropPeriod = learnRateDropPeriod;
                otherwise
                    error(message('nnet_cnn:TrainingOptionsSGDM:InvalidLearningRateScheduleMethod'));
            end
        end

        function this = manageDispatchInBackground(this, args)
            % Manage DispatchInBackground which is an alias of
            % PreprocessingEnvironment and a dependent property

            % Inspect varargin to respect the order of the options being
            % specified
            specifiedNvps = args(2:2:end);
            isSaveOptionStruct = isscalar(specifiedNvps) && isstruct(specifiedNvps{1});

            isDispatchInBackgroundOption = iIsOptionSpecified(specifiedNvps,"DispatchInBackground", isSaveOptionStruct);

            if any(isDispatchInBackgroundOption)
                % Find if PreprocessingEnvironment was also specified
                if any(iIsOptionSpecified(specifiedNvps,"PreprocessingEnvironment", isSaveOptionStruct))
                    % PreprocessingEnvironment specified, we use
                    % that value which was already set so nothing to do
                    return
                end
                % Find last position
                lastDispatch = find(isDispatchInBackgroundOption, true, 'last');
                % We need to set the DispatchInBackground value set by the
                % user. This also sets PreprocessingEnvironment as it is
                % dependent
                if isSaveOptionStruct
                    this.DispatchInBackground = args{2}.DispatchInBackground;
                else
                    this.DispatchInBackground = args{1 + lastDispatch*2};
                end
            end
        end
    end

    methods(Access=protected)
        function warnIfCheckpointFrequencyDoesNotMatchCheckpointPath(this)
            if isempty(this.CheckpointPath) ...
                    && ~(this.CheckpointFrequency == 1 && strcmp(this.CheckpointFrequencyUnit, 'epoch'))
                % Disable backtrace so that warning does not produce stack
                % trace if backtrace is on, then restore backtrace to its
                % original state when exiting the method
                backtraceState = warning('query', 'backtrace');
                restoreBacktraceState = onCleanup(@()warning(backtraceState));
                warning off backtrace
                warning(message('nnet_cnn:TrainingOptionsSGDM:CheckpointFreqAndPathMismatch'));
            end
        end
    end
end

function tf = iIsOptionSpecified(args, optionName, argsAsStruct)
% args is either the specified NVPs, or a struct from a saved object
if argsAsStruct
    % Input is a struct
    tf = isfield(args{1},optionName);
else
    % Input is a cell array of chars
    tf = false(size(args));
    for i = 1:numel(args)
        tf(i) = startsWith(optionName, args{i});
    end
end
end

function val = iGetDefaultValue(optname)
persistent optionsDefinition;
if isempty(optionsDefinition)
    optionsDefinition=nnet.internal.cnn.options.TrainingOptionsDefinitionMiniBatch;
end
val = optionsDefinition.getDefaultValue(optname);
end

function iValidatePropertyValue(x,optname)
persistent optionsDefinition;
if isempty(optionsDefinition)
    optionsDefinition=nnet.internal.cnn.options.TrainingOptionsDefinitionMiniBatch;
end
optionsDefinition.validateValue(x,optname);
end

function canonicalizedval = iCanonicalizePropertyValue(x,optname)
persistent optionsDefinition;
if isempty(optionsDefinition)
    optionsDefinition=nnet.internal.cnn.options.TrainingOptionsDefinitionMiniBatch;
end
canonicalizedval = optionsDefinition.canonicalizeValue(x,optname);
end