classdef LBFGSState < matlab.mixin.CustomDisplay
    % LBFGSState      L-BFGS solver state
    %
    %   To create an L-BFGS solver state object, use lbfgsState
    %
    %   LBFGSState properties:
    %       HistorySize                   - Maximum history size for the L-BFGS
    %                                       update
    %       InverseHessianFactor          - Scalar value that is used to
    %                                       construct the inverse Hessian
    %                                       approximation
    %       StepHistory                   - History of steps taken by L-BFGS
    %                                       solver
    %       GradientsDifferenceHistory    - History of gradient differences
    %       HistoryIndices                - Order of the indices in StepHistory
    %                                       and GradientsDifferenceHistory
    %       Loss                          - Loss at the last iteration
    %       Gradients                     - Gradients of the loss with
    %                                       respect to the parameters at
    %                                       the last iteration
    %       AdditionalLossFunctionOutputs - Additional outputs of the loss
    %                                       function at the last iteration
    %       StepNorm                      - L2 norm of last step
    %       GradientsNorm                 - L-infinity norm of last
    %                                       gradients
    %       NumLineSearchIterations       - Number of line search iterations
    %                                       performed at the last iteration
    %       LineSearchStatus              - String indicating the status of
    %                                       line search at the last iteration
    %       LineSearchMethod              - Line search method used in last
    %                                       iteration
    %       MaxNumLineSearchIterations    - Maximum number of line search
    %                                       iterations
    %       InitialInverseHessianFactor   - Value of InverseHessianFactor
    %                                       used in the first iteration
    %       InitialGradientsNorm          - Value of GradientsNorm used in
    %                                       the first iteration
    %
    %   See also lbfgsState, lbfgsupdate.
    
    %   Copyright 2022-2023 The MathWorks, Inc.

    properties(Dependent)
        % InverseHessianFactor  Scalar used to construct the inverse
        % Hessian approximation
        % A real positive scalar specifying the initial inverse Hessian
        % approximation as the value of InverseHessianFactor times
        % the identity matrix. A value of 1 is fine in most cases.
        InverseHessianFactor

        % StepHistory   History of steps taken by L-BFGS solver
        % A row cell array with each element specifying the history of
        % the step vectors taken.
        StepHistory

        % GradientsDifferenceHistory    History of gradient differences
        % A row cell array with each element specifying the history of
        % the gradient differences.
        GradientsDifferenceHistory

        % HistoryIndices    Order of the indices in StepHistory and
        % GradientsDifferenceHistory
        % A row vector with elements which are permuted values of 1:m
        % such that HistoryIndices(1) is the most recent correction
        % pair and HistoryIndices(end) is the oldest correction pair.
        HistoryIndices
    end

    properties(Dependent, SetAccess=private)
        % Loss  Loss at the last iteration
        % A real scalar specifying the loss at the last iteration.
        Loss

        % Gradients Gradients of the loss with respect to the parameters at
        % the last iteration.
        % A numeric array, dlarray, cell array, struct, or table specifying
        % the gradients of the loss with respect to the parameters at the
        % last iteration.
        Gradients

        % AdditionalLossFunctionOutputs Additional outputs of the loss
        % function at the last iteration
        % A cell array containing any additional outputs of the loss
        % function at the last iteration.
        AdditionalLossFunctionOutputs

        % StepNorm  L2 norm of last step
        % A real nonnegative scalar specifying the L2 norm of the step
        % taken at the last iteration.
        StepNorm

        % GradientsNorm  L-infinity norm of last gradients
        % A real nonnegative scalar specifying the L-infinity norm of the 
        % gradient at the last iteration.
        GradientsNorm

        % NumLineSearchIterations   Number of line search iterations
        % performed at the last iteration
        % An integer specifying the number of line search iterations used
        % in the last iteration.
        NumLineSearchIterations

        % LineSearchStatus  Status of line search at the last iteration
        % A string indicating the status of the line search procedure at
        % the last iteration.
        LineSearchStatus
    end

    properties(SetAccess=private)
        % LineSearchMethod  A string specifying the line search method used
        % at the last iteration
        % Possible values are 'weak-wolfe', 'strong-wolfe', and
        % 'backtracking'.
        LineSearchMethod = ""

        % MaxNumLineSearchIterations    Maximum number of line search
        % iterations
        % An integer specifying the maximum number of allowable line search
        % iterations that could have been used in the last iteration.
        MaxNumLineSearchIterations = 0

        % HistorySize   Maximum history size for the L-BFGS update
        % A nonnegative integer specifying the maximum allowable history
        % size that can be used during lbfgsupdate.
        HistorySize

        % InitialInverseHessianFactor   Scalar used to construct the inverse
        % Hessian approximation at the first iteration
        % A real positive scalar specifying the value of
        % InverseHessianFactor at the first iteration
        InitialInverseHessianFactor

        % InitialGradientsNorm   L-infinity norm of the gradients before
        % any updates were performed
        % A real nonnegative scalar specifying the L-infinity norm of the 
        % gradients before any updates were performed
        InitialGradientsNorm
    end

    properties(Access=private)
        % An mlearnlib.internal.utils.solver.LBFGS solver object.
        PrivateSolver

        % An mlearnlib.internal.utils.solver.CurrentIteration object.
        CurrentIteration

        % An mlearnlib.internal.utils.solver.LBFGSOperationsStrategy
        % object.
        LBFGSOperationsStrategy

        % An mlearnlib.internal.utils.solver.LineSearchStrategyFactory
        % object.
        LineSearchStrategyFactory

        % A boolean indicating whether the solver state has been
        % initialized
        Initialized

        % A boolean indicating whether the solver state variables need to
        % be validated.
        NeedsExtraValidation
    end

    methods
        function this = LBFGSState(historySize, initialInverseHessianFactor)
            this.HistorySize = historySize;
            this.InitialInverseHessianFactor = initialInverseHessianFactor;

            % Create an internal representation of the LBFGS solver.
            this.PrivateSolver = mlearnlib.internal.utils.solver.LBFGS(...
                historySize, ...
                initialInverseHessianFactor);

            this.CurrentIteration = mlearnlib.internal.utils.solver.CurrentIteration;
            this.LineSearchStrategyFactory = mlearnlib.internal.utils.solver.LineSearchStrategyFactory;
            
            this.NeedsExtraValidation = false;
            this.Initialized = false;
        end

        function this = set.StepHistory(this, stepHistory)
            arguments
                this
                stepHistory(1,:) cell
            end
            stepHistory = iConvertEmptyCellToCanonicalForm(stepHistory);

            % Verify that the number of elements is less than or equal to
            % HistorySize.
            this.assertNumelLessThanOrEqualTo(stepHistory, 'StepHistory')

            this.PrivateSolver.StepHistory = stepHistory;

            % Ensure that the solver state performs validation at the next
            % call to lbfgsupdate.
            this.NeedsExtraValidation = true;
        end

        function stepHistory = get.StepHistory(this)
            stepHistory = this.PrivateSolver.StepHistory;
        end

        function this = set.GradientsDifferenceHistory(this, gradientsDifferenceHistory)
            arguments
                this
                gradientsDifferenceHistory(1,:) cell
            end
            gradientsDifferenceHistory = iConvertEmptyCellToCanonicalForm(gradientsDifferenceHistory);

            % Verify that the number of elements is less than or equal to
            % HistorySize.
            this.assertNumelLessThanOrEqualTo(gradientsDifferenceHistory, 'GradientsDifferenceHistory')

            this.PrivateSolver.GradientsDifferenceHistory = gradientsDifferenceHistory;

            % Ensure that the solver state performs validation at the next
            % call to lbfgsupdate.
            this.NeedsExtraValidation = true;
        end

        function gradientsDifferenceHistory = get.GradientsDifferenceHistory(this)
            gradientsDifferenceHistory = this.PrivateSolver.GradientsDifferenceHistory;
        end

        function this = set.HistoryIndices(this, historyIndices)
            arguments
                this
                historyIndices(1,:) {mustBeInteger, mustBePositive, mustBeFinite}
            end

            % Verify that the number of elements is less than or equal to
            % HistorySize.
            this.assertNumelLessThanOrEqualTo(historyIndices, 'HistoryIndices')

            this.PrivateSolver.HistoryIndices = historyIndices;

            % Ensure that the solver state performs validation at the next
            % call to lbfgsupdate.
            this.NeedsExtraValidation = true;
        end

        function historyIndices = get.HistoryIndices(this)
            historyIndices = this.PrivateSolver.HistoryIndices;
        end

        function this = set.InverseHessianFactor(this, inverseHessianFactor)
            arguments
                this
                inverseHessianFactor(1,1)  {mustBeReal, mustBePositive, mustBeFinite}
            end
            this.PrivateSolver.InverseHessianApproximation = inverseHessianFactor;
        end

        function inverseHessianFactor = get.InverseHessianFactor(this)
            inverseHessianFactor = this.PrivateSolver.InverseHessianApproximation;
        end

        function loss = get.Loss(this)
            loss = this.CurrentIteration.Fun;
        end

        function gradients = get.Gradients(this)
            gradients = this.CurrentIteration.Gradient;
        end

        function lossFunctionOutputs = get.AdditionalLossFunctionOutputs(this)
            lossFunctionOutputs = this.CurrentIteration.AdditionalFunctionHandleOutputs;
        end

        function stepNorm = get.StepNorm(this)
            if ~isempty(this.PrivateSolver.LastStepTaken)
                stepNorm = this.LBFGSOperationsStrategy.calculateL2Norm(this.PrivateSolver.LastStepTaken);
            elseif this.Initialized
                stepNorm = 0;
            else
                stepNorm = [];
            end
        end

        function gradientsNorm = get.GradientsNorm(this)
            if ~isempty(this.CurrentIteration.Gradient)
                gradientsNorm = this.LBFGSOperationsStrategy.calculateLinfNorm(this.CurrentIteration.Gradient);
            elseif this.Initialized
                gradientsNorm = 0;
            else
                gradientsNorm = [];
            end
        end

        function numLineSearchIterations = get.NumLineSearchIterations(this)
            numLineSearchIterations = this.PrivateSolver.PreviousNumLineSearchIterations;
        end

        function lineSearchStatus = get.LineSearchStatus(this)
            if this.Initialized
                % Only retrieve the status from the private solver if the
                % solver state has already been initialized (i.e. an
                % iteration of L-BFGS has been performed).
                previousLineSearchFailed = this.PrivateSolver.PreviousLineSearchFailed;
                if previousLineSearchFailed
                    lineSearchStatus = "failed";
                else
                    lineSearchStatus = "completed";
                end
            else
                % If an iteration of L-BFGS has not been performed, display
                % an empty char.
                lineSearchStatus = "";
            end
        end

        function out = saveobj(this)
            % Specify version
            out.Version = 2;

            % Save variables needed to reconstruct the private solver
            out.HistorySize = this.HistorySize;
            out.InitialInverseHessianFactor = this.InitialInverseHessianFactor;

            % Save the private solver and current iteration
            out.PrivateSolver = this.PrivateSolver.toStruct();
            out.CurrentIteration = this.CurrentIteration.toStruct();

            % Save solver state variables
            out.StepHistory = this.StepHistory;
            out.GradientsDifferenceHistory = this.GradientsDifferenceHistory;
            out.HistoryIndices = this.HistoryIndices;
            out.InverseHessianFactor = this.InverseHessianFactor;

            % Save previous line search information
            out.LineSearchMethod = this.LineSearchMethod;
            out.MaxNumLineSearchIterations = this.MaxNumLineSearchIterations;
            out.NumLineSearchIterations = this.NumLineSearchIterations;
            out.LineSearchStatus = this.LineSearchStatus;

            % Save the norm of the initial gradients
            out.InitialGradientsNorm = this.InitialGradientsNorm;
        end
    end

    methods(Access=protected)
        function groups = getPropertyGroups( ~ )
            solverStateVariables = {
                'InverseHessianFactor'
                'StepHistory'
                'GradientsDifferenceHistory'
                'HistoryIndices'
                };

            iterationInfo = {
                'Loss'
                'Gradients'
                'AdditionalLossFunctionOutputs'
                'GradientsNorm'
                'StepNorm'
                'LineSearchStatus'
                };

            groups = [
                iPropertyGroupSolverStateVariables( solverStateVariables )
                iPropertyGroupIterationInfo( iterationInfo )
                ];
        end

        function footer = getFooter( this )
            variableName = inputname(1);
            footer = this.createShowAllPropertiesFooter( variableName );
        end
    end

    methods(Access=private)
        function footer = createShowAllPropertiesFooter( ~, variableName )
            text = iGetStringMessage( 'deep:lbfgsupdate:ShowAllProperties' );
            command = sprintf( 'try, displayAllProperties(%s), end', variableName);
            footer = sprintf( '  <a href="matlab:%s">%s</a>\n', command, text );
        end

        function this = initializeSolverState(this, p)
            % Specify which strategy will be used for performing linear
            % algebra operations (e.g. dot product, norms, etc.).
            if isnumeric(p)
                this.LBFGSOperationsStrategy = mlearnlib.internal.utils.solver.ArrayOperationsStrategy;
            else
                this.LBFGSOperationsStrategy = deep.internal.ContainerOperationsStrategy;
            end

            this.InitialGradientsNorm = this.LBFGSOperationsStrategy.calculateLinfNorm(this.CurrentIteration.Gradient);
            this.Initialized = true;
        end

        function gradientEmpty = validateLossFunctionOutputsAndCheckIfEmpty(this, currentIteration)
            fun = currentIteration.Fun;
            gradient = currentIteration.Gradient;

            % Validate that the first output of the function handle is a
            % scalar.
            if ~isnumeric(fun) || ~isreal(fun) || numel(fun) ~= 1
                error(message('deep:lbfgsupdate:invalidLossValue'));
            end

            % Validate that the second output of the function handle is not
            % complex.
            [gradientComplex, gradientEmpty] = ...
                this.LBFGSOperationsStrategy.checkGradientComplexOrEmpty(gradient);
            if gradientComplex
                error(message('nnet_cnn:solver:complexInputsLbfgs'));
            end
        end

        function this = validateSolverState(this)
            historyIndices = this.PrivateSolver.HistoryIndices;
            stepHistory = this.PrivateSolver.StepHistory;
            gradientsDifferenceHistory = this.PrivateSolver.GradientsDifferenceHistory;

            numelHistoryIndicesOrder = numel(historyIndices);
            numelStepHistory = numel(stepHistory);
            numelGradientsDifferenceHistory = numel(gradientsDifferenceHistory);

            % Verify that the number of elements is equal for all solver
            % state variables.
            if ~isequal(numelHistoryIndicesOrder, numelStepHistory, numelGradientsDifferenceHistory)
                error(message('deep:lbfgsupdate:numSolverStatesElementsNotEqual'))
            end
            this.PrivateSolver.HistoryUsed = numelStepHistory;

            % Validate elements in HistoryIndicesOrder are a permuted list
            % of 1:HistoryUsed.
            % NOTE: It is assumed that HistoryIndicesOrder is a circular
            % shift of HistoryUsed:-1:1 in this algorithm, but algorithm 
            % can still work with any permuted list of 1:HistoryUsed, so we
            % only check for that here.
            if ~isequal(sort(historyIndices), 1:numelStepHistory)
                error(message('deep:lbfgsupdate:invalidHistoryIndices'))
            end

            % Assert that elements of StepHistory and GradientsDifferenceHistory
            % are not complex.
            historiesComplex = this.LBFGSOperationsStrategy.checkHistoriesComplex(...
                stepHistory, gradientsDifferenceHistory);
            if historiesComplex
                error(message('nnet_cnn:solver:complexInputsLbfgs'));
            end
            
            % Recalculate Rho based on the new values of StepHistory and
            % GradientsDifferenceHistory.
            for i = historyIndices
                gradientDifferenceDotStep = this.LBFGSOperationsStrategy.calculateScaledDotProduct(...
                    stepHistory{i}, gradientsDifferenceHistory{i}, 1);
                this.PrivateSolver.Rho(i) = 1./nnet.internal.cnn.util.boundAwayFromZero(gradientDifferenceDotStep);
            end

            % Set NeedsExtraValidation to false so we do not validate the
            % solver state variables again unnecessarily.
            this.NeedsExtraValidation = false;
        end

        function assertNumelLessThanOrEqualTo(this, param, paramName)
            % Verify that the number of elements is less than or equal to
            % HistorySize.
            if numel(param) > this.HistorySize
                error(message('deep:lbfgsupdate:numSolverStatesElementsGreaterThanHistorySize', paramName))
            end
        end
    end

    methods(Static)
        function this = loadobj(in)
            if in.Version <= 1
                in = iUpgradeVersionOneToVersionTwo(in);
            end

            % Reconstruct the solver state
            this = deep.LBFGSState(in.HistorySize, in.InitialInverseHessianFactor);

            % Reconstruct the private solver and current iteration
            this.PrivateSolver = this.PrivateSolver.fromStruct(in.PrivateSolver);
            this.CurrentIteration = this.CurrentIteration.fromStruct(in.CurrentIteration);

            % Reinitialize the solver state if it was already initialized
            if ~isempty(this.CurrentIteration.Params)
                initGradientsNorm = in.InitialGradientsNorm;
                this = this.initializeSolverState(this.CurrentIteration.Params);
                this.InitialGradientsNorm = initGradientsNorm;
            end

            % Repopulate solver state variables
            this.StepHistory = in.StepHistory;
            this.GradientsDifferenceHistory = in.GradientsDifferenceHistory;
            this.HistoryIndices = in.HistoryIndices;
            this.InverseHessianFactor = in.InverseHessianFactor;

            % Reconstruct previous line search information
            this.LineSearchMethod = in.LineSearchMethod;
            this.MaxNumLineSearchIterations = in.MaxNumLineSearchIterations;
            this.PrivateSolver.PreviousNumLineSearchIterations = in.NumLineSearchIterations;
            lineSearchStatus = in.LineSearchStatus;
            if strcmp(lineSearchStatus, 'completed')
                this.PrivateSolver.PreviousLineSearchFailed = false;
            elseif strcmp(lineSearchStatus, 'failed')
                this.PrivateSolver.PreviousLineSearchFailed = true;
            end

            % Set NeedsExtraValidation to true so that solver state
            % variables are always validated
            this.NeedsExtraValidation = true;
        end
    end

    methods(Hidden)
        function displayAllProperties(this)
            proplist = properties( this );
            matlab.mixin.CustomDisplay.displayPropertyGroups( ...
                this, iPropertyGroupGeneral( proplist ) );
        end

        function tf = isequal(this, solverState, varargin)
            % isequal   Check equality of LBFGSState objects
            %
            %   tf = isequal(solverState1,solverState2) checks equality of
            %   the LBFGSState objects solverState1 and solverState2. The
            %   function returns 1 (true) when the states match and 0
            %   (false) otherwise.
            %
            %   tf = isequal(solverState1,...,solverStateN) checks equality
            %   of the N LBFGSState objects solverState1, …, solverStateN.

            tf = iCompare(this, solverState, @isequal, varargin{:});
        end

        function tf = isequaln(this, solverState, varargin)
            % isequaln   Check equality of LBFGSState objects ignoring NaNs
            %
            %   tf = isequaln(solverState1,solverState2) checks equality of
            %   the LBFGSState objects solverState1 and solverState2, 
            %   ignoring NaN values. The function returns 1 (true) when the
            %   states match and 0 (false) otherwise.
            %
            %   tf = isequaln(solverState1,...,solverStateN) checks equality
            %   of the N LBFGSState objects solverState1, …, solverStateN.
            tf = iCompare(this, solverState, @isequaln, varargin{:});
        end

        function [this, p] = step(this, p, lossFcn, lineSearchMethod, maxNumLineSearchIterations, numFunctionHandleOutputs)
            this.CurrentIteration = ...
                this.CurrentIteration.setCurrentIterationAndEvaluateFunctionHandleWithValidation(...
                p, lossFcn, numFunctionHandleOutputs);
            
            if ~this.Initialized
                this = this.initializeSolverState(p);
            end            

            % Validate function handle outputs and check if gradient is
            % empty. If it is, return the unchanged p and solver state.
            gradientEmpty = this.validateLossFunctionOutputsAndCheckIfEmpty(this.CurrentIteration);
            if gradientEmpty
                return
            end

            % Perform extra validation on solver state variables if their
            % values have been modified.
            if this.NeedsExtraValidation
                this = this.validateSolverState();
            end

            % Determine which strategy will be used for line search.
            lineSearchStrategy = this.LineSearchStrategyFactory.createLineSearchStrategy( ...
                lineSearchMethod, ...
                maxNumLineSearchIterations, ...
                lossFcn, ...
                numFunctionHandleOutputs);

            % Calculate the next iteration by performing the L-BFGS step.
            [this.PrivateSolver, this.CurrentIteration] = this.PrivateSolver.step(this.CurrentIteration, lineSearchStrategy, this.LBFGSOperationsStrategy);

            p = this.CurrentIteration.Params;

            % Set the values of the last line search method and the
            % maximum number of allowable line search iterations.
            this.LineSearchMethod = lineSearchMethod;
            this.MaxNumLineSearchIterations = maxNumLineSearchIterations;
        end
    end
