function Stop = plotObjectiveEvaluationTimeModel(bo, State)
%plotObjectiveEvaluationTimeModel     PlotFcn for plotting the
%ObjectiveEvaluationTime-model surface.
%   Stop = plotObjectiveEvaluationTimeModel(bo, State) plots
%   the ObjectiveEvaluationTime-model surface.
%
% See also: BAYESOPT, BAYESIANOPTIMIZATION
%

%   Copyright 2016-2022 The MathWorks, Inc.

persistent Axes Curves;
narginchk(2,2);
Stop = false;
if ~bo.PrivOptions.FitModels || bo.NumVars > 2
    return;
end
set(0, 'defaultfigureunits', 'pixels');     % Make sure units are pixels
switch State
    case 'initial'
        Axes = setupPlot(bo);
        Curves = {};
    case {'iteration', 'done'}
        if isvalid(Axes)
            [Axes, Curves] = updatePlot(bo, Axes, Curves);
        end
    case 'standalone'
        StandaloneAxes = setupPlot(bo);
        updatePlot(bo, StandaloneAxes, {});
end

    function Axes = setupPlot(Results)
        screenSize = get(groot,'ScreenSize');
        L = screenSize(1);
        B = screenSize(2);
        W = screenSize(3);
        H = screenSize(4);
        left = L;
        bottom = H/2 - 100;
        width = W/3 - 50;
        height = H/2;
        f = figure('Position',[left, bottom, width, height], 'Tag','bayesopt.ObjectiveEvaluationTimeModel');
        Axes = axes(f);
        title(Axes, bayesoptim.infoString('FEvalTimeModel'));
        zlabel(Axes, bayesoptim.infoString('EstFEvalTime'));
        if Results.NumVars == 2
            view(Axes, -45, 15);
        end
    end

    function [Axes, Curves] = updatePlot(Results, Axes, Curves)
        import bayesoptim.*
        if isempty(Results.ObjectiveEvaluationTimeModel)
            return;
        end

        set(0, 'defaultfigureunits', 'pixels');     % Make sure units are pixels
        HoldOff = bayesoptim.holdOn(Axes);
        Options = Results.PrivOptions;
        VarSpec = Options.VarSpec;
        Model = Results.ObjectiveEvaluationTimeModel;
        XTrace = Results.XTrace;
        ValTrace = Results.ObjectiveEvaluationTimeTrace;
        ValTrace(isnan(Results.ObjectiveTrace)) = NaN;      % Only plot valid function evaluations

        cellfun(@delete, Curves);
        Curves = {};
        switch Results.NumVars
            case 1
                XPlottable = makeXTablePlottable(Results, XTrace);
                % Create grid to evaluate ObjectiveFcn on (in design matrix shape)
                EvalGrid = NDGrid(VarSpec.LBTrans, VarSpec.UBTrans, Options.NumPlotGrid);
                LegalizedEvalGrid = legalizePoints(Results, EvalGrid);
                % Transform eval grid to native space for plotting
                PlotGrid = untransformPoints(Results, EvalGrid, false);
                PlotGrid = PlotGrid{:,1};
                % (1) Eval GPRs on eval grid
                YPred = exp(predict(Model, LegalizedEvalGrid));
                % (2) Plot training points on all axes
                plot(Axes, XPlottable, ValTrace, 'bo', 'MarkerFaceColor', 'blue');
                % (3) Plot Model's mean
                Curves{end+1} = plot(Axes, PlotGrid, YPred, 'r');
                % (6) Plot next point
                XNextToPlot = makeXTablePlottable(Results, Results.NextPoint);
                MarkerSpec = {'o', 'MarkerSize', 8, 'Color', 'black', ...
                    'MarkerFaceColor', 'black'};
                Curves{end+1} = plot(Axes, XNextToPlot, predictObjectiveEvaluationTime(...
                    Results, Results.NextPoint), MarkerSpec{:});
                % (7) Label axes
                xlabel(Axes, VarSpec.Names{1});
                ylabel(Axes, bayesoptim.infoString('FEvalTime'));
            case 2
                XPlottable = makeXTablePlottable(Results, XTrace);
                % Create grid to evaluate ObjectiveFcn on (in design matrix shape)
                EvalGrid = NDGrid(VarSpec.LBTrans, VarSpec.UBTrans, Options.NumPlotGrid);
                LegalizedEvalGrid = legalizePoints(Results, EvalGrid);
                % Transform eval grid to native space for plotting
                PlotGrid = untransformPoints(Results, EvalGrid, false);
                PlotGrid = PlotGrid{:,:};
                % Make a version of the PlotGrid in meshgrid shape:
                X1Plot = reshape(PlotGrid(:,1), reverse(Options.NumPlotGrid));
                X2Plot = reshape(PlotGrid(:,2), reverse(Options.NumPlotGrid));
                % (1) Eval GPRs on eval grid
                YPred = exp(predict(Model, LegalizedEvalGrid));
                % (2) Plot training points
                Curves{end+1} = plot3(Axes, XPlottable(:,1), XPlottable(:,2), ValTrace, ...
                    'bo', 'MarkerFaceColor', 'blue');
                % (3) Plot FModel's mean and minimum
                Curves{end+1} = surfc(Axes, X1Plot, X2Plot, reshape(YPred, reverse(Options.NumPlotGrid)), ...
                    'FaceColor', 'r', 'FaceAlpha', .5, 'LineStyle', '-');
                % (6) Plot next point
                XNextToPlot = makeXTablePlottable(Results, Results.NextPoint);
                MarkerSpec = {'o', 'MarkerSize', 8, 'Color', 'black', ...
                    'MarkerFaceColor', 'black'};
                Curves{end+1} = plot3(Axes, XNextToPlot(1), XNextToPlot(2), ...
                    predictObjectiveEvaluationTime(Results, Results.NextPoint), ...
                    MarkerSpec{:});
                % (7) Label axes
                xlabel(Axes, VarSpec.Names{1});
                ylabel(Axes, VarSpec.Names{2});
        end
        % set axes scale and limits
        Axes = bayesoptim.setAxesScaleAndLimits(Axes, VarSpec);
    end

end
