classdef BayesOptPrint
    %BayesOptPrint    Customized verbose display for bayesian optimization.
    % Protected methods of this class can be overridden to customize the
    % display.
    %
    % Each method corresponds to a part of the verbose display.
    % The display is broken down into the following parts
    %
    % Header Bar - A contiguous line of "=" that span the width of the
    % display
    %
    % Generic header line 1 - displays the first line of generic
    % bayesopt info like Iter, Eval etc
    %
    % Generic header line 2 - displays the second line of generic
    % bayesopt info like Iter, Eval etc. This line is usually empty
    % and is used when the field names are too long to fit in a single line
    %
    % Optional header line 1- displays the first line of optional header
    % info. Unused by default. Override if you need additioanl values that
    % are not present in bayesopt.
    %
    % Optional header line 2- displays the second line of optional header
    % info. Unused by default. Override if you need additional values that
    % are not present in bayesopt.
    %
    % hyperparameters header line 1- displays the first line of hyperparameter names
    % Override if you need to use different hyperparameter names.
    %
    % hyperparameters header line 2- displays the second line of hyperparameter names
    % Override if you need to use different hyperparameter names.
    %
    % generic optim values - displays the generic iteration result values
    % Override if you need to use different names instead of the generic names.
    %
    % additional optim values - displays the additional iteration result values
    % Override if you need display iteration result values for any additional columns.
    %
    % hyperparameter optim values - displays the hyperparmeter values
    % for the current iteration.
    % Override if you need to modify how the values are displayed.

    %   Copyright 2021-2022 The MathWorks, Inc.
    properties(SetAccess = protected, GetAccess = protected)
        Verbose
        NumVars
        VarNames
        UseParallel
        DisplayHeaderInterval
        IsRegression
        MandatoryFields
        IncludeEstimCol
        NumCoupledConstraints
        HeaderBar
        HeaderLine1
        HeaderLine2
        ObjectiveFieldLine1
        ObjectiveFieldLine2
        IsHeaderConstructed
    end

    properties(SetAccess = protected, GetAccess = protected)
        IterWidth = 7;
        EvalWidth = 9;
        ParallelWidth = 10;
        VarNameFieldWidth = 12;
        MandatoryFieldWidth = 14;
    end

    %% Overridable PUBLIC methods
    methods(Access = public)
        function this =  intializeBayesOptPrint(this,bayesianOptimization)
            this.Verbose = bayesianOptimization.PrivOptions.Verbose;
            this.UseParallel = bayesianOptimization.PrivOptions.UseParallel;
            this.DisplayHeaderInterval = bayesianOptimization.PrivOptions.DisplayHeaderInterval;
            this.IsRegression = bayesianOptimization.PrivOptions.IsClassregRegressionFunction;
            this.IncludeEstimCol = ~acquisitionFunctionIsGridOrRand(bayesianOptimization.PrivOptions.AcquisitionFunctionName);
            this.MandatoryFields = 3+this.IncludeEstimCol;
            this.NumCoupledConstraints = bayesianOptimization.PrivOptions.NumCoupledConstraints;
            this.IsHeaderConstructed = false;
        end

        function printProgress(this, bayesianOptimization,iteration, NumFinishedIterations,isResuming)
            if this.Verbose < 1
                return;
            end
            XTable = bayesianOptimization.XTrace(iteration,:);
            XTable = customizeHyperparameters(this, bayesianOptimization,XTable);
            % generate header if required
            if ~this.IsHeaderConstructed
                this = this.generateHeader(bayesianOptimization,XTable);
            end
            % Display header if required
            if rem(iteration, this.DisplayHeaderInterval) == 1 || this.Verbose == 2 || ...
                    isResuming
                this.printHeader(bayesianOptimization,XTable);
            end
            % Display results of the current iteration
            action = iDetermineAction(bayesianOptimization,iteration);
            this.printIterResults(bayesianOptimization,XTable,iteration,...
                NumFinishedIterations,action);
        end

        function verboseXNextDisplay(this, bayesianOptimization, XTable)
            % This gets called when Verbose >= 2
            XTable = customizeHyperparameters(this, bayesianOptimization, XTable);
            str = bayesoptim.infoString('IterationMessage');
            spaces = repmat(' ', 1, numCharsBeforeX(bayesianOptimization) - numel(str));
            Line = [str, spaces, hyperparametersDisplayLine(this, bayesianOptimization, XTable), '\n'];
            fprintf('\n');
            fprintf(Line);
        end

        function XTable = customizeHyperparameters(this,bayesianOptimization,XTable)
            % customization is a public method because it can be used in other
            % places of bayesan optimization. For example while displaying the
            % final results
        end

        function line = hyperparametersDisplayLine(this,bayesianOptimization,XTable)
            % generates a line with hyperparameter values. This is a public method
            % because it can be used in other places of bayesan optimization. For example
            % while displaying the next point in verbose level two.
            line = commandLineXDisplay(bayesianOptimization, XTable);
        end

        function printVerboseSummary(this, bo)
            % Printed at the start of iteration 1
        end

        function showStoppingReason(this, bo)
            if bo.OutputFcnStop
                bayesoptim.printInfo('OutputFcnStop');
            end
            if bo.PlotFcnStop
                bayesoptim.printInfo('PlotFcnStop');
            end
            if bo.NumObjectiveEvaluations >= bo.PrivOptions.MaxObjectiveEvaluations
                bayesoptim.printInfo('MaxEvalsReached', bo.PrivOptions.MaxObjectiveEvaluations);
            end
            if bo.TotalElapsedTime >= bo.PrivOptions.MaxTime
                bayesoptim.printInfo('MaxTimeReached', num2str(bo.PrivOptions.MaxTime));
            end
            if isempty(bo.XNext) && isequal(bo.PrivOptions.AcquisitionFunctionName, 'grid')
                bayesoptim.printInfo('GridSearched');
            end
            bayesoptim.printInfo('TotalEvaluations', bo.NumObjectiveEvaluations);
            bayesoptim.printInfo('TotalTime', num2str(bo.TotalElapsedTime));
            bayesoptim.printInfo('TotalEvalTime', num2str(nansum(bo.ObjectiveEvaluationTimeTrain)));
        end

        function showBestPoints(this, bo)
            if bo.PrivOptions.Verbose >= 1
                [MinObsXTable, MinObsObjective, MinObsLoc] = bestPoint(bo, 'Criterion', 'min-observed');
                if ~isfinite(MinObsLoc)
                    bayesoptim.printInfo('NoFeasibleResult');
                else
                    % Show min observed feasible point
                    if isequal('asha', bo.PrivOptions.AcquisitionFunctionName)
                        % Remove asha dummy hyperparams from X
                        MinObsXTable.TrainingSetSize = [];
                        MinObsXTable.ASHARung = [];
                    end
                    fprintf('\n');
                    bayesoptim.printInfo('BestObserved');
                    disp(customizeHyperparameters(bo.PrivOptions.BayesOptPrint, bo,MinObsXTable));
                    bayesoptim.printInfo('ObservedObjectiveValue', num2str(MinObsObjective));
                    if ~acqFcnIsNonBayes(bo) && ~isempty(MinObsXTable)
                        bayesoptim.printInfo('EstimatedObjectiveValue', num2str(predictObjective(bo, MinObsXTable)));
                    end
                    bayesoptim.printInfo('ObservedEvaluationTime', num2str(bo.ObjectiveEvaluationTimeTrace(MinObsLoc)));
                    if ~isempty(bo.ConstraintsTrace)
                        bayesoptim.printInfoNoReturn('ObservedConstraintViolations');
                        fprintf('[ ');
                        fprintf('%f ', bo.ConstraintsTrace(MinObsLoc,:));
                        fprintf(']\n');
                    end
                    % Unless using grid, random, or asha, show best point
                    % according to the default criterion
                    if ~acqFcnIsNonBayes(bo)
                        BestPoint = bestPoint(bo);
                        if ~isempty(BestPoint)
                            fprintf('\n');
                            bayesoptim.printInfo('BestEstimated');
                            disp(customizeHyperparameters(bo.PrivOptions.BayesOptPrint, bo,BestPoint));
                            bayesoptim.printInfo('EstimatedObjectiveValue', num2str(predictObjective(bo, BestPoint)));
                            bayesoptim.printInfo('EstimatedEvaluationTime', num2str(predictObjectiveEvaluationTime(bo, BestPoint)));
                            if bo.PrivOptions.NumCoupledConstraints > 0
                                ConstraintPred = cellfun(@(GP)predict(GP, transformPoints(bo, BestPoint)), ...
                                    bo.ConstraintModels);
                                bayesoptim.printInfoNoReturn('EstimatedConstraintViolations');
                                fprintf('[ ');
                                fprintf('%f ', ConstraintPred);
                                fprintf(']\n');
                            end
                        end
                    end
                    fprintf('\n');
                end
            end
        end

        function showVerboseDocHelp(this, bo)
        end

        % plots
        function [Title1, XLabel1, YLabel] = textInMinObjectivePlot(this, bo)
            Title1 = bayesoptim.infoString('MinObjTitle');
            XLabel1 = bayesoptim.infoString('FEvals');
            YLabel = bayesoptim.infoString('MinObj');
        end

        function Curves = plotMinObjective(this, bo, Axes, Curves)
            % Called each iteration. Plots on Axes, adds to Curves
            if isvalid(Axes)
                Curves{end+1} = plotCurve(1:numel(bo.ObjectiveMinimumTrace), ...
                    bo.ObjectiveMinimumTrace, Axes, 'b');
                if all(isnan(bo.EstimatedObjectiveMinimumTrace))
                    legend(Axes, {bayesoptim.infoString('MinObsObj')}, 'Location', 'best');
                else
                    Curves{end+1} = plotCurve(1:numel(bo.ObjectiveMinimumTrace), ...
                        bo.EstimatedObjectiveMinimumTrace, Axes, 'g');
                    legend(Axes, {bayesoptim.infoString('MinObsObj'), ...
                        bayesoptim.infoString('EstMinObj')}, 'Location', 'best');
                end
            end
        end

        % Format of the optimizationResults struct produced when using ASHA
        function optimizationResults = buildAshaOptimizationResults(this, bo)
            % TrainingSetSize appears immediately after the optimizable
            % variables
            optimizationResults = bo.XTrace;
            optimizationResults.ASHARung = [];
            optimizationResults.Objective = bo.ObjectiveTrace;
            optimizationResults.Rank = iRankVector(bo.ObjectiveTrace);
        end
    end

    %% Overridable PROTECTED methods
    methods(Access = protected)

        function printHeader(this,bayesianOptimization,XTable)
            % prints the verbose header.
            % Can be overridden for

            %print header bar
            fprintf(this.HeaderBar);
            fprintf('\n');

            %print header line 1
            fprintf(this.HeaderLine1);
            fprintf('\n');

            %print header line 2
            fprintf(this.HeaderLine2);
            fprintf('\n');

            %print header bottom bar
            fprintf(this.HeaderBar);
            fprintf('\n');
        end

        function this = generateHeader(this,bayesianOptimization,XTable)
            [this.ObjectiveFieldLine1, this.ObjectiveFieldLine2] = this.updateObjectiveFieldLines();
            this.NumVars = length(XTable.Properties.VariableNames);
            this.VarNames = XTable.Properties.VariableNames;
            this.HeaderBar = this.generateHeaderBar(XTable);

            %generate header line 1
            this.HeaderLine1 = [this.genericOptimHeaderLine1(bayesianOptimization),...
                this.additionalOptimHeaderLine1(bayesianOptimization),...
                this.hyperparametersHeaderLine1(bayesianOptimization,XTable)];

            %generate header line 2
            this.HeaderLine2 = [this.genericOptimHeaderLine2(bayesianOptimization),...
                this.additionalOptimHeaderLine2(bayesianOptimization),...
                this.hyperparametersHeaderLine2(bayesianOptimization,XTable)];

            this.IsHeaderConstructed = true;
        end

        function printIterResults(this,bayesianOptimization,XTable,iteration,numFinishedIterations,action)
            fprintf(this.genericOptimValues(bayesianOptimization,iteration,numFinishedIterations,action))
            fprintf(this.additionalOptimValues(bayesianOptimization))
            fprintf(this.hyperparameterValues(bayesianOptimization,XTable))
            fprintf('\n');
        end

        function [objectiveFieldLine1, objectiveFieldLine2] = updateObjectiveFieldLines(this)
            if this.IsRegression
                objectiveFieldLine1 = ' Objective:  ';
                objectiveFieldLine2 = ' log(1+loss) ';
            else
                objectiveFieldLine1 = ' Objective   ';
                objectiveFieldLine2 = '             ';
            end
        end

        function bar = generateHeaderBar(this,XTable)
            % generates the header bar used by bayesopt.
            % default headerbar is made of contiguous '=' signs that span the width
            % of the verbose display.
            % Can be orverridden to customize it.
            NumEquals = this.IterWidth + ...
                this.EvalWidth + ...
                this.MandatoryFields*(this.MandatoryFieldWidth) - 1 + ...
                (this.NumCoupledConstraints + this.NumVars)*(this.VarNameFieldWidth+3);
            if this.UseParallel
                NumEquals = NumEquals + this.ParallelWidth;
            end
            bar = [sprintf('|'), sprintf(repmat('=', 1, NumEquals)), sprintf('|')];
        end

        function out = genericOptimHeaderLine1(this,bayesianOptimization)
            % generates part of the first header line used by bayesopt.
            % For a classification function with 'UseParallel' enabled, is
            % shown below.
            % | Iter | Active  | Eval   | Objective   | Objective   | BestSoFar   | BestSoFar   |
            % Can be orverridden to customize it.
            out = [];
            if this.UseParallel
                out = [out,sprintf('| Iter | Active  | Eval   |%s| Objective   | BestSoFar   |', this.ObjectiveFieldLine1)];
            else
                out = [out,sprintf('| Iter | Eval   |%s| Objective   | BestSoFar   |', this.ObjectiveFieldLine1)];
            end

            if this.IncludeEstimCol
                out = [out,sprintf(' BestSoFar   |')];
            end

            for c = 1:this.NumCoupledConstraints
                out = [out,sprintf(' Constraint%d  |', c)];
            end
        end

        function out = genericOptimHeaderLine2(this,bayesianOptimization)
            % generates part of the second header line used by bayesopt.
            % For a classification function with 'UseParallel' enabled, is
            % shown below.
            % |      | workers | result |             | runtime     | (observed)  | (estim.)    |
            % Can be orverridden to customize it.
            out = [];
            if this.UseParallel
                out = [out,sprintf('|      | workers | result |%s| runtime     | (observed)  |', this.ObjectiveFieldLine2)];
            else
                out = [out,sprintf('|      | result |%s| runtime     | (observed)  |', this.ObjectiveFieldLine2)];
            end

            if this.IncludeEstimCol
                out = [out,sprintf(' (estim.)    |')];
            end

            for c = 1:this.NumCoupledConstraints
                out = [out,sprintf('              |')];
            end
        end

        function out = additionalOptimHeaderLine1(this,bayesianOptimization)
            % generates part of the first header line used by bayesopt.
            % This method can be used to add additional columns to the verbose
            % display like thos used by ASHA.
            % Unused by default.
            % Can be overridden to add additioanl columns it.
            out = '';
        end

        function out = additionalOptimHeaderLine2(this,bayesianOptimization)
            % generates part of the second header line used by bayesopt.
            % This method can be used to add additional columns to the verbose
            % display like thos used by ASHA.
            % Unused by default.
            % Can be overridden to add additioanl columns it.
            out = '';
        end

        function out = hyperparametersHeaderLine1(this,bayesianOptimization,XTable)
            % generates the part of the hyperparameter names left over from the
            % first line.
            % Can be orverridden to customize it.
            out = [];
            varNames = XTable.Properties.VariableNames;
            numVars = length(XTable.Properties.VariableNames);
            for v = 1:numVars
                Name = varNames{v};
                width = this.VarNameFieldWidth;
                if numel(Name) == width+1
                    width = width+1;
                    ControlString = sprintf(' %%%ds|', width);
                elseif numel(Name) > this.VarNameFieldWidth
                    ControlString = sprintf(' %%%ds-|', width);
                else
                    ControlString = sprintf(' %%%ds |', width);
                end
                ControlString = sprintf(ControlString, Name(1:min(numel(Name), width)));
                out = [out,ControlString];
            end
        end

        function out = hyperparametersHeaderLine2(this,bayesianOptimization,XTable)
            % generates the part of the hyperparameter names left over from the
            % first line.
            % Can be orverridden to customize it.
            varNames = XTable.Properties.VariableNames;
            numVars = length(XTable.Properties.VariableNames);
            ControlString = sprintf(' %%-%ds |', this.VarNameFieldWidth);
            out = [];
            for v = 1:numVars
                Name = varNames{v};
                if numel(Name) > this.VarNameFieldWidth+1
                    str = Name(this.VarNameFieldWidth+1 : min(numel(Name), 2*this.VarNameFieldWidth));
                else
                    str = '';
                end
                out = [out,sprintf(ControlString, str)];
            end
        end

        function out = genericOptimValues(this,bayesianOptimization, iteration,numFinishedIterations,action)
            % generates the iteration result values that are generic for
            % bayesian optimization.
            % Can be orverridden to customize it.
            if this.UseParallel
                % Number of jobs that will be running after we launch a new one:
                assert(numFinishedIterations > 0);  % Because printProgress should only be called when jobs finish.
                out = sprintf('| %4d | %7d | %s | %11.5g | %11.5g | %11.5g |', ...
                    iteration, 1+numPendingX(bayesianOptimization.Parallel), action, bayesianOptimization.FTrain(iteration), bayesianOptimization.ObjectiveEvaluationTimeTrain(iteration), ...
                    bayesianOptimization.MinObjective);
            else
                out = sprintf('| %4d | %s | %11.5g | %11.5g | %11.5g |', ...
                    iteration, action, bayesianOptimization.FTrain(iteration), bayesianOptimization.ObjectiveEvaluationTimeTrain(iteration), ...
                    bayesianOptimization.MinObjective);
            end

            if this.IncludeEstimCol
                out = [out,sprintf(' %11.5g |', bayesianOptimization.MinEstimatedObjective)];
            end

            for c = 1:this.NumCoupledConstraints
                out = [out,sprintf(' %12.3g |', bayesianOptimization.ConstraintTrain(iteration,c))];
            end
        end

        function out = additionalOptimValues(this,bayesianOptimization)
            % generates the iteration result values in additioan to the generic
            % bayseian optimization values.
            % Can be orverridden to customize it.
            out = '';
        end

        function out = hyperparameterValues(this,bayesianOptimization,XTable)
            % generates the hyperparameter values generated for the current iteration
            % of bayesian optimization.
            % Can be orverridden to customize it.
            nvars = size(XTable,2);
            out = '';
            for v = 1:nvars
                val = XTable.(v);
                if iscategorical(val)
                    if ismissing(val)
                        Name = '-';
                    else
                        Name = char(val);
                    end
                    out = [out, sprintf(' %12s |', Name(1:min(numel(Name),12)))];
                else
                    if ismissing(val)
                        Name = '-';
                        out = [out, sprintf(' %12s |', Name(1:min(numel(Name),12)))];
                    else
                        out = [out, sprintf(' %12.5g |', val)];
                    end
                end
            end
        end
    end
end

function Curve = plotCurve(X, TraceToPlot, Axes, color)
HoldOff = bayesoptim.holdOn(Axes);
Curve = plot(Axes, X, TraceToPlot, ['-' color 'o'], 'MarkerSize', 2);
end

function action = iDetermineAction(bayesianOptimization,iteration)
[~,~,MinObsLoc] = bestPoint(bayesianOptimization, 'Criterion', 'min-observed');
if isnan(bayesianOptimization.FTrain(iteration))
    action = 'Error ';
elseif ~isempty(bayesianOptimization.ConstraintsTrace) &&...
        any(bayesianOptimization.ConstraintsTrace(iteration,:) > 0)
    action = 'Infeas';
elseif MinObsLoc == iteration
    action = 'Best  ';
else
    action = 'Accept';
end
end

function tf = acquisitionFunctionIsGridOrRand(acquisitionFunctionName)
tf = ismember(acquisitionFunctionName, {'grid','random'});
end