end

function pgroup = iPropertyGroupGeneral( properties )
pgroup = iPropertyGroup( '', properties );
end

function pgroup = iPropertyGroupSolverStateVariables( properties )
pgroup = iPropertyGroup( '', properties );
end

function pgroup = iPropertyGroupIterationInfo( properties )
title = iGetStringMessage( 'deep:lbfgsupdate:IterationInfoGroupTitle' );
pgroup = iPropertyGroup( title, properties );
end

function stringMessage = iGetStringMessage(id, varargin)
stringMessage = getString( message( id, varargin{:} ) );
end

function pgroup = iPropertyGroup( title, properties )
% iPropertyGroup   Create a property group from title and properties.
pgroup = matlab.mixin.util.PropertyGroup(properties,title);
end

function cellArg = iConvertEmptyCellToCanonicalForm(cellArg)
if numel(cellArg) == 0
    cellArg = {};
end
end

function tf = iCompare(solverState1, solverState2, fun, varargin)
% Compare solverState1 and solverState2.
tf = deep.internal.compareLBFGSStateObjects(solverState1, solverState2, fun);
if nargin > 3
    % Deal with the case of more than two inputs by comparing solverState1
    % with each input separately.
    ii = 1;
    while tf && ii <= nargin-3
        tf = deep.internal.compareLBFGSStateObjects(solverState1, varargin{ii}, fun);
        ii = ii+1;
    end
    return;
end
end

function S = iUpgradeVersionOneToVersionTwo(S)
% iUpgradeVersionOneToVersionTwo   Upgrade a v1 (2023a) saved struct to a v2 saved struct
S.Version = 2;

% As of R2023b, the LBFGSState object should also have a
% InitialGradientsNorm property that specifies the norm of the gradients at
% the 0th iteration (i.e. before performing any updates). Considering this
% value was not stored previously, we set a dummy value of [] for
% LBFGSState objects that have been loaded in previous releases.
S.InitialGradientsNorm = [];
end